code: def testVariableWeightsDefaultTemp(
      self, observation_shape, batch_size, weights
  ):
      observation_spec = tensor_spec.TensorSpec(
          shape=observation_shape,
          dtype=tf.float32,
          name="observation_spec",
      )
      time_step_spec = time_step.time_step_spec(observation_spec)
  
      weights = tf.compat.v2.Variable(weights, dtype=tf.float32)
      action_spec = tensor_spec.BoundedTensorSpec(
          shape=(),
          dtype=tf.int32,
          minimum=0,
          maximum=tf.compat.dimension_value(weights.shape[0]) - 1,
          name="action",
      )
      policy = categorical_policy.CategoricalPolicy(
          weights, time_step_spec, action_spec
      )
      observation_step = _get_dummy_observation_step(
          observation_shape, batch_size
      )
      action_time_step = policy.action(observation_step)
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.assertAllEqual(
          action_time_step.action.shape.as_list(), [batch_size]
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.policies import categorical_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.utils import test_utils
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 2 for calling the CategoricalPolicyTest obj's underlying member method code with example usage: obj.testVariableWeightsDefaultTemp(observation_shape, batch_size, weights) without return types