code: def jax_funcify_choice(op):
      """JAX implementation of `ChoiceRV`."""
  
      def sample_fn(rng, size, dtype, *parameters):
          rng_key = rng["jax_state"]
          rng_key, sampling_key = jax.random.split(rng_key, 2)
          (a, p, replace) = parameters
          smpl_value = jax.random.choice(
              sampling_key, a, size, replace, p
          )
          rng["jax_state"] = rng_key
          return (rng, smpl_value)
  
      return sample_fn
  
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 13 for calling a global function with example usage: jax_funcify_choice(op) and returns: sample_fn, rng,  smpl_value