code: def enumerate(semiring, edge, lengths=None):
      model = torch_struct.LinearChain(semiring)
      semiring = semiring
      ssize = semiring.size()
      edge, batch, N, C, lengths = model._check_potentials(
          edge, lengths
      )
      chains = [
          [
              (
                  [c],
                  semiring.fill(
                      torch.zeros(ssize, batch),
                      torch.tensor(True),
                      semiring.one,
                  ),
              )
              for c in range(C)
          ]
      ]
  
      enum_lengths = torch.LongTensor(lengths.shape)
      for n in range(1, N):
          new_chains = []
          for chain, score in chains[-1]:
              for c in range(C):
                  new_chains.append(
                      (
                          chain + [c],
                          semiring.mul(
                              score, edge[:, :, n - 1, c, chain[-1]]
                          ),
                      )
                  )
          chains.append(new_chains)
  
          for b in range(lengths.shape[0]):
              if lengths[b] == n + 1:
                  enum_lengths[b] = len(new_chains)
  
      edges = model.to_parts(
          torch.stack([torch.tensor(c) for (c, _) in chains[-1]]), C
      )
      # Sum out non-batch
      a = torch.einsum("ancd,sbncd->sbancd", edges.float(), edge)
      a = semiring.prod(a.view(*a.shape[:3] + (-1,)), dim=3)
      a = semiring.sum(a, dim=2)
      ret = semiring.sum(
          torch.stack([s for (_, s) in chains[-1]], dim=1), dim=1
      )
      assert torch.isclose(a, ret).all(), "%s %s" % (a, ret)
  
      edges = torch.zeros(len(chains[-1]), batch, N - 1, C, C)
      for b in range(lengths.shape[0]):
          edges[
              : enum_lengths[b], b, : lengths[b] - 1
          ] = model.to_parts(
              torch.stack(
                  [torch.tensor(c) for (c, _) in chains[lengths[b] - 1]]
              ),
              C,
          )
  
      return (
          semiring.unconvert(ret),
          [s for (_, s) in chains[-1]],
          edges,
          enum_lengths,
      )
  
imports: import torch_struct
  import torch
  from torch_struct import LogSemiring
  import itertools
  from hypothesis.strategies import integers, composite, floats
  from hypothesis.extra.numpy import arrays
  import numpy
label_desc: Use image node_id 2 for calling the LinearChainTest obj's underlying member method code with example usage: obj.enumerate(semiring, edge, lengths) and returns: semiring,  unknown,  edges,  enum_lengths