# 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/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/components/dense.py
# def dense(features, name, inputs, kernel_axis_names):
#   if self.weight_params is not None or self.act_params is not None:
#     # TODO: Push the "quantized vs not" decision down into the
#     # AQT library. Currently we make that decision here, because the AQT
#     # library doesn't support DenseGeneral, so there's extra reshapes here
#     # whose performance impact I don't know.
#     aqt_context = aqt_config.QuantContext(
#         update_bounds=False, collect_acts_stats=False)
#     aqt_hparams = aqt_flax_layers.DenseAqt.HParams(
#         weight_prec=self.weight_params.prec,
#         weight_half_shift=self.weight_params.half_shift,
#         quant_act=self.act_params,  # currently supports fixed bounds only.
#         quant_type=aqt.QuantType.aqt,
#         weight_quant_granularity=aqt_config.QuantGranularity.per_channel,
#     )
#     batch, seq_len, channels = inputs.shape
#     inputs = inputs.reshape((batch * seq_len, channels))
#     result = aqt_flax_layers.DenseAqt(
#         features=features,
#         hparams=aqt_hparams,
#         train=enable_dropout,
#         quant_context=aqt_context,
#         paxis_name=None,
#         # No "cross-replica" reduction expressed in the XLA graph at this
#         # stage. Will be imposed later, automatically, by XLA SPMD.
#         use_bias=self.use_bias,
#         kernel_init=self.kernel_init,
#         bias_init=self.bias_init,
#         dtype=self.dtype,
#         kernel_axis_names=kernel_axis_names,
#         name=name,
#     )(inputs, padding_mask=None)
#     return result.reshape((batch, seq_len, features))
#   else:
#     return DenseGeneral(
#         features,
#         use_bias=self.use_bias,
#         dtype=self.dtype,
#         kernel_init=self.kernel_init,
#         bias_init=self.bias_init,
#         kernel_axis_names=kernel_axis_names,
#         name=name)(
#             inputs)


import functools
import jax
import jax.numpy as jnp
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from flax import linen as nn
from flax.linen import partitioning
from jax import dtypes
from jax import random
from jax.nn import initializers
from flaxformer import sharding
from flaxformer import testing_utils
from flaxformer.components import dense
from aqt import quantization as aqt
    y, variables = model.init_with_output(rng, x)
    # We transform the last input dimension to two output dimensions (2, 2).
    np.testing.assert_allclose(y, np.full((1, 2, 2), 3.))

    # The output sharding dimensions have been collapsed.
    sharding.check_params_and_axis_names_match(variables)
    self.assertEqual(variables['params_axes']['kernel_axes'],
                     sharding.axis_names('a', 'b * c'))

  def test_dense_general_two_axes(self):
    rng = random.PRNGKey(0)
    x = jnp.ones((1, 2, 2))
    model = dense.DenseGeneral(
        features=3,
        use_bias=False,
        axis=(-2, 2),  # Note: this is the same as (1, 2).
        kernel_init=initializers.ones,
        bias_init=initializers.ones,
        kernel_axis_names=('a', 'b', 'c'),
    )
    y, variables = model.init_with_output(rng, x)
    # We transform the last two input dimensions (2, 2) to one output dimension.
    np.testing.assert_allclose(y, np.full((1, 3), 4.))

    # The input sharding dimensions have been collapsed.
    sharding.check_params_and_axis_names_match(variables)
    self.assertEqual(variables['params_axes']['kernel_axes'],
                     sharding.axis_names('a * b', 'c'))

  def test_mlp_same_out_dim(self):

# 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/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/components/dense.py
# def dense(features, name, inputs, kernel_axis_names):
#   if self.weight_params is not None or self.act_params is not None:
#     # TODO: Push the "quantized vs not" decision down into the
#     # AQT library. Currently we make that decision here, because the AQT
#     # library doesn't support DenseGeneral, so there's extra reshapes here
#     # whose performance impact I don't know.
#     aqt_context = aqt_config.QuantContext(
#         update_bounds=False, collect_acts_stats=False)
#     aqt_hparams = aqt_flax_layers.DenseAqt.HParams(
#         weight_prec=self.weight_params.prec,
#         weight_half_shift=self.weight_params.half_shift,
#         quant_act=self.act_params,  # currently supports fixed bounds only.
#         quant_type=aqt.QuantType.aqt,
#         weight_quant_granularity=aqt_config.QuantGranularity.per_channel,
#     )
#     batch, seq_len, channels = inputs.shape
#     inputs = inputs.reshape((batch * seq_len, channels))
#     result = aqt_flax_layers.DenseAqt(
#         features=features,
#         hparams=aqt_hparams,
#         train=enable_dropout,
#         quant_context=aqt_context,
#         paxis_name=None,
#         # No "cross-replica" reduction expressed in the XLA graph at this
#         # stage. Will be imposed later, automatically, by XLA SPMD.
#         use_bias=self.use_bias,
#         kernel_init=self.kernel_init,
#         bias_init=self.bias_init,
#         dtype=self.dtype,
#         kernel_axis_names=kernel_axis_names,
#         name=name,
#     )(inputs, padding_mask=None)
#     return result.reshape((batch, seq_len, features))
#   else:
#     return DenseGeneral(
#         features,
#         use_bias=self.use_bias,
#         dtype=self.dtype,
#         kernel_init=self.kernel_init,
#         bias_init=self.bias_init,
#         kernel_axis_names=kernel_axis_names,
#         name=name)(
#             inputs)


# Path: flaxformer/components/dense_test.py
import functools
import jax
import jax.numpy as jnp
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from flax import linen as nn
from flax.linen import partitioning
from jax import dtypes
from jax import random
from jax.nn import initializers
from flaxformer import sharding
from flaxformer import testing_utils
from flaxformer.components import dense
from aqt import quantization as aqt

    y, variables = model.init_with_output(rng, x)
    # We transform the last input dimension to two output dimensions (2, 2).
    np.testing.assert_allclose(y, np.full((1, 2, 2), 3.))

    # The output sharding dimensions have been collapsed.
    sharding.check_params_and_axis_names_match(variables)
    self.assertEqual(variables['params_axes']['kernel_axes'],
                     sharding.axis_names('a', 'b * c'))

  def test_dense_general_two_axes(self):
    rng = random.PRNGKey(0)
    x = jnp.ones((1, 2, 2))
    model = dense.DenseGeneral(
        features=3,
        use_bias=False,
        axis=(-2, 2),  # Note: this is the same as (1, 2).
        kernel_init=initializers.ones,
        bias_init=initializers.ones,
        kernel_axis_names=('a', 'b', 'c'),
    )
    y, variables = model.init_with_output(rng, x)
    # We transform the last two input dimensions (2, 2) to one output dimension.
    np.testing.assert_allclose(y, np.full((1, 3), 4.))

    # The input sharding dimensions have been collapsed.
    sharding.check_params_and_axis_names_match(variables)
    self.assertEqual(variables['params_axes']['kernel_axes'],
                     sharding.axis_names('a * b', 'c'))

  def test_mlp_same_out_dim(self):

    module = dense.MlpBlock(