code: def test_rsu_forward_shape(self):
      batch_size = 3
      seq_len = 32
      d_model = 17
      model = rse.ResidualSwitchUnit(
          d_model=d_model, dropout=0.1, mode="train"
      )
      x = np.ones((batch_size, seq_len, d_model)).astype(np.int32)
      _, _ = model.init(shapes.signature(x))
      y = model(x)
      self.assertEqual(y.shape, (batch_size, seq_len, d_model))
  
file_path: trax/trax/models/research/rse_test.py
imports: from absl.testing import absltest
  import numpy
  from trax import shapes
  from trax.models.research import rse
label_desc: Use image node_id 1 for calling the RSETest obj's underlying member method code with example usage: obj.test_rsu_forward_shape() without return types