code: def jax_sample_fn_loc_scale(op):
      """JAX implementation of random variables in the loc-scale families.
  
      JAX only implements the standard version of random variables in the
      loc-scale family. We thus need to translate and rescale the results
      manually.
  
      """
      name = op.name
      jax_op = getattr(jax.random, name)
  
      def sample_fn(rng, size, dtype, *parameters):
          rng_key = rng["jax_state"]
          rng_key, sampling_key = jax.random.split(rng_key, 2)
          loc, scale = parameters
          sample = loc + jax_op(sampling_key, size, dtype) * scale
          rng["jax_state"] = rng_key
          return (rng, sample)
  
      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 7 for calling a global function with example usage: jax_sample_fn_loc_scale(op) and returns: sample_fn, rng,  sample