code: def test_extra_ops():
      a = matrix("a")
      a.tag.test_value = np.arange(6, dtype=config.floatX).reshape(
          (3, 2)
      )
  
      out = pt_extra_ops.cumsum(a, axis=0)
      fgraph = FunctionGraph([a], [out])
      compare_jax_and_py(
          fgraph, [get_test_value(i) for i in fgraph.inputs]
      )
  
      out = pt_extra_ops.cumprod(a, axis=1)
      fgraph = FunctionGraph([a], [out])
      compare_jax_and_py(
          fgraph, [get_test_value(i) for i in fgraph.inputs]
      )
  
      out = pt_extra_ops.diff(a, n=2, axis=1)
      fgraph = FunctionGraph([a], [out])
      compare_jax_and_py(
          fgraph, [get_test_value(i) for i in fgraph.inputs]
      )
  
      out = pt_extra_ops.repeat(a, (3, 3), axis=1)
      fgraph = FunctionGraph([a], [out])
      compare_jax_and_py(
          fgraph, [get_test_value(i) for i in fgraph.inputs]
      )
  
      c = ptb.as_tensor(5)
  
      out = pt_extra_ops.fill_diagonal(a, c)
      fgraph = FunctionGraph([a], [out])
      compare_jax_and_py(
          fgraph, [get_test_value(i) for i in fgraph.inputs]
      )
  
      with pytest.raises(NotImplementedError):
          out = pt_extra_ops.fill_diagonal_offset(a, c, c)
          fgraph = FunctionGraph([a], [out])
          compare_jax_and_py(
              fgraph, [get_test_value(i) for i in fgraph.inputs]
          )
  
      with pytest.raises(NotImplementedError):
          out = pt_extra_ops.Unique(axis=1)(a)
          fgraph = FunctionGraph([a], [out])
          compare_jax_and_py(
              fgraph, [get_test_value(i) for i in fgraph.inputs]
          )
  
      indices = np.arange(np.prod((3, 4)))
      out = pt_extra_ops.unravel_index(indices, (3, 4), order="C")
      fgraph = FunctionGraph([], out)
      compare_jax_and_py(
          fgraph,
          [get_test_value(i) for i in fgraph.inputs],
          must_be_device_array=False,
      )
  
file_path: pytensor/tests/link/jax/test_extra_ops.py
imports: import numpy
  import pytest
  from packaging.version import parse
  import pytensor.tensor.basic
  from pytensor.configdefaults import config
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.op import get_test_value
  from pytensor.tensor import extra_ops
  from pytensor.tensor.type import matrix
  from tests.link.jax.test_basic import compare_jax_and_py
label_desc: Use image node_id 2 for calling a global function with example usage: test_extra_ops() without return types