code: def test_extra_ops_omni():
      a = matrix("a")
      a.tag.test_value = np.arange(6, dtype=config.floatX).reshape(
          (3, 2)
      )
  
      # This function also cannot take symbolic input.
      c = ptb.as_tensor(5)
      out = pt_extra_ops.bartlett(c)
      fgraph = FunctionGraph([], [out])
      compare_jax_and_py(
          fgraph, [get_test_value(i) for i in fgraph.inputs]
      )
  
      multi_index = np.unravel_index(np.arange(np.prod((3, 4))), (3, 4))
      out = pt_extra_ops.ravel_multi_index(multi_index, (3, 4))
      fgraph = FunctionGraph([], [out])
      compare_jax_and_py(
          fgraph,
          [get_test_value(i) for i in fgraph.inputs],
          must_be_device_array=False,
      )
  
      # The inputs are "concrete", yet it still has problems?
      out = pt_extra_ops.Unique()(
          ptb.as_tensor(
              np.arange(6, dtype=config.floatX).reshape((3, 2))
          )
      )
      fgraph = FunctionGraph([], [out])
      compare_jax_and_py(fgraph, [])
  
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 3 for calling a global function with example usage: test_extra_ops_omni() without return types