code: def _calculate_log_probability(outer_dims, action_spec):
      """Helper function for calculating log prob of a uniform distribution.
  
      Each item in the returned tensor will be equal to:
      |action_spec.shape| * log_prob_of_each_component_of_action_spec.
  
      Note that this method expects the same value for all outer_dims because
      we're sampling uniformly from the same distribution for each batch row.
  
      Args:
        outer_dims: TensorShape.
        action_spec: BoundedTensorSpec.
  
      Returns:
        A tensor of type float32 with shape outer_dims.
      """
      # Equivalent of what a tfp.distribution.Categorical would return.
      if action_spec.dtype.is_integer:
          log_prob = -tf.math.log(
              action_spec.maximum - action_spec.minimum + 1.0
          )
      # Equivalent of what a tfp.distribution.Uniform would return.
      else:
          log_prob = -tf.math.log(
              action_spec.maximum - action_spec.minimum
          )
  
      # Note that log_prob may be a vector. We first reduce it to a scalar, and then
      # adjust by the number of times that vector is repeated in action_spec.
      log_prob = tf.reduce_sum(log_prob) * (
          action_spec.shape.num_elements()
          / log_prob.shape.num_elements()
      )
      # Regardless of the type of the action, the log_prob should be float32.
      return tf.cast(tf.fill(outer_dims, log_prob), tf.float32)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import cast
  import tensorflow
  from tf_agents.distributions import masked
  from tf_agents.policies import tf_policy
  from tf_agents.policies import utils
  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.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 1 for calling a global function with example usage: _calculate_log_probability(outer_dims, action_spec) and returns: tf