docstr: Tests for a single element.
docstr_tok: ["Tests","for","a","single","element","."]
code: def test_discounted_returns_batch(self):
      """Tests for a single element."""
      discounted_returns = self.variant(
          jax.vmap(multistep.discounted_returns)
      )
      # Compute discounted return.
      actual_scalar = discounted_returns(
          self.r_t, self.discount_t, self.bootstrap_v
      )
      actual_vector = discounted_returns(
          self.r_t, self.discount_t, self.v_t
      )
      # Test output.
      np.testing.assert_allclose(
          self.expected, actual_scalar, rtol=1e-5
      )
      np.testing.assert_allclose(
          self.expected, actual_vector, rtol=1e-5
      )
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import chex
  import jax
  import jax.numpy
  import numpy
  from rlax._src import multistep
label_desc: Use image node_id 2 for calling the DiscountedReturnsTest obj's underlying member method code with example usage: obj.test_discounted_returns_batch() without return types