code: def get_context_dims_from_spec(
      context_spec: types.NestedTensorSpec,
      accepts_per_arm_features: bool,
  ) -> Tuple[int, int]:
      """Returns the global and per-arm context dimensions.
  
      If the policy accepts per-arm features, this function returns the tuple of
      the global and per-arm context dimension. Otherwise, it returns the (global)
      context dim and zero.
  
      Args:
        context_spec: A nest of tensor specs, containing the observation spec.
        accepts_per_arm_features: (bool) Whether the context_spec is for a policy
          that accepts per-arm features.
  
      Returns: A 2-tuple of ints, the global and per-arm context dimension. If the
        policy does not accept per-arm features, the per-arm context dim is 0.
      """
      if accepts_per_arm_features:
          global_context_dim = context_spec[
              GLOBAL_FEATURE_KEY
          ].shape.as_list()[0]
          arm_context_dim = context_spec[
              PER_ARM_FEATURE_KEY
          ].shape.as_list()[1]
      else:
          assert hasattr(context_spec, "shape")
          spec_shape = context_spec.shape.as_list()
          global_context_dim = spec_shape[0] if spec_shape else 1
          arm_context_dim = 0
      return global_context_dim, arm_context_dim
  
file_path: agents/tf_agents/specs/bandit_spec_utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  from typing import Optional, Tuple
  import tensorflow
  from tf_agents.specs import tensor_spec
  from tf_agents.typing import types
label_desc: Use image node_id 2 for calling a global function with example usage: get_context_dims_from_spec(context_spec, accepts_per_arm_features) and returns: global_context_dim,  arm_context_dim