code: def test_checkpoint(a, b, c):
      torch.manual_seed(0)
      t1 = torch.rand(a, 1, c).requires_grad_(True)
      t2 = torch.rand(1, b, c).requires_grad_(True)
  
      r1 = LogSemiring.dot(t1, t2)
      r2 = CheckpointSemiring(LogSemiring).dot(t1, t2)
      r2 = CheckpointShardSemiring(LogSemiring, 2).dot(t2, t1)
      assert torch.isclose(r1, r2).all()
  
      (a1, b1) = torch.autograd.grad(r1.sum(), (t1, t2))
      (a2, b2) = torch.autograd.grad(r2.sum(), (t1, t2))
  
      assert torch.isclose(a1, a2).all()
      assert torch.isclose(b1, b2).all()
  
imports: import torch
  from hypothesis import given
  from hypothesis.strategies import integers
  from torch_struct import LogSemiring, CheckpointSemiring, CheckpointShardSemiring, KMaxSemiring, MaxSemiring, StdSemiring
label_desc: Use image node_id 2 for calling a global function with example usage: test_checkpoint(a, b, c) without return types