# Path: vmoe/checkpoints/serialization.py
# def from_bytes(target: PyTree, encoded_bytes: bytes) -> PyTree:
# def msgpack_restore(encoded_pytree: bytes) -> PyTree:
# def msgpack_serialize(pytree: PyTree, in_place: bool = False) -> bytes:
# def to_bytes(target: PyTree) -> bytes:
# def _shaped_array_to_bytes(x: jax.ShapedArray) -> bytes:
# def _shaped_array_from_bytes(data: bytes) -> jax.ShapedArray:
# def _slice_to_bytes(x: Slice) -> bytes:
# def _slice_from_bytes(data: bytes) -> Slice:
# def _slice_nd_to_bytes(x: SliceNd) -> bytes:
# def _slice_nd_from_bytes(data: bytes) -> SliceNd:
# def _slice_nd_array_to_bytes(x: SliceNdArray) -> bytes:
# def _slice_nd_array_from_bytes(data: bytes) -> SliceNdArray:
# def _lazy_array_chunks_encode(
#     lazy_array_chunks: LazyArrayChunks) -> Dict[str, Any]:
#   def _encode_chunk(pair):
# def _array_chunks_decode(mapping: Mapping[str, Any]) -> Any:
#   def _decode_chunk(pair):
# def _index_info_to_bytes(index_info: IndexInfo) -> bytes:
# def _index_info_from_bytes(data: bytes) -> IndexInfo:
# def _msgpack_ext_pack(x):
# def _msgpack_ext_unpack(code, data):
# _MAX_CHUNK_SIZE = flax.serialization.MAX_CHUNK_SIZE
# _ARRAY_CHUNKS_MAGIC_KEY = '__msgpack_array_chunk_a87ca2__'
# class _MsgpackExtType(enum.IntEnum):


import functools
import itertools
import multiprocessing.pool
import os
import re
from typing import Any, Callable, Iterable, Iterator, Sequence, Mapping, Optional
from tensorflow.io import gfile
from vmoe.checkpoints import serialization
             reverse=decreasing),
      lambda x: x[0]):
    if set(x[1] for x in group) == suffixes:
      yield step


def remove_checkpoints(filepaths: Iterable[str],
                       match_fn: Callable[[str], bool],
                       *,
                       thread_pool: Optional[ThreadPool] = None):
  """Removes checkpoints for which `match_fn` returns True."""

  def remove(filepath):
    if match_fn(filepath):
      gfile.remove(filepath)

  thread_pool = ThreadPool() if thread_pool is None else thread_pool
  thread_pool.map(remove, filepaths)


def remove_shard_suffix(filepath: str) -> str:
  return CHECKPOINT_REGEX.sub(r'\1\2\3', filepath)


def restore_checkpoint(filepath: str, tree: Optional[Any] = None) -> Any:
  with gfile.GFile(filepath, 'rb') as fp:
    checkpoint_contents = fp.read()
  if tree is None:
    return serialization.msgpack_restore(checkpoint_contents)
  else:

# Path: vmoe/checkpoints/serialization.py
# def from_bytes(target: PyTree, encoded_bytes: bytes) -> PyTree:
# def msgpack_restore(encoded_pytree: bytes) -> PyTree:
# def msgpack_serialize(pytree: PyTree, in_place: bool = False) -> bytes:
# def to_bytes(target: PyTree) -> bytes:
# def _shaped_array_to_bytes(x: jax.ShapedArray) -> bytes:
# def _shaped_array_from_bytes(data: bytes) -> jax.ShapedArray:
# def _slice_to_bytes(x: Slice) -> bytes:
# def _slice_from_bytes(data: bytes) -> Slice:
# def _slice_nd_to_bytes(x: SliceNd) -> bytes:
# def _slice_nd_from_bytes(data: bytes) -> SliceNd:
# def _slice_nd_array_to_bytes(x: SliceNdArray) -> bytes:
# def _slice_nd_array_from_bytes(data: bytes) -> SliceNdArray:
# def _lazy_array_chunks_encode(
#     lazy_array_chunks: LazyArrayChunks) -> Dict[str, Any]:
#   def _encode_chunk(pair):
# def _array_chunks_decode(mapping: Mapping[str, Any]) -> Any:
#   def _decode_chunk(pair):
# def _index_info_to_bytes(index_info: IndexInfo) -> bytes:
# def _index_info_from_bytes(data: bytes) -> IndexInfo:
# def _msgpack_ext_pack(x):
# def _msgpack_ext_unpack(code, data):
# _MAX_CHUNK_SIZE = flax.serialization.MAX_CHUNK_SIZE
# _ARRAY_CHUNKS_MAGIC_KEY = '__msgpack_array_chunk_a87ca2__'
# class _MsgpackExtType(enum.IntEnum):


# Path: vmoe/checkpoints/base.py
import functools
import itertools
import multiprocessing.pool
import os
import re
from typing import Any, Callable, Iterable, Iterator, Sequence, Mapping, Optional
from tensorflow.io import gfile
from vmoe.checkpoints import serialization

             reverse=decreasing),
      lambda x: x[0]):
    if set(x[1] for x in group) == suffixes:
      yield step


def remove_checkpoints(filepaths: Iterable[str],
                       match_fn: Callable[[str], bool],
                       *,
                       thread_pool: Optional[ThreadPool] = None):
  """Removes checkpoints for which `match_fn` returns True."""

  def remove(filepath):
    if match_fn(filepath):
      gfile.remove(filepath)

  thread_pool = ThreadPool() if thread_pool is None else thread_pool
  thread_pool.map(remove, filepaths)


def remove_shard_suffix(filepath: str) -> str:
  return CHECKPOINT_REGEX.sub(r'\1\2\3', filepath)


def restore_checkpoint(filepath: str, tree: Optional[Any] = None) -> Any:
  with gfile.GFile(filepath, 'rb') as fp:
    checkpoint_contents = fp.read()
  if tree is None:
    return serialization.msgpack_restore(checkpoint_contents)
  else:

    return serialization.from_bytes(tree, checkpoint_contents)