code: def test_random_custom_implementation():
      """We can register a JAX implementation for user-defined `RandomVariable`s"""
  
      class CustomRV(RandomVariable):
          name = "non-existent"
          ndim_supp = 0
          ndims_params = []
          dtype = "floatX"
  
          def __call__(self, size=None, **kwargs):
              return super().__call__(size=size, **kwargs)
  
          def rng_fn(cls, rng, size):
              return 0
  
      from pytensor.link.jax.dispatch.random import jax_sample_fn
  
      @jax_sample_fn.register(CustomRV)
      def jax_sample_fn_custom(op):
          def sample_fn(rng, size, dtype, *parameters):
              return (rng, 0)
  
          return sample_fn
  
      nonexistentrv = CustomRV()
      rng = shared(np.random.RandomState(123))
      out = nonexistentrv(rng=rng)
      fgraph = FunctionGraph([out.owner.inputs[0]], [out], clone=False)
      with pytest.warns(
          UserWarning,
          match=r"The RandomType SharedVariables \[.+\] will not be used",
      ):
          compare_jax_and_py(fgraph, [])
  
imports: import numpy
  import pytest
  import scipy.stats
  import pytensor
  import pytensor.tensor
  import pytensor.tensor.random.basic
  from pytensor.compile.function import function
  from pytensor.compile.sharedvalue import SharedVariable, shared
  from pytensor.graph.basic import Constant
  from pytensor.graph.fg import FunctionGraph
  from pytensor.tensor.random.basic import RandomVariable
  from pytensor.tensor.random.type import RandomType
  from pytensor.tensor.random.utils import RandomStream
  from tests.link.jax.test_basic import compare_jax_and_py, jax_mode, set_test_value
  from pytensor.link.jax.dispatch.random import numpyro_available
label_desc: Use image node_id 20 for calling a global function with example usage: test_random_custom_implementation() and returns: sample_fn, super, int, rng,  int