code: def jax_funcify_ExtractDiag(op, **kwargs):
      offset = op.offset
      axis1 = op.axis1
      axis2 = op.axis2
  
      def extract_diag(x, offset=offset, axis1=axis1, axis2=axis2):
          return jnp.diagonal(
              x, offset=offset, axis1=axis1, axis2=axis2
          )
  
      return extract_diag
  
file_path: pytensor/pytensor/link/jax/dispatch/tensor_basic.py
imports: import warnings
  import jax.numpy
  import numpy
  from pytensor.graph.basic import Constant
  from pytensor.link.jax.dispatch.basic import jax_funcify
  from pytensor.tensor import get_vector_length
  from pytensor.tensor.basic import Alloc, AllocEmpty, ARange, ExtractDiag, Eye, Join, MakeVector, ScalarFromTensor, Split, TensorFromScalar, Tri, get_underlying_scalar_constant_value
  from pytensor.tensor.exceptions import NotScalarConstantError
  from pytensor.tensor.shape import Shape_i
label_desc: Use image node_id 6 for calling a global function with example usage: jax_funcify_ExtractDiag(op) and returns: extract_diag, jnp