docstr: Tests that the forward pass runs and returns the expected shape.
docstr_tok: ["Tests","that","the","forward","pass","runs","and","returns","the","expected","shape","."]
code: def test_layerdrop_layerwise_skip_fraction(self):
      """Tests that the forward pass runs and returns the expected shape."""
      vocab_size = 16
      model = layerdrop_transformer.LayerDropTransformerLM(
          vocab_size,
          d_model=16,
          d_ff=32,
          n_layers=2,
          n_heads=2,
          max_len=16,
          skip_fraction=[0.2, 0.8],
      )
      xs = [
          np.ones((1, 8)).astype(np.int32),
          np.ones((1, 8)).astype(np.int32),
      ]
      _, _ = model.init(shapes.signature(xs))
      ys = model(xs)
      self.assertEqual([y.shape for y in ys], [(1, 8, 16), (1, 8)])
  
file_path: trax/trax/models/research/layerdrop_transformer_test.py
imports: from absl.testing import absltest
  import numpy
  from trax import shapes
  from trax.models.research import layerdrop_transformer
label_desc: Use image node_id 2 for calling the LayerDropTransformerTest obj's underlying member method code with example usage: obj.test_layerdrop_layerwise_skip_fraction() without return types