code: def test_transformer_noencdec_forward_shape(self):
      input_vocab_size = 16
      output_vocab_size = 16
  
      model = transformer2.Transformer2(
          input_vocab_size,
          output_vocab_size,
          d_model=32,
          d_ff=64,
          n_encoder_layers=2,
          n_decoder_layers=2,
          n_heads=2,
      )
  
      enc_toks = np.array([[6, 2, 0, 0, 0, 0], [6, 3, 7, 0, 0, 0]])
      dec_toks = np.array([[4, 2, 0, 0], [8, 5, 0, 0]])
  
      xs = [enc_toks, dec_toks]
      _, _ = model.init(shapes.signature(xs))
  
      # decoder output, decoder mask
      ys = model(xs)
  
      # (B, L2, H)
      self.assertEqual(
          ys[0].shape,
          (dec_toks.shape[0], dec_toks.shape[1], output_vocab_size),
      )
  
      self.assertEqual(ys[1].shape, dec_toks.shape)
  
file_path: trax/trax/models/research/transformer2_test.py
imports: from absl.testing import absltest
  import numpy
  from trax import shapes
  from trax.models.research import transformer2
label_desc: Use image node_id 6 for calling the Transformer2Test obj's underlying member method code with example usage: obj.test_transformer_noencdec_forward_shape() without return types