code: def setUp(self):
      super().setUp()
      self.n = 2
  
      self.q_tm1 = np.array(
          [
              [[1.1, 2.1], [-1.1, 1.1], [3.1, -3.1]],
              [[2.1, 3.1], [-1.1, 0.1], [-2.1, -1.1]],
          ],
          dtype=np.float32,
      )
      self.a_tm1 = np.array([[0, 1, 0], [1, 0, 0]], dtype=np.int32)
      self.discount_t = np.array(
          [[0.0, 0.89, 0.85], [0.88, 1.0, 0.83]], dtype=np.float32
      )
      self.r_t = np.array(
          [[-1.3, -1.3, 2.3], [1.3, 5.3, -3.3]], dtype=np.float32
      )
      self.target_q_t = np.array(
          [
              [[1.2, 2.2], [-1.2, 0.2], [2.2, -1.2]],
              [[4.2, 2.2], [1.2, 1.2], [-1.2, -2.2]],
          ],
          dtype=np.float32,
      )
      self.a_t = np.array([[0, 1, 0], [1, 0, 0]], dtype=np.int32)
  
      self.expected_td = np.array(
          [
              [
                  [-2.4, 1.3112999, 1.0700002],
                  [3.9199996, 2.104, -2.196],
              ],
              [
                  [-1.9329091, 0.9564189, -0.7853615],
                  [-0.9021418, 1.1716722, 0.2713145],
              ],
              [
                  [-1.6178751, 0.85600746, -0.75762916],
                  [-0.87689304, 0.6246443, 0.3598088],
              ],
              [
                  [-2.178451, 1.02313, -0.593768],
                  [-0.415362, 1.790864, -0.120749],
              ],
          ],
          dtype=np.float32,
      )
  
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 1 for calling the TransformedNStepQLearningTest obj's underlying member method code with example usage: obj.setUp() without return types