code: def _get_dist_class(
      policy: Policy,
      config: AlgorithmConfigDict,
      action_space: gym.spaces.Space,
  ) -> Type[TFActionDistribution]:
      """Helper function to return a dist class based on config and action space.
  
      Args:
          policy: The policy for which to return the action
              dist class.
          config: The Algorithm's config dict.
          action_space (gym.spaces.Space): The action space used.
  
      Returns:
          Type[TFActionDistribution]: A TF distribution class.
      """
      if (
          hasattr(policy, "dist_class")
          and policy.dist_class is not None
      ):
          return policy.dist_class
      elif config["model"].get("custom_action_dist"):
          action_dist_class, _ = ModelCatalog.get_action_dist(
              action_space, config["model"], framework="tf"
          )
          return action_dist_class
      elif isinstance(action_space, Discrete):
          return Categorical
      elif isinstance(action_space, Simplex):
          return Dirichlet
      else:
          assert isinstance(action_space, Box)
          if config["normalize_actions"]:
              return (
                  SquashedGaussian
                  if not config["_use_beta_distribution"]
                  else Beta
              )
          else:
              return DiagGaussian
  
imports: import copy
  import gymnasium
  from gymnasium.spaces import Box, Discrete
  from functools import partial
  import logging
  from typing import Dict, List, Optional, Tuple, Type, Union
  import ray
  import ray.experimental.tf_utils
  from ray.rllib.algorithms.dqn.dqn_tf_policy import postprocess_nstep_and_prio, PRIO_WEIGHTS
  from ray.rllib.algorithms.sac.sac_tf_model import SACTFModel
  from ray.rllib.algorithms.sac.sac_torch_model import SACTorchModel
  from ray.rllib.evaluation.episode import Episode
  from ray.rllib.models import ModelCatalog
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.tf.tf_action_dist import Beta, Categorical, DiagGaussian, Dirichlet, SquashedGaussian, TFActionDistribution
  from ray.rllib.policy.policy import Policy
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.tf_mixins import TargetNetworkMixin
  from ray.rllib.policy.tf_policy_template import build_tf_policy
  from ray.rllib.utils.error import UnsupportedSpaceException
  from ray.rllib.utils.framework import get_variable, try_import_tf
  from ray.rllib.utils.spaces.simplex import Simplex
  from ray.rllib.utils.tf_utils import huber_loss, make_tf_callable
  from ray.rllib.utils.typing import AgentID, LocalOptimizer, ModelGradients, TensorType, AlgorithmConfigDict
label_desc: Use image node_id 3 for calling a global function with example usage: _get_dist_class(policy, config, action_space) and returns: policy, action_dist_class, Categorical, Dirichlet, unknown, DiagGaussian