docstr: Tests correctness for full batch.
docstr_tok: ["Tests","correctness","for","full","batch","."]
code: def test_transformed_retrace_batch(self, tx_pair, td_index):
      """Tests correctness for full batch."""
      transformed_retrace = self.variant(
          jax.vmap(
              functools.partial(
                  nonlinear_bellman.transformed_retrace,
                  tx_pair=tx_pair,
                  lambda_=self._lambda,
              )
          )
      )
      # Compute transformed vtrace td errors in batch.
      actual_td = transformed_retrace(
          self._qs[:, :-1],
          self._targnet_qs[:, 1:],
          self._actions[:, :-1],
          self._actions[:, 1:],
          self._rewards[:, :-1],
          self._pcontinues[:, :-1],
          self._target_policy_probs[:, 1:],
          self._behavior_policy_probs[:, 1:],
      )
      # Test output.
      np.testing.assert_allclose(
          self.expected_td[td_index], actual_td, rtol=1e-3
      )
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import chex
  import jax
  import numpy
  from rlax._src import nonlinear_bellman
label_desc: Use image node_id 2 for calling the TransformedRetraceTest obj's underlying member method code with example usage: obj.test_transformed_retrace_batch(tx_pair, td_index) without return types