# Path: flaxformer/components/layer_norm.py
# class T5LayerNorm(nn.Module, param_remapping.ParameterRemappable):
#   def __call__(self,
#                x: Array,
#                decode: bool = False,
#                prefill: bool = False,
#                prefill_lengths: Optional[Array] = None) -> Array:


from absl.testing import absltest
from flax import linen as nn
from jax import numpy as jnp
from jax import random
from flaxformer.components import layer_norm
import chex
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for layer_norm."""




class T5LayerNormTest(absltest.TestCase):

  def test_layer_norm(self):
    rng = random.PRNGKey(0)
    key1, key2, key3 = random.split(rng, 3)
    x = random.normal(key1, (2, 3, 4))
    model_fn = lambda dtype: layer_norm.T5LayerNorm(dtype=dtype)
    y, _ = model_fn(jnp.float32).init_with_output(key2, x)
    self.assertEqual(x.shape, y.shape)
    self.assertEqual(y.dtype, jnp.float32)

    y, _ = model_fn(jnp.int32).init_with_output(key3, x)
    self.assertEqual(y.dtype, jnp.int32)

  def test_default_axis_name(self):

# Path: flaxformer/components/layer_norm.py
# class T5LayerNorm(nn.Module, param_remapping.ParameterRemappable):
#   def __call__(self,
#                x: Array,
#                decode: bool = False,
#                prefill: bool = False,
#                prefill_lengths: Optional[Array] = None) -> Array:


# Path: flaxformer/components/layer_norm_test.py
from absl.testing import absltest
from flax import linen as nn
from jax import numpy as jnp
from jax import random
from flaxformer.components import layer_norm
import chex

# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for layer_norm."""




class T5LayerNormTest(absltest.TestCase):

  def test_layer_norm(self):
    rng = random.PRNGKey(0)
    key1, key2, key3 = random.split(rng, 3)
    x = random.normal(key1, (2, 3, 4))
    model_fn = lambda dtype: layer_norm.T5LayerNorm(dtype=dtype)
    y, _ = model_fn(jnp.float32).init_with_output(key2, x)
    self.assertEqual(x.shape, y.shape)
    self.assertEqual(y.dtype, jnp.float32)

    y, _ = model_fn(jnp.int32).init_with_output(key3, x)
    self.assertEqual(y.dtype, jnp.int32)

  def test_default_axis_name(self):

    module = layer_norm.T5LayerNorm()