code: def test_replaced_shared_rng_storage_ordering_equality():
      """Test case described in issue #314.
  
      This happened when we tried to update the input storage after we clone the shared RNG.
      We used to call `input_storage.index(old_input_storage)` which would fail when the input_storage contained
      numpy arrays before the RNG value, which would fail the equality check.
  
      """
      pt_rng = RandomStream(1)
  
      batchshape = (3, 1, 4, 4)
      inp_shared = pytensor.shared(
          np.zeros(batchshape, dtype="float64"), name="inp_shared"
      )
  
      inp = pt.tensor4(dtype="float64", name="inp")
      inp_update = inp + pt_rng.normal(
          size=inp.shape, loc=5, scale=1e-5
      )
  
      # This function replaces inp by input_shared in the update expression
      # This is what caused the RNG to appear later than inp_shared in the input_storage
  
      fn = compile_random_function(
          inputs=[],
          outputs=[],
          updates={inp_shared: inp_update},
          givens={inp: inp_shared},
          mode="JAX",
      )
      fn()
      np.testing.assert_allclose(inp_shared.get_value(), 5, rtol=1e-3)
      fn()
      np.testing.assert_allclose(inp_shared.get_value(), 10, rtol=1e-3)
  
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 5 for calling a global function with example usage: test_replaced_shared_rng_storage_ordering_equality() without return types