code: def _test_transformer_forward_shape(
      self, input_vocab_size, output_vocab_size
  ):
      model = transformer.Transformer(
          input_vocab_size,
          output_vocab_size,
          d_model=32,
          d_ff=64,
          n_encoder_layers=2,
          n_decoder_layers=2,
          n_heads=2,
      )
      xs = [
          np.ones((3, 5)).astype(np.int32),
          np.ones((3, 5)).astype(np.int32),
      ]
      _, _ = model.init(shapes.signature(xs))
      y, _ = model(xs)
  
      vocab_size = output_vocab_size or input_vocab_size
      self.assertEqual(y.shape, (3, 5, vocab_size))
  
file_path: trax/trax/models/transformer_test.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  from trax import fastmath
  from trax import shapes
  from trax.layers import test_utils
  from trax.models import transformer
label_desc: Use image node_id 2 for calling the TransformerTest obj's underlying member method code with example usage: obj._test_transformer_forward_shape(input_vocab_size, output_vocab_size) without return types