docstr: Tests for a full batch.
docstr_tok: ["Tests","for","a","full","batch","."]
code: def test_lambda_returns_batch(self):
      """Tests for a full batch."""
      lambda_returns = self.variant(
          jax.vmap(
              functools.partial(
                  multistep.lambda_returns, lambda_=self.lambda_
              )
          )
      )
      # Compute lambda return in batch.
      actual = lambda_returns(self.r_t, self.discount_t, self.v_t)
      # Test return estimate.
      np.testing.assert_allclose(self.expected, actual, 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 LambdaReturnsTest obj's underlying member method code with example usage: obj.test_lambda_returns_batch() without return types