code: def _setup_pgops_mock(
      sequence_length=3, batch_size=2, num_policies=3
  ):
      """Setup ops using mock distribution for numerical tests."""
      t, b = sequence_length, batch_size
      policies = [
          MockDistribution((t, b), i + 1) for i in xrange(num_policies)
      ]
      actions = [
          tf.constant(np.arange(t * b).reshape((t, b)))
          for i in xrange(num_policies)
      ]
      if num_policies == 1:
          policies, actions = policies[0], actions[0]
      entropy_scale_op = lambda policies: len(nest.flatten(policies))
      return policies, actions, entropy_scale_op
  
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 1 for calling a global function with example usage: _setup_pgops_mock(sequence_length, batch_size, num_policies) and returns: policies,  actions,  entropy_scale_op