class_docstr: Tests for modules in masked_layers.py.
class_docstr_tok: ["Tests","for","modules","in","masked_layers.py","."]
docstr: Test Batch/Group Norm unaffected by padding.
docstr_tok: ["Test","Batch\/Group","Norm","unaffected","by","padding","."]
code: def test_norm(self, norm1_spec, norm2_spec):
      """Test Batch/Group Norm unaffected by padding."""
      norm1 = norm1_spec.cls(**norm1_spec.ctor_kwargs)
      norm2 = norm2_spec.cls(**norm2_spec.ctor_kwargs)
  
      # It is OK to re-init since we're keeping the rng constant.
      rng = random.PRNGKey(0)
      norm1_outputs, norm1_params = norm1.init_with_output(
          rng, **norm1_spec.init_kwargs
      )
      norm2_outputs, norm2_params = norm2.init_with_output(
          rng, **norm2_spec.init_kwargs
      )
  
      # Inferred spatial shape has the right shape.
      norm1_unpad, norm1_pad = norm1_spec.process_fn(norm1_outputs)
      norm2_unpad, norm2_pad = norm2_spec.process_fn(norm2_outputs)
  
      # Unpadded parts of both outputs are the same.
      np.testing.assert_allclose(norm1_unpad, norm2_unpad, atol=1e-5)
  
      # All padded parts are zero.
      for part in norm1_pad + norm2_pad:
          np.testing.assert_allclose(
              part, np.zeros_like(part), atol=1e-5
          )
  
      # Run a second time and repeat all the checks. This is necessary for BN
      # because it is stateful; and does not change the output for GN.
      norm1_outputs, norm1_params = norm1.apply(
          norm1_params,
          mutable=["batch_stats"],
          **norm1_spec.init_kwargs,
      )
      norm2_outputs, norm2_params = norm2.apply(
          norm2_params,
          mutable=["batch_stats"],
          **norm2_spec.init_kwargs,
      )
  
      norm1_unpad, norm1_pad = norm1_spec.process_fn(norm1_outputs)
      norm2_unpad, norm2_pad = norm2_spec.process_fn(norm2_outputs)
  
      # Unpadded parts of both outputs are the same.
      np.testing.assert_allclose(norm1_unpad, norm2_unpad, atol=1e-5)
  
      # All padded parts are zero.
      for part in norm1_pad + norm2_pad:
          np.testing.assert_allclose(
              part, np.zeros_like(part), atol=1e-5
          )
  
imports: import dataclasses
  from typing import Any, Callable, Dict, Sequence, Tuple, Type, Union
  from absl.testing import absltest
  from absl.testing import parameterized
  import flax.linen
  from jax import random
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import masked_layers
label_desc: Use image node_id 6 for calling the MaskedLayersTest obj's underlying member method code with example usage: obj.test_norm(norm1_spec, norm2_spec) without return types