class_docstr: Test the `pixel_control_loss` op.
docstr: Compute loss for given observations, actions, values, tensor discount.
docstr_tok: ["Compute","loss","for","given","observations",",","actions",",","values",",","tensor","discount","."]
code: def testPixelControlLossTensorDiscount(self):
      """Compute loss for given observations, actions, values, tensor discount."""
      zero_discount = np.zeros((1,))
      non_zero_discount = self.discount * np.ones(self.seq_length - 1)
      discount = np.concatenate(
          [zero_discount, non_zero_discount], axis=0
      )
      loss_fn = self.variant(
          functools.partial(
              losses.pixel_control_loss, cell_size=self.cell_size
          )
      )
      loss = loss_fn(
          self.observations, self.actions, self.action_values, discount
      )
      loss = jnp.sum(loss)
      np.testing.assert_allclose(loss, self.error_term, rtol=1e-3)
  
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 losses
label_desc: Use image node_id 3 for calling the PixelControlLossTest obj's underlying member method code with example usage: obj.testPixelControlLossTensorDiscount() without return types