class_docstr: Tests ActionDistribution classes.
code: def _stability_test(
      self,
      distribution_cls,
      network_output_shape,
      fw,
      sess=None,
      bounds=None,
      extra_kwargs=None,
  ):
      extreme_values = [
          0.0,
          float(LARGE_INTEGER),
          -float(LARGE_INTEGER),
          1.1e-34,
          1.1e34,
          -1.1e-34,
          -1.1e34,
          SMALL_NUMBER,
          -SMALL_NUMBER,
      ]
      inputs = np.zeros(shape=network_output_shape, dtype=np.float32)
      for batch_item in range(network_output_shape[0]):
          for num in range(len(inputs[batch_item]) // 2):
              inputs[batch_item][num] = np.random.choice(extreme_values)
          else:
              # For Gaussians, the second half of the vector is
              # log standard deviations, and should therefore be
              # the log of a positive number >= 1.
              inputs[batch_item][num] = np.log(
                  max(1, np.random.choice((extreme_values)))
              )
  
      dist = distribution_cls(inputs, {}, **(extra_kwargs or {}))
      for _ in range(100):
          sample = dist.sample()
          if fw != "tf":
              sample_check = sample.numpy()
          else:
              sample_check = sess.run(sample)
          assert not np.any(np.isnan(sample_check))
          assert np.all(np.isfinite(sample_check))
          if bounds:
              assert np.min(sample_check) >= bounds[0]
              assert np.max(sample_check) <= bounds[1]
              # Make sure bounds make sense and are actually also being
              # sampled.
              if isinstance(bounds[0], int):
                  assert isinstance(bounds[1], int)
                  assert bounds[0] in sample_check
                  assert bounds[1] in sample_check
          logp = dist.logp(sample)
          if fw != "tf":
              logp_check = logp.numpy()
          else:
              logp_check = sess.run(logp)
          assert not np.any(np.isnan(logp_check))
          assert np.all(np.isfinite(logp_check))
  
imports: from functools import partial
  from gymnasium.spaces import Box, Dict, Tuple
  import numpy
  from scipy.stats import beta, norm
  import tree
  import unittest
  from ray.rllib.models.tf.tf_action_dist import Beta, Categorical, DiagGaussian, GumbelSoftmax, MultiActionDistribution, MultiCategorical, SquashedGaussian
  from ray.rllib.models.torch.torch_action_dist import TorchBeta, TorchCategorical, TorchDiagGaussian, TorchMultiActionDistribution, TorchMultiCategorical, TorchSquashedGaussian
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.numpy import MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT, softmax, SMALL_NUMBER, LARGE_INTEGER
  from ray.rllib.utils.test_utils import check, framework_iterator
label_desc: Use image node_id 2 for calling the TestActionDistributions obj's underlying member method code with example usage: obj._stability_test(distribution_cls, network_output_shape, fw, sess, bounds, extra_kwargs) without return types