code: def test_embed_shape(self):
      # Test output shape [T?, B, D+A+1].
      emb = embedding.embed_oar(
          self._features,
          self._actions,
          self._rewards,
          self._num_actions,
      )
      np.testing.assert_array_equal(
          emb.shape[-1],
          self._features.shape[-1] + self._num_actions + 1,
      )
  
file_path: rlax/rlax/_src/embedding_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import numpy
  from rlax._src import embedding
label_desc: Use image node_id 3 for calling the EmbeddingTest obj's underlying member method code with example usage: obj.test_embed_shape() without return types