class_docstr: Autoregressive sequence model utilizing beam search.
  
  * batch_shape -> Given by initializer
  * event_shape -> N x T sequence of choices
  
  Parameters:
      model (AutoregressiveModel): A lazily computed autoregressive model.
      init (tuple of tensors, batch_shape x ...): initial state of autoregressive model.
      n_classes (int): number of classes in each time step
      n_length (int): max length of sequence
class_docstr_tok: ["Autoregressive","sequence","model","utilizing","beam","search",".","*","batch_shape","-",">","Given","by","initializer","*","event_shape","-",">","N","x","T","sequence","of","choices","Parameters",":","model","(","AutoregressiveModel",")",":","A","lazily","computed","autoregressive","model",".","init","(","tuple","of","tensors",",","batch_shape","x","...",")",":","initial","state","of","autoregressive","model",".","n_classes","(","int",")",":","number","of","classes","in","each","time","step","n_length","(","int",")",":","max","length","of","sequence"]
code: def _beam_search(self, semiring, gumbel=False):
      beam = semiring.fill(
          torch.zeros(
              (semiring.size(),) + self.batch_shape, device=self.device
          ),
          torch.tensor(True),
          semiring.one,
      )
      ssize = semiring.size()
  
      def take(state, indices):
          return tuple(
              (
                  s.contiguous()[
                      (
                          indices * self.batch_shape[0]
                          + torch.arange(
                              self.batch_shape[0], device=self.device
                          ).unsqueeze(0)
                      )
                      .contiguous()
                      .view(-1)
                  ]
                  for s in state
              )
          )
  
      tokens = (
          torch.zeros((ssize * self.batch_shape[0]))
          .long()
          .fill_(self.start_class)
      )
      state = tuple(
          (
              unwrap(i.unsqueeze(0).expand((ssize,) + i.shape))
              for i in self.init
          )
      )
  
      # Beam Search
      all_beams = []
      all_logits = []
      for t in range(0, self.n_length):
          logits, state = self.model(unwrap(tokens).unsqueeze(1), state)
          b2, n2, c2 = logits.shape
          assert (
              (b2 == ssize * self.batch_shape[0])
              and (n2 == 1)
              and (c2 == self.n_classes)
          ), "Model should return logits of shape `batch x N x C` "
          for s in state:
              assert (
                  s.shape[0] == ssize * self.batch_shape[0]
              ), "Model should return state tuple with shapes `batch x ...` "
          logits = wrap(logits.squeeze(1), ssize)
          if gumbel:
              logits = logits + torch.distributions.Gumbel(
                  0.0, 1.0
              ).sample(logits.shape)
          if self.normalize:
              logits = logits.log_softmax(-1)
          all_logits.append(logits)
          ex_beam = beam.unsqueeze(-1) + logits
          ex_beam.requires_grad_(True)
          all_beams.append(ex_beam)
          beam, (positions, tokens) = semiring.sparse_sum(ex_beam)
          state = take(state, positions)
  
      # Back pointers
      v = beam
      all_m = []
      for k in range(v.shape[0]):
          obj = v[k].sum(dim=0)
          marg = torch.autograd.grad(
              obj,
              all_beams,
              create_graph=True,
              only_inputs=True,
              allow_unused=False,
          )
          marg = torch.stack(marg, dim=2)
          all_m.append(marg.sum(0))
      return (
          torch.stack(all_m, dim=0),
          v,
          torch.stack(all_logits, dim=2),
      )
  
file_path: pytorch-struct/torch_struct/autoregressive.py
imports: import torch
  from .semirings import MaxSemiring, KMaxSemiring, TempMax
  from torch.distributions.distribution import Distribution
label_desc: Use image node_id 3 for calling the Autoregressive obj's underlying member method code with example usage: obj._beam_search(semiring, gumbel) and returns: torch,  v,  torch, tuple