code: def unaccumulate_(a, b, grad_output, preserve, fn, step=10000):
      slices = []
      total = 1
      size = grad_output.shape[:preserve]
      for s in grad_output.shape[:preserve]:
          slices.append(slice(s))
          total *= s
  
      if step > total:
          with torch.enable_grad():
              a_in = a.clone().requires_grad_(True)
              b_in = b.clone().requires_grad_(True)
              q = fn(a, b)
          ag, bg = torch.autograd.grad(q, (a, b), grad_output)
          return ag, bg
  
      a2 = a.expand(*size[:-2], a.shape[-2], a.shape[-1])
      b2 = b.expand(*size[:-2], b.shape[-2], b.shape[-1])
      a2 = a2.contiguous().view(-1, a.shape[-2], a.shape[-1])
      b2 = b2.contiguous().view(-1, b.shape[-2], b.shape[-1])
  
      a_grad = a2.clone().fill_(0)
      b_grad = b2.clone().fill_(0)
  
      grad_output = grad_output.view(-1, a.shape[-2], b.shape[-1])
      for p in range(0, grad_output.shape[0], step):
          with torch.enable_grad():
              a_in = a2[p : p + step].clone().requires_grad_(True)
              b_in = b2[p : p + step].clone().requires_grad_(True)
              q = fn(a_in, b_in)
          ag, bg = torch.autograd.grad(
              q, (a_in, b_in), grad_output[p : p + step]
          )
          a_grad[p : p + step] += ag
          b_grad[p : p + step] += bg
  
      a_grad = a_grad.view(*size[:-2], a.shape[-2], a.shape[-1])
      b_grad = b_grad.view(*size[:-2], b.shape[-2], b.shape[-1])
      a_ones = ones(a)
      b_ones = ones(b)
      f1, f2 = a_grad.sum(a_ones, keepdim=True), b_grad.sum(
          b_ones, keepdim=True
      )
      return f1, f2
  
file_path: pytorch-struct/torch_struct/semirings/checkpoint.py
imports: import torch
label_desc: Use image node_id 8 for calling a global function with example usage: unaccumulate_(a, b, grad_output, preserve, fn, step) and returns: f1,  f2, ag,  bg