# Path: flaxformer/testing_utils.py
# def param_shapes(params: Mapping[str, Any]) -> Dict[str, Any]:
# def param_dtypes_shapes(params: Mapping[str, Any]) -> Dict[str, Any]:
# def param_dtypes_shapes_axes(params: Mapping[str, Any],
#                              params_axes: Mapping[str, Any]) -> Dict[str, Any]:
#   def remove_axes_suffix(ks):
#   def _create_entry(param, param_axes):
# def format_params_shapes(params_shapes: Dict[str, Any]) -> str:
#   def re_compact_arrays(array_match) -> str:
#   def __init__(self, base_path: str):
#   def get_params(self, filename: str) -> Mapping[str, Any]:
#   def check_params(
#       self,
#       actual_params: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
#   def check_params_and_axes(
#       self,
#       actual_params: Mapping[str, Any],
#       actual_params_axes: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
#   def check_params_shapes_only(
#       self,
#       actual_params: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
# class ExpectedJsonFiles:


from absl.testing import absltest
from flax import linen as nn
from flax.linen import partitioning as nn_partitioning
from flaxformer import testing_utils
import jax.numpy as jnp
  "foo[bar]": ["baz", 1, 2, 3]
}""")

  def test_param_dtypes_shapes_axes(self):
    params = nn.FrozenDict({
        "a": {
            "b": jnp.zeros([3, 7], dtype=jnp.float32),
            "c": {
                "d": jnp.zeros([9], dtype=jnp.float32),
            },
        },
        "b": {
            "c": jnp.zeros([3, 7, 4], dtype=jnp.float32),
        },
    })

    params_axes = nn.FrozenDict({
        "a": {
            "b_axes": nn_partitioning.AxisMetadata(names=("vocab", "embed")),
            "c": {
                "d_axes": nn_partitioning.AxisMetadata(names=("embed",)),
            },
        },
        "b": {
            "c_axes":
                nn_partitioning.AxisMetadata(names=("embed", "mlp", "output")),
        },
    })

    result = testing_utils.format_params_shapes(

# Path: flaxformer/testing_utils.py
# def param_shapes(params: Mapping[str, Any]) -> Dict[str, Any]:
# def param_dtypes_shapes(params: Mapping[str, Any]) -> Dict[str, Any]:
# def param_dtypes_shapes_axes(params: Mapping[str, Any],
#                              params_axes: Mapping[str, Any]) -> Dict[str, Any]:
#   def remove_axes_suffix(ks):
#   def _create_entry(param, param_axes):
# def format_params_shapes(params_shapes: Dict[str, Any]) -> str:
#   def re_compact_arrays(array_match) -> str:
#   def __init__(self, base_path: str):
#   def get_params(self, filename: str) -> Mapping[str, Any]:
#   def check_params(
#       self,
#       actual_params: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
#   def check_params_and_axes(
#       self,
#       actual_params: Mapping[str, Any],
#       actual_params_axes: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
#   def check_params_shapes_only(
#       self,
#       actual_params: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
# class ExpectedJsonFiles:


# Path: flaxformer/testing_utils_test.py
from absl.testing import absltest
from flax import linen as nn
from flax.linen import partitioning as nn_partitioning
from flaxformer import testing_utils
import jax.numpy as jnp

  "foo[bar]": ["baz", 1, 2, 3]
}""")

  def test_param_dtypes_shapes_axes(self):
    params = nn.FrozenDict({
        "a": {
            "b": jnp.zeros([3, 7], dtype=jnp.float32),
            "c": {
                "d": jnp.zeros([9], dtype=jnp.float32),
            },
        },
        "b": {
            "c": jnp.zeros([3, 7, 4], dtype=jnp.float32),
        },
    })

    params_axes = nn.FrozenDict({
        "a": {
            "b_axes": nn_partitioning.AxisMetadata(names=("vocab", "embed")),
            "c": {
                "d_axes": nn_partitioning.AxisMetadata(names=("embed",)),
            },
        },
        "b": {
            "c_axes":
                nn_partitioning.AxisMetadata(names=("embed", "mlp", "output")),
        },
    })

    result = testing_utils.format_params_shapes(

        testing_utils.param_dtypes_shapes_axes(params, params_axes))