code: def test_ar2():
      batch, N, C, H = 3, 5, 2, 5
      layer = 1
  
      def t(a):
          return tuple((t.transpose(0, 1) for t in a))
  
      init = (torch.zeros(batch, layer, H),)
  
      class AR(torch.nn.Module):
          def __init__(self, sparse=True):
              super().__init__()
              self.sparse = sparse
              self.rnn = torch.nn.RNN(H, H, batch_first=True)
              self.proj = torch.nn.Linear(H, C)
              if sparse:
                  self.embed = torch.nn.Embedding(C, H)
              else:
                  self.embed = torch.nn.Linear(C, H)
  
          def forward(self, inputs, state):
              if not self.sparse and inputs.dim() == 2:
                  inputs = torch.nn.functional.one_hot(
                      inputs, C
                  ).float()
              inputs = self.embed(inputs)
              out, state = self.rnn(inputs, t(state)[0])
              out = self.proj(out)
              return out, t((state,))
  
      dist2 = Autoregressive(
          AR(sparse=False), init, C, N, normalize=False
      )
      path, _, _ = dist2.greedy_tempmax(1)
  
      dist = Autoregressive(AR(), init, C, N, normalize=False)
  
      path, scores, _ = dist.greedy_max()
  
      assert torch.isclose(
          scores, dist.log_prob(path.unsqueeze(0))
      ).all()
      scores = dist._beam_max(7)
      path = dist.beam_topk(7)
  
      a = torch.tensor([[1, 2, 5], [3, 4, 6]])
      print(KMaxSemiring(2).sparse_sum(a))
  
      import itertools
  
      for i in range(5):
          print(dist.log_prob(path)[i])
          print(scores[i])
          print(path[i, 0])
  
      print(path.shape, scores.shape)
      assert torch.isclose(scores, dist.log_prob(path)).all()
  
      v = torch.tensor(list(itertools.product([0, 1], repeat=N)))
  
      v = v.unsqueeze(1).expand(v.shape[0], batch, N)
      all_scores = dist.log_prob(v, sparse=True)
      best, ind = torch.max(all_scores, dim=0)
      assert torch.isclose(scores[0, 0], best[0]).all()
  
      print(v[ind[0], 0].shape, path[0, 0].shape)
      assert (
          torch.nn.functional.one_hot(v[ind, 0], C) == path[0, 0].long()
      ).all()
  
      dist = Autoregressive(AR(), init, C, N, normalize=True)
      path = dist.sample((7,))
  
      init = (
          torch.zeros(batch, layer, H),
          torch.zeros(batch, layer, H),
      )
  
      class AR(torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.rnn = torch.nn.LSTM(H, H, batch_first=True)
              self.proj = torch.nn.Linear(H, C)
              self.embed = torch.nn.Embedding(C, H)
  
          def forward(self, inputs, state):
              inputs = self.embed(inputs)
              out, state = self.rnn(inputs, t(state))
              out = self.proj(out)
              return out, t(state)
  
      dist = Autoregressive(AR(), init, C, N)
      dist.greedy_max()
      dist.beam_topk(5)
  
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 4 for calling a global function with example usage: test_ar2() and returns: tuple, out,  t, out,  t