code: def jacobian(func, input):
      output = func(input)
  
      ji = input.numpy().reshape(-1).shape[-1]
      jo = output.numpy().reshape(-1).shape[-1]
      J = np.zeros((jo, ji), dtype=np.float32)
  
      for o in range(jo):
          input.grad = None
          output = func(input)
  
          # tinygrad doesn't support slicing, tiny-hack to select
          # the needed scalar an backpropagate only through it
          o_scalar = (
              Tensor(mask_like(output.numpy(), o, 1.0))
              .mul(output)
              .sum()
          )
          o_scalar.backward()
  
          for i, grad in enumerate(input.grad.numpy().reshape(-1)):
              J[o, i] = grad
      return J
  
imports: import numpy
  from tinygrad.tensor import Tensor
label_desc: Use image node_id 2 for calling a global function with example usage: jacobian(func, input) and returns: J