code: def numba_funcify_ExtractDiag(op, node, **kwargs):
      view = op.view
      axis1, axis2, offset = op.axis1, op.axis2, op.offset
  
      if node.inputs[0].type.ndim == 2:
  
          @numba_basic.numba_njit(inline="always")
          def extract_diag(x):
              out = np.diag(x, k=offset)
  
              if not view:
                  out = out.copy()
  
              return out
  
      else:
          axis1p1 = axis1 + 1
          axis2p1 = axis2 + 1
          leading_dims = (slice(None),) * axis1
          middle_dims = (slice(None),) * (axis2 - axis1 - 1)
  
          @numba_basic.numba_njit(inline="always")
          def extract_diag(x):
              if offset >= 0:
                  diag_len = min(
                      x.shape[axis1], max(0, x.shape[axis2] - offset)
                  )
              else:
                  diag_len = min(
                      x.shape[axis2], max(0, x.shape[axis1] + offset)
                  )
              base_shape = (
                  x.shape[:axis1]
                  + x.shape[axis1p1:axis2]
                  + x.shape[axis2p1:]
              )
              out_shape = base_shape + (diag_len,)
              out = np.empty(out_shape)
  
              for i in range(diag_len):
                  if offset >= 0:
                      new_entry = x[
                          leading_dims
                          + (i,)
                          + middle_dims
                          + (i + offset,)
                      ]
                  else:
                      new_entry = x[
                          leading_dims
                          + (i - offset,)
                          + middle_dims
                          + (i,)
                      ]
                  out[..., i] = new_entry
              return out
  
      return extract_diag
  
file_path: pytensor/pytensor/link/numba/dispatch/tensor_basic.py
imports: from textwrap import indent
  import numpy
  from pytensor.link.numba.dispatch import basic
  from pytensor.link.numba.dispatch.basic import create_tuple_string, numba_funcify
  from pytensor.link.utils import compile_function_src, unique_name_generator
  from pytensor.tensor.basic import Alloc, AllocEmpty, ARange, ExtractDiag, Eye, Join, MakeVector, ScalarFromTensor, Split, TensorFromScalar
  from pytensor.tensor.shape import Unbroadcast
label_desc: Use image node_id 6 for calling a global function with example usage: numba_funcify_ExtractDiag(op, node) and returns: extract_diag, out, out