code: def set_bandit_policy_type(
      info: Optional[Union[PolicyInfo, PerArmPolicyInfo]],
      bandit_policy_type: types.Tensor,
  ) -> Union[PolicyInfo, PerArmPolicyInfo]:
      """Sets the InfoFields.BANDIT_POLICY_TYPE on info to bandit_policy_type.
  
      If policy `info` does not support InfoFields.BANDIT_POLICY_TYPE, this method
      returns `info` as-is (without any modification).
  
      Args:
        info: Policy info on which to set bandit policy type.
        bandit_policy_type: Tensor containing BanditPolicyType enums or TensorSpec
          from `create_bandit_policy_type_tensor_spec()`.
  
      Returns:
        Policy info with modified field (if possible).
      """
      if not info:
          return PolicyInfo(bandit_policy_type=bandit_policy_type)
      fields = getattr(info, "_fields", None)
      if fields is not None and InfoFields.BANDIT_POLICY_TYPE in fields:
          return info._replace(bandit_policy_type=bandit_policy_type)
      try:
          info[InfoFields.BANDIT_POLICY_TYPE] = bandit_policy_type
      except TypeError:
          pass
      return info
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from typing import Optional, Sequence, Text, Union
  import tensorflow
  from tf_agents.specs import bandit_spec_utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 7 for calling a global function with example usage: set_bandit_policy_type(info, bandit_policy_type) and returns: info, PolicyInfo, info