code: def test_ExtractDiag_exhaustive(k, axis1, axis2, reverse_axis):
      if reverse_axis:
          axis1, axis2 = axis2, axis1
  
      x = pt.tensor4("x")
      x_shape = (2, 3, 4, 5)
      x_test = np.arange(np.prod(x_shape)).reshape(x_shape)
      out = pt.diagonal(x, k, axis1, axis2)
      numba_fn = numba_funcify(out.owner.op, out.owner)
      np.testing.assert_allclose(
          numba_fn(x_test), np.diagonal(x_test, k, axis1, axis2)
      )
  
imports: import numpy
  import pytest
  import pytensor.scalar
  import pytensor.tensor
  import pytensor.tensor.basic
  from pytensor import config, function
  from pytensor.compile import get_mode
  from pytensor.compile.sharedvalue import SharedVariable
  from pytensor.graph.basic import Constant
  from pytensor.graph.fg import FunctionGraph
  from pytensor.scalar import Add
  from pytensor.tensor.shape import Unbroadcast
  from tests.link.numba.test_basic import compare_numba_and_py, compare_shape_dtype, set_test_value
  from tests.tensor.test_basic import TestAlloc
  from pytensor.link.numba.dispatch import numba_funcify
label_desc: Use image node_id 14 for calling a global function with example usage: test_ExtractDiag_exhaustive(k, axis1, axis2, reverse_axis) without return types