# Path: vmoe/train/initialization.py
# def initialize_from_vmoe_release(
#     *,
#     params: PyTree,
#     axis_resources: PyTree,
#     prefix: str,
#     mapping: Sequence[Tuple[str, str]] = (),
#     keep: Sequence[str] = (),
#     thread_pool: Optional[ThreadPool] = None,
#     mesh: Optional[maps.Mesh] = None,
# ) -> PyTree:
# def initialize_from_vit(
#     *,
#     params: PyTree,
#     axis_resources: PyTree,
#     filepath: str,
#     mapping: Sequence[Tuple[str, str]] = (),
#     keep: Sequence[str] = (),
#     broadcast: Sequence[str] = (),
#     mesh: Optional[maps.Mesh] = None,
# ) -> PyTree:
# def _is_broadcastable(shape1: Tuple[int, ...], shape2: Tuple[int, ...]) -> bool:
# def _map_name(name: str, map_regexes: Sequence[Tuple[re.Pattern, str]]) -> str:
# def _pjit_donate_to_device(data, axis_resources, mesh):
# def _zoom_position_embedding(source, target):
#   def _get_tok_and_grid_emb(value):


from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from jax.experimental import maps
from jax.experimental import pjit
from vmoe.train import initialization
import flax.core
import jax
import numpy as np
    super().setUp()
    self.mock_load = self.enter_context(
        mock.patch.object(initialization.vmoe_partitioned_checkpoint,
                          'restore_checkpoint'))
    self.mock_load.return_value = flax.core.FrozenDict({
        'foo': {
            'a': np.ones((4, 3), dtype=np.float64),
            'b': np.ones((4, 4), dtype=np.float64),
            'c': np.ones((3, 3), dtype=np.float64),
            'posembed_input': {
                'pos_embedding': np.ones((1, 4, 3), np.float64),
            },
        },
    })
    self.params = flax.core.FrozenDict({
        'foo': {
            'A': np.zeros((4, 3), dtype=np.float32),
            'B': np.zeros((4, 4), dtype=np.float32),
            'C': np.zeros((3, 3), dtype=np.float32),
            'posembed_input': {
                'pos_embedding': np.zeros((1, 16, 3), np.float32),
            },
        },
    })
    self.axis_resources = jax.tree_map(lambda _: pjit.PartitionSpec(),
                                       self.params)

  def test_success(self):
    with self.assertLogs() as logs:
      with maps.mesh(np.asarray(jax.local_devices()), ('d',)):

# Path: vmoe/train/initialization.py
# def initialize_from_vmoe_release(
#     *,
#     params: PyTree,
#     axis_resources: PyTree,
#     prefix: str,
#     mapping: Sequence[Tuple[str, str]] = (),
#     keep: Sequence[str] = (),
#     thread_pool: Optional[ThreadPool] = None,
#     mesh: Optional[maps.Mesh] = None,
# ) -> PyTree:
# def initialize_from_vit(
#     *,
#     params: PyTree,
#     axis_resources: PyTree,
#     filepath: str,
#     mapping: Sequence[Tuple[str, str]] = (),
#     keep: Sequence[str] = (),
#     broadcast: Sequence[str] = (),
#     mesh: Optional[maps.Mesh] = None,
# ) -> PyTree:
# def _is_broadcastable(shape1: Tuple[int, ...], shape2: Tuple[int, ...]) -> bool:
# def _map_name(name: str, map_regexes: Sequence[Tuple[re.Pattern, str]]) -> str:
# def _pjit_donate_to_device(data, axis_resources, mesh):
# def _zoom_position_embedding(source, target):
#   def _get_tok_and_grid_emb(value):


# Path: vmoe/train/initialization_test.py
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from jax.experimental import maps
from jax.experimental import pjit
from vmoe.train import initialization
import flax.core
import jax
import numpy as np

    super().setUp()
    self.mock_load = self.enter_context(
        mock.patch.object(initialization.vmoe_partitioned_checkpoint,
                          'restore_checkpoint'))
    self.mock_load.return_value = flax.core.FrozenDict({
        'foo': {
            'a': np.ones((4, 3), dtype=np.float64),
            'b': np.ones((4, 4), dtype=np.float64),
            'c': np.ones((3, 3), dtype=np.float64),
            'posembed_input': {
                'pos_embedding': np.ones((1, 4, 3), np.float64),
            },
        },
    })
    self.params = flax.core.FrozenDict({
        'foo': {
            'A': np.zeros((4, 3), dtype=np.float32),
            'B': np.zeros((4, 4), dtype=np.float32),
            'C': np.zeros((3, 3), dtype=np.float32),
            'posembed_input': {
                'pos_embedding': np.zeros((1, 16, 3), np.float32),
            },
        },
    })
    self.axis_resources = jax.tree_map(lambda _: pjit.PartitionSpec(),
                                       self.params)

  def test_success(self):
    with self.assertLogs() as logs:
      with maps.mesh(np.asarray(jax.local_devices()), ('d',)):

        params = initialization.initialize_from_vmoe_release(