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