code: def jax_funcify_FillDiagonalOffset(op, **kwargs):
      # def filldiagonaloffset(a, val, offset):
      #     height, width = a.shape
      #
      #     if offset >= 0:
      #         start = offset
      #         num_of_step = min(min(width, height), width - offset)
      #     else:
      #         start = -offset * a.shape[1]
      #         num_of_step = min(min(width, height), height + offset)
      #
      #     step = a.shape[1] + 1
      #     end = start + step * num_of_step
      #     a.flat[start:end:step] = val
      #
      #     return a
      #
      # return filldiagonaloffset
  
      raise NotImplementedError("flatiter not implemented in JAX")
  
file_path: pytensor/pytensor/link/jax/dispatch/extra_ops.py
imports: import warnings
  import jax.numpy
  from pytensor.link.jax.dispatch.basic import jax_funcify
  from pytensor.tensor.extra_ops import Bartlett, CumOp, FillDiagonal, FillDiagonalOffset, RavelMultiIndex, Repeat, Unique, UnravelIndex
label_desc: Use image node_id 8 for calling a global function with example usage: jax_funcify_FillDiagonalOffset(op) without return types