code: def test_autoregressive(data, seed):
      n_classes = 2
      n_length = 5
      batch = 3
      BATCH = 3
  
      values = torch.rand(batch, n_length, n_classes)
  
      values2 = (
          values.unsqueeze(-1)
          .expand(batch, n_length, n_classes, n_classes)
          .clone()
      )
      values2[:, 0, :, :] = -1e9
      values2[
          :, 0, torch.arange(n_classes), torch.arange(n_classes)
      ] = values[:, 0]
  
      init = (torch.zeros(batch, 5).long(),)
  
      class Model(torch.nn.Module):
          def forward(self, inputs, state):
              if inputs.shape[1] == 1:
                  (state,) = state
                  in_batch, hidden = state.shape
                  t = state[0, 0]
                  batch = values.shape[0]
                  x = (
                      values[:, t, :]
                      .unsqueeze(0)
                      .expand(in_batch // batch, batch, n_classes)
                  )
                  state = (state + 1,)
                  return (
                      x.contiguous().view(in_batch, 1, n_classes),
                      state,
                  )
              else:
                  return (
                      torch.cat(
                          (values, torch.zeros(BATCH, 1, n_classes)),
                          dim=1,
                      ),
                      state,
                  )
  
      auto = Autoregressive(
          Model(), init, n_classes, n_length, normalize=False
      )
      v, _, _ = auto.greedy_max()
      batch, n, c = v.shape
      assert n == n_length
      assert c == n_classes
  
      assert (v == LinearChainCRF(values2).argmax.sum(-1)).all()
      crf = LinearChainCRF(values2)
      v2 = auto.beam_topk(K=5)
  
      assert (v2.nonzero() == crf.topk(5).sum(-1).nonzero()).all()
      assert (v2[0] == LinearChainCRF(values2).argmax.sum(-1)).all()
  
      print(auto.log_prob(v.unsqueeze(0)))
      print(crf.struct().score(crf.argmax, values2))
      assert (
          torch.isclose(
              auto.log_prob(v.unsqueeze(0)),
              crf.struct().score(crf.argmax, values2),
          )
      ).all()
      assert auto.sample((7,)).shape == (7, batch, n_length, n_classes)
  
      assert auto.sample_without_replacement((7,)).shape == (
          7,
          batch,
          n_length,
          n_classes,
      )
  
file_path: pytorch-struct/tests/test_distributions.py
imports: from torch_struct import LinearChainCRF, Autoregressive, KMaxSemiring, LogSemiring
  import torch
  from hypothesis import given, settings
  from hypothesis.strategies import integers, data, sampled_from
  from .extensions import test_lookup
label_desc: Use image node_id 3 for calling a global function with example usage: test_autoregressive(data, seed) and returns: x,  state, torch,  state