code: def test_jax_Join():
      a = matrix("a")
      b = matrix("b")
  
      x = ptb.join(0, a, b)
      x_fg = FunctionGraph([a, b], [x])
      compare_jax_and_py(
          x_fg,
          [
              np.c_[[1.0, 2.0, 3.0]].astype(config.floatX),
              np.c_[[4.0, 5.0, 6.0]].astype(config.floatX),
          ],
      )
      compare_jax_and_py(
          x_fg,
          [
              np.c_[[1.0, 2.0, 3.0]].astype(config.floatX),
              np.c_[[4.0, 5.0]].astype(config.floatX),
          ],
      )
  
      x = ptb.join(1, a, b)
      x_fg = FunctionGraph([a, b], [x])
      compare_jax_and_py(
          x_fg,
          [
              np.c_[[1.0, 2.0, 3.0]].astype(config.floatX),
              np.c_[[4.0, 5.0, 6.0]].astype(config.floatX),
          ],
      )
      compare_jax_and_py(
          x_fg,
          [
              np.c_[[1.0, 2.0], [3.0, 4.0]].astype(config.floatX),
              np.c_[[5.0, 6.0]].astype(config.floatX),
          ],
      )
  
imports: import numpy
  import pytest
  from pytensor.compile import get_mode
  import jax.errors
  import pytensor
  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.type import iscalar, matrix, scalar, vector
  from tests.link.jax.test_basic import compare_jax_and_py
  from tests.tensor.test_basic import TestAlloc
label_desc: Use image node_id 7 for calling a global function with example usage: test_jax_Join() without return types