code: def run():
      # set a deterministic seed so that both ranks generate the same random tensor
      Tensor.manual_seed(42)
  
      rank = getenv("RANK")
  
      # loop 3 times to make sure it works with the jit
      for _ in range(3):
          # create a tensor to send
          t = (
              Tensor.zeros(SIZE, SIZE)
              if rank != 0
              else Tensor.ones(SIZE, SIZE)
          )
          t2 = allreduce_jit(t.contiguous().realize())
          assert np.allclose(
              np.ones((SIZE, SIZE)), t2.numpy()
          ), f"{t2.numpy()} wasn't ones"
  
      # reset jit
      allreduce_jit.cnt = 0
  
      # test uneven chunk sizes
      for _ in range(3):
          # create a tensor to send
          t = (
              Tensor.ones(SIZE_2, SIZE_2, SIZE_2)
              if rank == 0
              else Tensor.zeros(SIZE_2, SIZE_2, SIZE_2)
          )
          t2 = allreduce_jit(t.contiguous().realize())
          assert np.allclose(
              np.ones((SIZE_2, SIZE_2, SIZE_2)), t2.numpy()
          ), f"{t2.numpy()} wasn't ones"
  
      print(f"rank {rank} passed")
  
imports: from extra import dist
  from tinygrad.jit import TinyJit
  from extra.dist import collectives
  from tinygrad.helpers import CI, getenv
  from tinygrad.tensor import Tensor
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: run() without return types