code: def compare_jax_and_py(
      fgraph: FunctionGraph,
      test_inputs: Iterable,
      assert_fn: Optional[Callable] = None,
      must_be_device_array: bool = True,
      jax_mode=jax_mode,
      py_mode=py_mode,
  ):
      """Function to compare python graph output and jax compiled output for testing equality
  
      In the tests below computational graphs are defined in PyTensor. These graphs are then passed to
      this function which then compiles the graphs in both jax and python, runs the calculation
      in both and checks if the results are the same
  
      Parameters
      ----------
      fgraph: FunctionGraph
          PyTensor function Graph object
      test_inputs: iter
          Numerical inputs for testing the function graph
      assert_fn: func, opt
          Assert function used to check for equality between python and jax. If not
          provided uses np.testing.assert_allclose
      must_be_device_array: Bool
          Checks for instance of jax.interpreters.xla.DeviceArray. For testing purposes
          if this device array is found it indicates if the result was computed by jax
  
      Returns
      -------
      jax_res
  
      """
      if assert_fn is None:
          assert_fn = partial(np.testing.assert_allclose, rtol=1e-4)
  
      fn_inputs = [
          i for i in fgraph.inputs if not isinstance(i, SharedVariable)
      ]
      pytensor_jax_fn = function(
          fn_inputs, fgraph.outputs, mode=jax_mode
      )
      jax_res = pytensor_jax_fn(*test_inputs)
  
      if must_be_device_array:
          if isinstance(jax_res, list):
              assert all(isinstance(res, jax.Array) for res in jax_res)
          else:
              assert isinstance(
                  jax_res, jax.interpreters.xla.DeviceArray
              )
  
      pytensor_py_fn = function(fn_inputs, fgraph.outputs, mode=py_mode)
      py_res = pytensor_py_fn(*test_inputs)
  
      if len(fgraph.outputs) > 1:
          for j, p in zip(jax_res, py_res):
              assert_fn(j, p)
      else:
          assert_fn(jax_res, py_res)
  
      return pytensor_jax_fn, jax_res
  
imports: from collections.abc import Iterable
  from functools import partial
  from typing import Callable, Optional
  import numpy
  import pytest
  from pytensor.compile.function import function
  from pytensor.compile.mode import get_mode
  from pytensor.compile.sharedvalue import SharedVariable, shared
  from pytensor.configdefaults import config
  from pytensor.graph.basic import Apply
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.op import Op, get_test_value
  from pytensor.ifelse import ifelse
  from pytensor.raise_op import assert_op
  from pytensor.tensor.type import dscalar, scalar, vector
label_desc: Use image node_id 2 for calling a global function with example usage: compare_jax_and_py(fgraph, test_inputs, assert_fn, must_be_device_array, jax_mode, py_mode) and returns: pytensor_jax_fn,  jax_res