class_docstr: Tests for policy_gradient op.
class_docstr_tok: ["Tests","for","policy_gradient","op","."]
code: def testLoss(self, multi_actions):
      batch_size = 2
      self._setUp_loss(multi_actions, batch_size=batch_size)
      self.assertEqual(
          self._loss.get_shape(), tf.TensorShape(batch_size)
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  from six.moves import xrange
  import tensorflow.compat.v1
  import tensorflow_probability
  import tree
  from trfl import policy_gradient_ops
label_desc: Use image node_id 2 for calling the PolicyGradientLossTest obj's underlying member method code with example usage: obj.testLoss(multi_actions) without return types