# Path: flaxformer/sharding.py
# class AxisNames(tuple):
# def axis_names(*names: str) -> partitioning.AxisMetadata:
# def reduce_fn(x, y):
# def _get_single_sowed_value(value) -> AxisNames:
# def get_axis_names(variables):
# def check_params_and_axis_names_match(variables):
#   def _flatten_with_joined_names(xs):


from absl.testing import absltest
from flax import linen as nn
from flax.linen import partitioning
from jax import numpy as jnp
from flaxformer import sharding
import jax
      sharding.check_params_and_axis_names_match(
          variables={
              "params": {
                  "foo": {
                      "bar": jnp.array([1, 2, 3])
                  }
              },
              "params_axes": {
                  "foo": {
                      "bar_axes":
                          sharding.axis_names("unsharded", "model", "vocab")
                  }
              },
          })

  def test_get_axis_names(self):
    variables = {
        "params": {
            "foo": {
                # Make sure the method is not distracted by actual parameters.
                "bar": jnp.array([1, 2, 3])
            }
        },
        "params_axes": {
            "foo": {
                "bar_axes": sharding.axis_names("unsharded", "model", "vocab")
            }
        },
    }
    self.assertEqual(

# Path: flaxformer/sharding.py
# class AxisNames(tuple):
# def axis_names(*names: str) -> partitioning.AxisMetadata:
# def reduce_fn(x, y):
# def _get_single_sowed_value(value) -> AxisNames:
# def get_axis_names(variables):
# def check_params_and_axis_names_match(variables):
#   def _flatten_with_joined_names(xs):


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

      sharding.check_params_and_axis_names_match(
          variables={
              "params": {
                  "foo": {
                      "bar": jnp.array([1, 2, 3])
                  }
              },
              "params_axes": {
                  "foo": {
                      "bar_axes":
                          sharding.axis_names("unsharded", "model", "vocab")
                  }
              },
          })

  def test_get_axis_names(self):
    variables = {
        "params": {
            "foo": {
                # Make sure the method is not distracted by actual parameters.
                "bar": jnp.array([1, 2, 3])
            }
        },
        "params_axes": {
            "foo": {
                "bar_axes": sharding.axis_names("unsharded", "model", "vocab")
            }
        },
    }
    self.assertEqual(

        sharding.get_axis_names(variables),