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_rezero_forward_shape(self):
      """Tests that the forward pass runs and returns the expected shape."""
      vocab_size = 16
      model = rezero.ReZeroTransformer(
          vocab_size,
          d_model=32,
          d_ff=64,
          n_encoder_layers=2,
          n_decoder_layers=2,
          n_heads=2,
          max_len=16,
      )
      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/rezero_test.py
imports: from absl.testing import absltest
  import numpy
  from trax import layers
  from trax import shapes
  from trax.models.research import rezero
label_desc: Use image node_id 1 for calling the ReZeroTransformerTest obj's underlying member method code with example usage: obj.test_rezero_forward_shape() without return types