code: def test_gae_shape_on_random(self):
      # create random data, simulating 4 parallel envs and 20 time_steps
      envs, steps = 10, 100
      rewards = np.random.choice(
          [-1.0, 0.0, 1.0], size=(steps, envs), p=[0.01, 0.98, 0.01]
      )
      terminal_masks = np.ones(shape=(steps, envs), dtype=np.float64)
      values = np.random.random(size=(steps + 1, envs))
      discount = 0.99
      gae_param = 0.95
      adv = ppo_lib.gae_advantages(
          rewards, terminal_masks, values, discount, gae_param
      )
      self.assertEqual(adv.shape, (steps, envs))
  
imports: from absl.testing import absltest
  from flax.training import train_state
  import jax
  import ml_collections
  import numpy
  import numpy.testing
  import agent
  import env_utils
  import models
  import ppo_lib
label_desc: Use image node_id 1 for calling the TestGAE obj's underlying member method code with example usage: obj.test_gae_shape_on_random() without return types