code: def test_categorical_sample():
      """distribution test"""
      logits = torch.from_numpy(np.random.randn(5))
      dist = CategoricalProbabilityDistribution(
          logits=logits,
          action_space=spaces.Discrete(5),
          temperature=1.0,
      )
      assert dist.sample().numpy().ndim == 0
      assert dist.deterministic_sample().numpy().ndim == 0
  
      logits = torch.from_numpy(np.random.randn(100, 5))
      dist = CategoricalProbabilityDistribution(
          logits=logits,
          action_space=spaces.Discrete(5),
          temperature=1.0,
      )
      assert dist.sample().numpy().ndim == 1
      assert dist.deterministic_sample().numpy().ndim == 1
      assert dist.sample().numpy().shape == (100,)
      assert dist.deterministic_sample().numpy().shape == (100,)
  
      logits = torch.from_numpy(np.random.randn(100, 8, 5))
      dist = CategoricalProbabilityDistribution(
          logits=logits,
          action_space=spaces.Discrete(5),
          temperature=1.0,
      )
      assert dist.sample().numpy().ndim == 2
      assert dist.deterministic_sample().numpy().ndim == 2
      assert dist.sample().numpy().shape == (100, 8)
      assert dist.deterministic_sample().numpy().shape == (100, 8)
  
file_path: maze/maze/test/distributions/test_categorical.py
imports: import numpy
  import torch
  from gym import spaces
  from maze.distributions.categorical import CategoricalProbabilityDistribution
label_desc: Use image node_id 1 for calling a global function with example usage: test_categorical_sample() without return types