docstr: Tests for a full batch.
docstr_tok: ["Tests","for","a","full","batch","."]
code: def test_policy_gradient_loss_batch(self):
      """Tests for a full batch."""
      policy_gradient_loss = self.variant(
          jax.vmap(policy_gradients.policy_gradient_loss)
      )
      # Test outputs.
      actual = policy_gradient_loss(
          self.logits, self.actions, self.advantages, self.weights
      )
      np.testing.assert_allclose(self.expected, 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 PolicyGradientLossTest obj's underlying member method code with example usage: obj.test_policy_gradient_loss_batch() without return types