class_docstr: Tests for modules in nn_layers.py.
class_docstr_tok: ["Tests","for","modules","in","nn_layers.py","."]
docstr: Test the StochasticDepth module.
docstr_tok: ["Test","the","StochasticDepth","module","."]
code: def test_stochastic_depth(self):
      """Test the StochasticDepth module."""
      rng = random.PRNGKey(0)
      rngs = {"dropout": rng}
  
      inputs_shape = (
          1024,
          8,
          8,
          8,
      )  # Use many batches so averages work out.
      inputs = jnp.array(np.random.normal(size=inputs_shape))
      inputs_np = np.asarray(inputs)
  
      drop_none = nn_layers.StochasticDepth(rate=0.0)
      out_none = drop_none.apply(
          {}, inputs, deterministic=False, rngs=rngs
      )
      np.testing.assert_equal(np.asarray(out_none), inputs_np)
  
      # Make sure we zero out roughly half the samples when rate = 0.5.
      drop_half = nn_layers.StochasticDepth(rate=0.5)
      ones = jnp.ones_like(inputs)
      out_half = drop_half.apply(
          {}, ones, deterministic=False, rngs=rngs
      )
      self.assertAlmostEqual(jnp.mean(out_half), 1.0, places=1)
  
      # Make sure that we always drop full samples.
      # Note that the samples kept are scaled by 1 / (1 - rate).
      for row in out_half:
          assert jnp.all(row == 0.0) or jnp.all(row == 2.0)
  
      out_half_det = drop_half.apply(
          {}, inputs, deterministic=True, rngs=rngs
      )
      np.testing.assert_equal(np.asarray(out_half_det), inputs_np)
  
      drop_all = nn_layers.StochasticDepth(rate=1.0)
      out_all = drop_all.apply(
          {}, inputs, deterministic=False, rngs=rngs
      )
      np.testing.assert_equal(np.asarray(out_all), 0.0)
  
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 3 for calling the NNLayersTest obj's underlying member method code with example usage: obj.test_stochastic_depth() without return types