code: def test_squashed_gaussian_kl():
      """distribution test"""
      logits_0 = torch.from_numpy(np.random.randn(100, 8, 6))
      logits_1 = torch.from_numpy(np.random.randn(100, 8, 6))
  
      dist_0 = SquashedGaussianProbabilityDistribution(
          logits=logits_0,
          action_space=Box(
              low=0.0, high=10.0, shape=(3,), dtype=np.float64
          ),
          temperature=1.0,
      )
      dist_1 = SquashedGaussianProbabilityDistribution(
          logits=logits_1,
          action_space=Box(
              low=0.0, high=10.0, shape=(3,), dtype=np.float64
          ),
          temperature=1.0,
      )
  
      assert dist_0.kl(dist_1).numpy().ndim == 2
      assert dist_0.kl(dist_1).numpy().shape == (100, 8)
  
file_path: maze/maze/test/distributions/test_squashed_gaussian.py
imports: import numpy
  import torch
  from gym.spaces import Box
  from maze.distributions.squashed_gaussian import SquashedGaussianProbabilityDistribution
label_desc: Use image node_id 4 for calling a global function with example usage: test_squashed_gaussian_kl() without return types