code: def KMaxSemiring(k):
      "Implements the k-max semiring (kmax, +, [-inf, -inf..], [0, -inf, ...])."
  
      class KMaxSemiring(_BaseLog):
          zero = torch.tensor([-1e5 for i in range(k)])
          one = torch.tensor([0 if i == 0 else -1e5 for i in range(k)])
  
          @staticmethod
          def size():
              return k
  
          @classmethod
          def convert(cls, orig_potentials):
              potentials = torch.zeros(
                  (k,) + orig_potentials.shape,
                  dtype=orig_potentials.dtype,
                  device=orig_potentials.device,
              )
              potentials = cls.fill(
                  potentials, torch.tensor(True), cls.zero
              )
              potentials[0] = orig_potentials
              return potentials
  
          @staticmethod
          def unconvert(potentials):
              return potentials[0]
  
          @staticmethod
          def sum(xs, dim=-1):
              if dim == -1:
                  xs = xs.permute(tuple(range(1, xs.dim())) + (0,))
                  xs = xs.contiguous().view(xs.shape[:-2] + (-1,))
                  xs = torch.topk(xs, k, dim=-1)[0]
                  xs = xs.permute(
                      (xs.dim() - 1,) + tuple(range(0, xs.dim() - 1))
                  )
                  assert xs.shape[0] == k
                  return xs
              assert False
  
          @staticmethod
          def sparse_sum(xs, dim=-1):
              if dim == -1:
                  xs = xs.permute(tuple(range(1, xs.dim())) + (0,))
                  xs = xs.contiguous().view(xs.shape[:-2] + (-1,))
                  xs, xs2 = torch.topk(xs, k, dim=-1)
                  xs = xs.permute(
                      (xs.dim() - 1,) + tuple(range(0, xs.dim() - 1))
                  )
                  xs2 = xs2.permute(
                      (xs.dim() - 1,) + tuple(range(0, xs.dim() - 1))
                  )
                  assert xs.shape[0] == k
                  return xs, (xs2 % k, xs2 // k)
              assert False
  
          @staticmethod
          def mul(a, b):
              a = a.view((k, 1) + a.shape[1:])
              b = b.view((1, k) + b.shape[1:])
              c = a + b
              c = c.contiguous().view((k * k,) + c.shape[2:])
              ret = torch.topk(c, k, 0)[0]
              assert ret.shape[0] == k
              return ret
  
      return KMaxSemiring
  
imports: import torch
label_desc: Use image node_id 2 for calling a global function with example usage: KMaxSemiring(k) and returns: KMaxSemiring, k, potentials, potentials, ret, xs, xs,  unknown,  unknown