code: def test_categorical_shape():
      probs = torch.ones(3, 2) / 2
      d = dist.Categorical(probs)
      assert d.batch_shape == (3,)
      assert d.event_shape == ()
      assert d.shape() == (3,)
      assert d.sample().size() == d.shape()
  
file_path: pyro/tests/distributions/test_shapes.py
imports: import torch
  import pyro.distributions
label_desc: Use image node_id 1 for calling a global function with example usage: test_categorical_shape() without return types