text: import torch
  from .semirings import MaxSemiring, KMaxSemiring, TempMax
  from torch.distributions.distribution import Distribution
  
  
  class AutoregressiveModel(torch.nn.Module):
      """
      User should implement as their favorite RNN / Transformer / etc.
      """
  
      def forward(self, inputs, state=None):
          r"""
          Compute the logits for all tokens in a batched sequence :math:`p(y_{t+1}, ... y_{T}| y_1 \ldots t)`
  
          Parameters:
              inputs (batch_size x N x C ): next tokens to update representation
              state (tuple of batch_size x ...): everything needed for conditioning.
  
          Retuns:
              logits (*batch_size x C*): next set of logits.
  
              state (*tuple of batch_size x ...*): next set of logits.
          """
          pass
  
  
  def wrap(state, ssize):
      return state.contiguous().view(ssize, -1, *state.shape[1:])
  
  
  def unwrap(state):
      return state.contiguous().view(-1, *state.shape[2:])
  
  
  class Autoregressive(Distribution):
      """
      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
      """
  
      def __init__(
          self,
          model,
          initial_state,
          n_classes,
          n_length,
          normalize=True,
          start_class=0,
          end_class=None,
      ):
          self.model = model
          self.init = initial_state
          self.n_length = n_length
          self.n_classes = n_classes
          self.start_class = start_class
          self.normalize = normalize
          event_shape = (n_length, n_classes)
          batch_shape = initial_state[0].shape[:1]
          self.device = initial_state[0].device
          super().__init__(batch_shape=batch_shape, event_shape=event_shape)
  
      def log_prob(self, value, sparse=False):
          """
          Compute log probability over values :math:`p(z)`.
  
          Parameters:
              value (tensor): One-hot events (*sample_shape x batch_shape x N*)
  
          Returns:
              log_probs (*sample_shape x batch_shape*)
          """
          value = value.long()
          if not sparse:
              sample, batch_shape, n_length, n_classes = value.shape
              value = (
                  (value * torch.arange(n_classes).view(1, 1, n_classes)).sum(-1).long()
              )
          else:
              sample, batch_shape, n_length = value.shape
  
          value = torch.cat(
              [
                  torch.zeros(sample, batch_shape, 1, device=value.device)
                  .fill_(self.start_class)
                  .long(),
                  value,
              ],
              dim=2,
          )
          value = unwrap(value)
          state = tuple(
              (unwrap(i.unsqueeze(0).expand((sample,) + i.shape)) for i in self.init)
          )
  
          logits, _ = self.model(value, state)
          b2, n2, c2 = logits.shape
          assert (
              (b2 == sample * batch_shape)
              and (n2 == n_length + 1)
              and (c2 == self.n_classes)
          ), "Model should return logits of shape `batch x N x C` "
  
          if self.normalize:
              log_probs = logits.log_softmax(-1)
          else:
              log_probs = logits
  
          scores = log_probs[:, :-1].gather(2, value[:, 1:].unsqueeze(-1)).sum(-1).sum(-1)
          return wrap(scores, sample)
  
      def _beam_search(self, semiring, gumbel=False):
          beam = semiring.one_(
              torch.zeros((semiring.size(),) + self.batch_shape, device=self.device)
          )
          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)
  
      def greedy_max(self):
          """
          Compute "argmax" using greedy search.
  
          Returns:
              greedy_path (*batch x N x C*)
              greedy_max (*batch*)
              logits (*batch x N x C*)
          """
          a, b, c = self._beam_search(MaxSemiring)
          return a.squeeze(0), b.squeeze(0), c.squeeze(0)
  
      def greedy_tempmax(self, alpha):
          """
          Compute differentiable scheduled sampling using greedy search.
  
          Based on:
  
          * Differentiable Scheduled Sampling for Credit Assignment :cite:`goyal2017differentiable`
  
          Returns:
              greedy_path (*batch x N x C*)
              greedy_max (*batch*)
              logits (*batch x N x C*)
          """
          a, b, c = self._beam_search(TempMax(alpha), alpha)
          return a.squeeze(0), b.squeeze(0), c.squeeze(0)
  
      def beam_topk(self, K):
          """
          Compute "top-k" using beam search
  
          Returns:
              paths (*K x batch x N x C*)
  
          """
          return self._beam_search(KMaxSemiring(K))[0]
  
      def _beam_max(self, K):
          return self._beam_search(KMaxSemiring(K))[1]
  
      def sample_without_replacement(self, sample_shape=torch.Size()):
          """
          Compute sampling without replacement using Gumbel trick.
  
          Based on:
  
          * Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for
                 Sampling Sequences Without Replacement :cite:`DBLP:journals/corr/abs-1903-06059`
  
          Parameters:
              sample_shape (torch.Size): batch_size
  
          Returns:
              paths (*K x batch x N x C*)
  
          """
          K = sample_shape[0]
          return self._beam_search(KMaxSemiring(K), gumbel=True)[0]
  
      def sample(self, sample_shape=torch.Size()):
          r"""
          Compute structured samples from the distribution :math:`z \sim p(z)`.
  
          Parameters:
              sample_shape (torch.Size): number of samples
  
          Returns:
              samples (*sample_shape x batch_shape x event_shape*)
          """
          sample_shape = sample_shape[0]
          state = tuple(
              (
                  unwrap(i.unsqueeze(0).expand((sample_shape,) + i.shape))
                  for i in self.init
              )
          )
          all_tokens = []
          tokens = (
              torch.zeros((sample_shape * self.batch_shape[0]))
              .long()
              .fill_(self.start_class)
          )
          for t in range(0, self.n_length):
              logits, state = self.model(tokens.unsqueeze(-1), state)
              logits = logits.squeeze(1)
              tokens = torch.distributions.Categorical(logits=logits).sample((1,))[0]
              all_tokens.append(tokens)
          v = wrap(torch.stack(all_tokens, dim=1), sample_shape)
          return torch.nn.functional.one_hot(v, self.n_classes)
  