docstr: Tests for a full batch.
docstr_tok: ["Tests","for","a","full","batch","."]
code: def test_qpg_loss_batch(self, compile_fn, place_fn):
      """Tests for a full batch."""
      # Vmap and optionally compile.
      qpg_loss = compile_fn(policy_gradients.qpg_loss)
  
      # Optionally convert to device array.
      policy_logits, q_values = jax.tree_map(
          place_fn, (self.policy_logits, self.q_values)
      )
      # Test outputs.
      actual = qpg_loss(policy_logits, q_values)
      np.testing.assert_allclose(
          self.expected_policy_loss, actual, atol=1e-4
      )
  
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 distributions
  from rlax._src import policy_gradients
label_desc: Use image node_id 2 for calling the QPGLossTest obj's underlying member method code with example usage: obj.test_qpg_loss_batch(compile_fn, place_fn) without return types