code: def check_stability(
      dist_class, *, sample_input=None, constraints=None
  ):
      max_tries = 100
      extreme_values = [
          0.0,
          float(LARGE_INTEGER),
          -float(LARGE_INTEGER),
          1.1e-34,
          1.1e34,
          -1.1e-34,
          -1.1e34,
          SMALL_NUMBER,
          -SMALL_NUMBER,
      ]
  
      input_kwargs = copy(sample_input)
      for key, array in input_kwargs.items():
          arr_sampled = np.random.choice(
              extreme_values, replace=True, size=array.shape
          )
          input_kwargs[key] = torch.from_numpy(arr_sampled).float()
  
          if constraints:
              constraint = constraints.get(key, None)
              if constraint:
                  if constraint == "positive_not_inf":
                      input_kwargs[key] = torch.minimum(
                          SMALL_NUMBER
                          + torch.log(1 + torch.exp(input_kwargs[key])),
                          torch.tensor([LARGE_INTEGER]),
                      )
                  elif constraint == "probability":
                      input_kwargs[key] = torch.softmax(
                          input_kwargs[key], dim=-1
                      )
  
      dist = dist_class(**input_kwargs)
      for _ in range(max_tries):
          sample = dist.sample()
  
          assert not torch.isnan(sample).any()
          assert torch.all(torch.isfinite(sample))
  
          logp = dist.logp(sample)
          assert not torch.isnan(logp).any()
          assert torch.all(torch.isfinite(logp))
  
file_path: ray/rllib/models/tests/test_distributions.py
imports: from copy import copy
  import numpy
  import unittest
  import math
  from ray.rllib.models.torch.torch_distributions import TorchCategorical, TorchDiagGaussian, TorchDeterministic
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.numpy import softmax, SMALL_NUMBER, LARGE_INTEGER
  from ray.rllib.utils.test_utils import check
label_desc: Use image node_id 1 for calling a global function with example usage: check_stability(dist_class) without return types