code: def assert_shape_argument_jax_compatible(shape):
      """Assert whether the current node can be JIT-compiled by JAX.
  
      JAX can JIT-compile functions with a `shape` or `size` argument if it is
      given a concrete value, i.e. either a constant or the shape of any traced
      value.
  
      """
      shape_op = shape.owner.op
      if not isinstance(shape_op, (Shape, Shape_i, JAXShapeTuple)):
          raise NotImplementedError(SHAPE_NOT_COMPATIBLE)
  
file_path: pytensor/pytensor/link/jax/dispatch/shape.py
imports: import jax.numpy
  from pytensor.graph import Constant
  from pytensor.graph.basic import Apply
  from pytensor.graph.op import Op
  from pytensor.link.jax.dispatch.basic import jax_funcify
  from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape, Unbroadcast
  from pytensor.tensor.type import TensorType
label_desc: Use image node_id 2 for calling a global function with example usage: assert_shape_argument_jax_compatible(shape) without return types