docstr: Tests correctness for full batch.
docstr_tok: ["Tests","correctness","for","full","batch","."]
code: def test_transformed_q_lambda_batch(self, tx_pair, td_index):
      """Tests correctness for full batch."""
      transformed_n_step_q_learning = self.variant(
          jax.vmap(
              functools.partial(
                  nonlinear_bellman.transformed_n_step_q_learning,
                  tx_pair=tx_pair,
                  n=self.n,
              )
          )
      )
      actual_td = transformed_n_step_q_learning(
          self.q_tm1,
          self.a_tm1,
          self.target_q_t,
          self.a_t,
          self.r_t,
          self.discount_t,
      )
      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 TransformedNStepQLearningTest obj's underlying member method code with example usage: obj.test_transformed_q_lambda_batch(tx_pair, td_index) without return types