class_docstr: Tests Distribution classes.
class_docstr_tok: ["Tests","Distribution","classes","."]
code: def test_diag_gaussian(self):
      batch_size = 128
      ndim = 4
      sample_shape = 100000
  
      loc = np.random.randn(batch_size, ndim)
      scale = np.exp(np.random.randn(batch_size, ndim))
  
      loc_tens = torch.from_numpy(loc).float()
      scale_tens = torch.from_numpy(scale).float()
  
      dist = TorchDiagGaussian(loc=loc_tens, scale=scale_tens)
      sample = dist.sample(sample_shape=(sample_shape,))
  
      # check shape of samples
      self.assertEqual(sample.shape, (sample_shape, batch_size, ndim))
      self.assertEqual(sample.dtype, torch.float32)
      # check that none of the samples are nan
      self.assertFalse(torch.isnan(sample).any())
  
      # check that mean and std are approximately correct
      check(sample.mean(0), loc, decimals=1)
      check(sample.std(0), scale, decimals=1)
  
      # check logp values
      expected = (
          -0.5 * ((sample - loc_tens) / scale_tens).pow(2).sum(-1)
          + -0.5 * ndim * math.log(2 * math.pi)
          - scale_tens.log().sum(-1)
      )
      check(dist.logp(sample), expected)
  
      # check entropy
      expected = 0.5 * ndim * (
          1 + math.log(2 * math.pi)
      ) + scale_tens.log().sum(-1)
      check(dist.entropy(), expected)
  
      # check kl
      loc2 = torch.from_numpy(np.random.randn(batch_size, ndim)).float()
      scale2 = torch.from_numpy(
          np.exp(np.random.randn(batch_size, ndim))
      )
      dist2 = TorchDiagGaussian(loc=loc2, scale=scale2)
      expected = (
          scale2.log()
          - scale_tens.log()
          + (scale_tens.pow(2) + (loc_tens - loc2).pow(2))
          / (2 * scale2.pow(2))
          - 0.5
      ).sum(-1)
      check(dist.kl(dist2), expected, decimals=4)
  
      # check rsample
      loc_tens.requires_grad = True
      scale_tens.requires_grad = True
      dist = TorchDiagGaussian(loc=2 * loc_tens, scale=2 * scale_tens)
      sample1 = dist.rsample()
      sample2 = dist.sample()
  
      self.assertRaises(
          RuntimeError,
          lambda: sample2.mean().backward(retain_graph=True),
      )
      sample1.mean().backward(retain_graph=True)
  
      # check stablity against skewed inputs
      check_stability(
          TorchDiagGaussian,
          sample_input={"loc": loc_tens, "scale": scale_tens},
          constraints={"scale": "positive_not_inf"},
      )
  
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 3 for calling the TestDistributions obj's underlying member method code with example usage: obj.test_diag_gaussian() without return types