code: def CheckpointShardSemiring(cls, max_size, min_size=0):
      class _Check(torch.autograd.Function):
          @staticmethod
          def forward(ctx, a, b):
              ctx.save_for_backward(a, b)
              size = matmul_size(a, b)
              return accumulate_(
                  a,
                  b,
                  size,
                  lambda a, b: cls.matmul(a, b),
                  preserve=len(size),
                  step=max_size // (b.shape[-2] * a.shape[-1]) + 2,
              )
  
          @staticmethod
          def backward(ctx, grad_output):
              a, b = ctx.saved_tensors
              grad_a, grad_b = unaccumulate_(
                  a,
                  b,
                  grad_output,
                  len(grad_output.shape),
                  lambda a, b: cls.matmul(a, b),
                  step=max_size // (b.shape[-2] * a.shape[-1]) + 2,
              )
              return grad_a, grad_b
  
      class _CheckpointSemiring(cls):
          @staticmethod
          def matmul(a, b):
              size = torch.tensor(
                  [max(i, j) for i, j in zip(a.shape, b.shape)]
              ).prod()
              if size < min_size:
                  return cls.matmul(a, b)
              else:
                  return _Check.apply(a, b)
  
      return _CheckpointSemiring
  
file_path: pytorch-struct/torch_struct/semirings/checkpoint.py
imports: import torch
label_desc: Use image node_id 4 for calling a global function with example usage: CheckpointShardSemiring(cls, max_size, min_size) and returns: _CheckpointSemiring, accumulate_, grad_a,  grad_b, cls, _Check