code: def testSetBanditPolicyType(self):
      dims = (10, 1)
      bandit_policy_spec = (
          policy_utilities.create_bandit_policy_type_tensor_spec(dims)
      )
      info = policy_utilities.set_bandit_policy_type(
          None, bandit_policy_spec
      )
      self.assertIsInstance(info, policy_utilities.PolicyInfo)
      self.assertIsInstance(
          info.bandit_policy_type, tensor_spec.BoundedTensorSpec
      )
      self.assertEqual(info.bandit_policy_type.shape, dims)
      self.assertEqual(info.bandit_policy_type.dtype, tf.int32)
      # Set to tensor.
      input_tensor = tf.fill(dims, value=_GREEDY)
      info = policy_utilities.set_bandit_policy_type(info, input_tensor)
      self.assertIsInstance(info.bandit_policy_type, tf.Tensor)
      self.assertEqual(
          info.bandit_policy_type.shape, input_tensor.shape
      )
      expected = [[_GREEDY] for _ in range(dims[0])]
      self.assertAllEqual(info.bandit_policy_type, expected)
  
file_path: agents/tf_agents/bandits/policies/policy_utilities_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import tensorflow
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.utils import test_utils
label_desc: Use image node_id 5 for calling the PolicyUtilitiesTest obj's underlying member method code with example usage: obj.testSetBanditPolicyType() without return types