code: def assert_size_argument_jax_compatible(node):
      """Assert whether the current node can be JIT-compiled by JAX.
  
      JAX can JIT-compile `jax.random` functions when the `size` argument
      is a concrete value, i.e. either a constant or the shape of any
      traced value.
  
      """
      size = node.inputs[1]
      size_node = size.owner
      if (size_node is not None) and (
          not isinstance(size_node.op, (Shape, Shape_i, JAXShapeTuple))
      ):
          raise NotImplementedError(SIZE_NOT_COMPATIBLE)
  
imports: from functools import singledispatch
  import jax
  import numpy
  from numpy.random import Generator, RandomState
  from numpy.random.bit_generator import _coerce_to_uint32_array
  import pytensor.tensor.random.basic
  from pytensor.link.jax.dispatch.basic import jax_funcify, jax_typify
  from pytensor.link.jax.dispatch.shape import JAXShapeTuple
  from pytensor.tensor.shape import Shape, Shape_i
label_desc: Use image node_id 1 for calling a global function with example usage: assert_size_argument_jax_compatible(node) without return types