code: def test_shapes(self, advantage_fn):
      rewards = np.array([[1, 1, 1]], dtype=np.float32)
      returns = np.array([[3, 2, 1]], dtype=np.float32)
      values = np.array([[2, 2, 2]], dtype=np.float32)
      dones = np.array([[False, False, True]])
      discount_mask = np.ones_like(rewards)
      adv1 = advantage_fn(gamma=1, margin=1)(
          rewards, returns, values, dones, discount_mask
      )
      self.assertEqual(adv1.shape, (1, 2))
      adv2 = advantage_fn(gamma=1, margin=2)(
          rewards, returns, values, dones, discount_mask
      )
      self.assertEqual(adv2.shape, (1, 1))
  
file_path: trax/trax/rl/advantages_test.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  from trax.rl import advantages
label_desc: Use image node_id 1 for calling the AdvantagesTest obj's underlying member method code with example usage: obj.test_shapes(advantage_fn) without return types