code: def test_multi_categorical_sample():
      """distribution test"""
      logits = torch.from_numpy(np.random.randn(8))
      dist = MultiCategoricalProbabilityDistribution(
          logits=logits,
          action_space=spaces.MultiDiscrete((3, 5)),
          temperature=1.0,
      )
      for multi_sample in [dist.sample(), dist.deterministic_sample()]:
          for sample in multi_sample:
              assert sample.numpy().ndim == 0
  
      logits = torch.from_numpy(np.random.randn(100, 8))
      dist = MultiCategoricalProbabilityDistribution(
          logits=logits,
          action_space=spaces.MultiDiscrete((3, 5)),
          temperature=1.0,
      )
      for multi_sample in [dist.sample(), dist.deterministic_sample()]:
          assert multi_sample.numpy().ndim == 2
          assert multi_sample.numpy().shape == (100, 2)
  
      logits = torch.from_numpy(np.random.randn(100, 12, 8))
      dist = MultiCategoricalProbabilityDistribution(
          logits=logits,
          action_space=spaces.MultiDiscrete((3, 5)),
          temperature=1.0,
      )
      for multi_sample in [dist.sample(), dist.deterministic_sample()]:
          assert multi_sample.numpy().ndim == 3
          assert multi_sample.numpy().shape == (100, 12, 2)
  
file_path: maze/maze/test/distributions/test_multi_categorical.py
imports: import numpy
  import torch
  from gym import spaces
  from maze.distributions.multi_categorical import MultiCategoricalProbabilityDistribution
label_desc: Use image node_id 1 for calling a global function with example usage: test_multi_categorical_sample() without return types