code: def test_transformer_lm_forward_shape(self):
      vocab_size = 16
      model = ct.ConfigurableTransformerLM(
          vocab_size, d_model=32, d_ff=64, n_layers=2, n_heads=2
      )
      x = np.ones((3, 5)).astype(np.int32)
      _, _ = model.init(shapes.signature(x))
      y = model(x)
      self.assertEqual(y.shape, (3, 5, vocab_size))
  
file_path: trax/trax/models/research/configurable_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 layers
  from trax import shapes
  from trax.layers import test_utils
  from trax.models.research import configurable_transformer
label_desc: Use image node_id 1 for calling the ConfigurableTransformerTest obj's underlying member method code with example usage: obj.test_transformer_lm_forward_shape() without return types