code: def test_jax_ifelse():
      true_vals = np.r_[1, 2, 3]
      false_vals = np.r_[-1, -2, -3]
  
      x = ifelse(np.array(True), true_vals, false_vals)
      x_fg = FunctionGraph([], [x])
  
      compare_jax_and_py(x_fg, [])
  
      a = dscalar("a")
      a.tag.test_value = np.array(0.2, dtype=config.floatX)
      x = ifelse(a < 0.5, true_vals, false_vals)
      x_fg = FunctionGraph([a], [x])  # I.e. False
  
      compare_jax_and_py(x_fg, [get_test_value(i) for i in x_fg.inputs])
  
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 6 for calling a global function with example usage: test_jax_ifelse() without return types