class_docstr: Tests for modules in nn_layers.py.
class_docstr_tok: ["Tests","for","modules","in","nn_layers.py","."]
docstr: Test Residual module.
docstr_tok: ["Test","Residual","module","."]
code: def test_residual(self, residual_type):
      """Test Residual module."""
      rng = random.PRNGKey(0)
      inputs_shape = (16, 32, 32, 3)
  
      residual_module_def = nn_layers.Residual(
          residual_type=residual_type
      )
      x = jnp.array(np.random.normal(size=inputs_shape))
      y = jnp.array(np.random.normal(size=inputs_shape))
      outputs, _ = residual_module_def.init_with_output(rng, x, y)
  
      # test output shape for 4d inputs
      self.assertEqual(outputs.shape, inputs_shape)
  
      # test the residual connection
      if residual_type == "add":
          self.assertTrue(jnp.array_equal(outputs, x + y))
  
      # make sure the residual connection is an identity mapping at initialization
      else:
          np.testing.assert_allclose(outputs, x, atol=1e-3)
  
file_path: scenic/scenic/model_lib/layers/tests/test_nn_layers.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from jax import random
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling the NNLayersTest obj's underlying member method code with example usage: obj.test_residual(residual_type) without return types