code: def test_jax_Reshape_shape_graph_input():
      a = vector("a")
      shape_pt = iscalar("b")
      x = reshape(a, (shape_pt, shape_pt))
      x_fg = FunctionGraph([a, shape_pt], [x])
      compare_jax_and_py(
          x_fg, [np.r_[1.0, 2.0, 3.0, 4.0].astype(config.floatX), 2]
      )
  
file_path: pytensor/tests/link/jax/test_shape.py
imports: import numpy
  import pytest
  import pytensor.tensor
  from pytensor.compile.ops import DeepCopyOp, ViewOp
  from pytensor.configdefaults import config
  from pytensor.graph.fg import FunctionGraph
  from pytensor.tensor.shape import Shape, Shape_i, Unbroadcast, reshape
  from pytensor.tensor.type import iscalar, vector
  from tests.link.jax.test_basic import compare_jax_and_py
label_desc: Use image node_id 5 for calling a global function with example usage: test_jax_Reshape_shape_graph_input() without return types