code: def test_jax_specify_shape():
      in_pt = pt.matrix("in")
      x = pt.specify_shape(in_pt, (4, None))
      x_fg = FunctionGraph([in_pt], [x])
      compare_jax_and_py(x_fg, [np.ones((4, 5)).astype(config.floatX)])
  
      # When used to assert two arrays have similar shapes
      in_pt = pt.matrix("in")
      shape_pt = pt.matrix("shape")
      x = pt.specify_shape(in_pt, shape_pt.shape)
      x_fg = FunctionGraph([in_pt, shape_pt], [x])
      compare_jax_and_py(
          x_fg,
          [
              np.ones((4, 5)).astype(config.floatX),
              np.ones((4, 5)).astype(config.floatX),
          ],
      )
  
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 2 for calling a global function with example usage: test_jax_specify_shape() without return types