docstr: Tests for a full batch.
docstr_tok: ["Tests","for","a","full","batch","."]
code: def test_leaky_and_non_leaky_vtrace(self):
      """Tests for a full batch."""
      vtrace_fn = self.variant(
          jax.vmap(
              functools.partial(vtrace.vtrace, lambda_=self._lambda)
          )
      )
      leaky_vtrace_fn = self.variant(
          jax.vmap(
              functools.partial(
                  vtrace.leaky_vtrace, alpha_=1.0, lambda_=self._lambda
              )
          )
      )
      # Get function arguments.
      r_t, discount_t, rho_tm1, v_tm1, bootstrap_value = self._inputs
      v_t = np.concatenate(
          [v_tm1[:, 1:], bootstrap_value[:, None]], axis=1
      )
      # Compute vtrace and leaky vtrace output.
      vtrace_output = vtrace_fn(v_tm1, v_t, r_t, discount_t, rho_tm1)
      leaky_vtrace_output = leaky_vtrace_fn(
          v_tm1, v_t, r_t, discount_t, rho_tm1
      )
      # Test output.
      np.testing.assert_allclose(
          vtrace_output, leaky_vtrace_output, rtol=1e-3
      )
  
file_path: rlax/rlax/_src/vtrace_test.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import chex
  import jax
  import numpy
  from rlax._src import distributions
  from rlax._src import vtrace
label_desc: Use image node_id 5 for calling the VTraceTest obj's underlying member method code with example usage: obj.test_leaky_and_non_leaky_vtrace() without return types