class_docstr: Tests Distribution classes.
class_docstr_tok: ["Tests","Distribution","classes","."]
code: def test_categorical(self):
      batch_size = 10000
      num_categories = 4
      sample_shape = 2
  
      # Create categorical distribution with n categories.
      logits = np.random.randn(batch_size, num_categories)
      probs = torch.from_numpy(softmax(logits)).float()
      logits = torch.from_numpy(logits).float()
  
      # check stability against skewed inputs
      check_stability(TorchCategorical, sample_input={"logits": logits})
      check_stability(
          TorchCategorical,
          sample_input={"probs": logits},
          constraints={"probs": "probability"},
      )
  
      dist_with_logits = TorchCategorical(logits=logits)
      dist_with_probs = TorchCategorical(probs=probs)
  
      samples = dist_with_logits.sample(sample_shape=(sample_shape,))
  
      # check shape of samples
      self.assertEqual(
          samples.shape,
          (
              sample_shape,
              batch_size,
          ),
      )
      self.assertEqual(samples.dtype, torch.int64)
      # check that none of the samples are nan
      self.assertFalse(torch.isnan(samples).any())
      # check that all samples are in the range of the number of categories
      self.assertTrue((samples >= 0).all())
      self.assertTrue((samples < num_categories).all())
  
      # resample to remove the first batch dim
      samples = dist_with_logits.sample()
      # check that the two distributions are the same
      check(
          dist_with_logits.logp(samples), dist_with_probs.logp(samples)
      )
  
      # check logp values
      expected = (
          probs.log().gather(dim=-1, index=samples.view(-1, 1)).view(-1)
      )
      check(dist_with_logits.logp(samples), expected)
  
      # check entropy
      expected = -(probs * probs.log()).sum(dim=-1)
      check(dist_with_logits.entropy(), expected)
  
      # check kl
      probs2 = softmax(np.random.randn(batch_size, num_categories))
      probs2 = torch.from_numpy(probs2).float()
      dist2 = TorchCategorical(probs=probs2)
      expected = (probs * (probs / probs2).log()).sum(dim=-1)
      check(dist_with_probs.kl(dist2), expected)
  
file_path: ray/python/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 2 for calling the TestDistributions obj's underlying member method code with example usage: obj.test_categorical() without return types