code: def test_custom_action_distribution(self):
      class Model:
          pass
  
      ray.init(
          object_store_memory=1000 * 1024 * 1024,
          ignore_reinit_error=True,
      )  # otherwise fails sometimes locally
      # registration
      ModelCatalog.register_custom_action_dist(
          "test", CustomActionDistribution
      )
      action_space = Box(0, 1, shape=(5, 3), dtype=np.float32)
  
      # test retrieving it
      model_config = MODEL_DEFAULTS.copy()
      model_config["custom_action_dist"] = "test"
      dist_cls, param_shape = ModelCatalog.get_action_dist(
          action_space, model_config
      )
      self.assertEqual(str(dist_cls), str(CustomActionDistribution))
      self.assertEqual(param_shape, action_space.shape)
  
      # test the class works as a distribution
      dist_input = tf1.placeholder(tf.float32, (None,) + param_shape)
      model = Model()
      model.model_config = model_config
      dist = dist_cls(dist_input, model=model)
      self.assertEqual(dist.sample().shape[1:], dist_input.shape[1:])
      self.assertIsInstance(dist.sample(), tf.Tensor)
      with self.assertRaises(NotImplementedError):
          dist.entropy()
  
      # test passing the options to it
      model_config["custom_model_config"].update({"output_dim": (3,)})
      dist_cls, param_shape = ModelCatalog.get_action_dist(
          action_space, model_config
      )
      self.assertEqual(param_shape, (3,))
      dist_input = tf1.placeholder(tf.float32, (None,) + param_shape)
      model.model_config = model_config
      dist = dist_cls(dist_input, model=model)
      self.assertEqual(dist.sample().shape[1:], dist_input.shape[1:])
      self.assertIsInstance(dist.sample(), tf.Tensor)
      with self.assertRaises(NotImplementedError):
          dist.entropy()
  
imports: from functools import partial
  from gymnasium.spaces import Box, Dict, Discrete, Tuple
  import numpy
  import unittest
  import ray
  from ray.rllib.models import ActionDistribution, ModelCatalog, MODEL_DEFAULTS
  from ray.rllib.models.preprocessors import Preprocessor, TupleFlatteningPreprocessor
  from ray.rllib.models.tf.tf_action_dist import MultiActionDistribution, TFActionDistribution
  from ray.rllib.models.tf.tf_modelv2 import TFModelV2
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.spaces.space_utils import get_dummy_batch_for_space
  from ray.rllib.utils.torch_utils import convert_to_torch_tensor
label_desc: Use image node_id 4 for calling the TestModelCatalog obj's underlying member method code with example usage: obj.test_custom_action_distribution() without return types