text: import functools
  
  import numpy as np
  
  import aesara
  from aesara import tensor as aet
  from aesara.configdefaults import config
  from aesara.gpuarray.rng_mrg import GPUA_mrg_uniform
  from aesara.gpuarray.type import gpuarray_shared_constructor
  from aesara.sandbox import rng_mrg
  from aesara.sandbox.rng_mrg import MRG_RandomStream
  from tests import unittest_tools as utt
  from tests.gpuarray.config import mode_with_gpu as mode
  from tests.sandbox.test_rng_mrg import java_samples, rng_mrg_overflow
  from tests.sandbox.test_rng_mrg import test_f16_nonzero as cpu_f16_nonzero
  
  
  def test_consistency_GPUA_serial():
      # Verify that the random numbers generated by GPUA_mrg_uniform, serially,
      # are the same as the reference (Java) implementation by L'Ecuyer et al.
  
      seed = 12345
      n_samples = 5
      n_streams = 12
      n_substreams = 7
  
      samples = []
      curr_rstate = np.array([seed] * 6, dtype="int32")
  
      for i in range(n_streams):
          stream_rstate = curr_rstate.copy()
          for j in range(n_substreams):
              substream_rstate = np.array([stream_rstate.copy()], dtype="int32")
              # Transfer to device
              rstate = gpuarray_shared_constructor(substream_rstate)
  
              new_rstate, sample = GPUA_mrg_uniform.new(
                  rstate, ndim=None, dtype="float32", size=(1,)
              )
              rstate.default_update = new_rstate
  
              # Not really necessary, just mimicking
              # rng_mrg.MRG_RandomStream' behavior
              sample.rstate = rstate
              sample.update = (rstate, new_rstate)
  
              # We need the sample back in the main memory
              cpu_sample = aet.as_tensor_variable(sample)
              f = aesara.function([], cpu_sample, mode=mode)
              for k in range(n_samples):
                  s = f()
                  samples.append(s)
  
              # next substream
              stream_rstate = rng_mrg.ff_2p72(stream_rstate)
  
          # next stream
          curr_rstate = rng_mrg.ff_2p134(curr_rstate)
  
      samples = np.array(samples).flatten()
      assert np.allclose(samples, java_samples)
  
  
  def test_consistency_GPUA_parallel():
      # Verify that the random numbers generated by GPUA_mrg_uniform, in
      # parallel, are the same as the reference (Java) implementation by
      # L'Ecuyer et al.
      seed = 12345
      n_samples = 5
      n_streams = 12
      n_substreams = 7  # 7 samples will be drawn in parallel
  
      samples = []
      curr_rstate = np.array([seed] * 6, dtype="int32")
  
      for i in range(n_streams):
          stream_samples = []
          rstate = [curr_rstate.copy()]
          for j in range(1, n_substreams):
              rstate.append(rng_mrg.ff_2p72(rstate[-1]))
          rstate = np.asarray(rstate)
          rstate = gpuarray_shared_constructor(rstate)
  
          new_rstate, sample = GPUA_mrg_uniform.new(
              rstate, ndim=None, dtype="float32", size=(n_substreams,)
          )
          rstate.default_update = new_rstate
  
          # Not really necessary, just mimicking
          # rng_mrg.MRG_RandomStream' behavior
          sample.rstate = rstate
          sample.update = (rstate, new_rstate)
  
          # We need the sample back in the main memory
          cpu_sample = aet.as_tensor_variable(sample)
          f = aesara.function([], cpu_sample, mode=mode)
  
          for k in range(n_samples):
              s = f()
              stream_samples.append(s)
  
          samples.append(np.array(stream_samples).T.flatten())
  
          # next stream
          curr_rstate = rng_mrg.ff_2p134(curr_rstate)
  
      samples = np.array(samples).flatten()
      assert np.allclose(samples, java_samples)
  
  
  def test_GPUA_full_fill():
      # Make sure the whole sample buffer is filled.  Also make sure
      # large samples are consistent with CPU results.
  
      # This needs to be large to trigger the problem on GPU
      size = (10, 1000)
  
      R = MRG_RandomStream(234)
      uni = R.uniform(size, nstreams=60 * 256)
      f_cpu = aesara.function([], uni)
  
      rstate_gpu = gpuarray_shared_constructor(R.state_updates[-1][0].get_value())
      new_rstate, sample = GPUA_mrg_uniform.new(
          rstate_gpu, ndim=None, dtype="float32", size=size
      )
      rstate_gpu.default_update = new_rstate
      f_gpu = aesara.function([], sample, mode=mode)
  
      utt.assert_allclose(f_cpu(), f_gpu())
  
  
  def test_overflow_gpu_new_backend():
      seed = 12345
      n_substreams = 7
      curr_rstate = np.array([seed] * 6, dtype="int32")
      rstate = [curr_rstate.copy()]
      for j in range(1, n_substreams):
          rstate.append(rng_mrg.ff_2p72(rstate[-1]))
      rstate = np.asarray(rstate)
      rstate = gpuarray_shared_constructor(rstate)
      fct = functools.partial(GPUA_mrg_uniform.new, rstate, ndim=None, dtype="float32")
      # should raise error as the size overflows
      sizes = [
          (2 ** 31,),
          (2 ** 32,),
          (
              2 ** 15,
              2 ** 16,
          ),
          (2, 2 ** 15, 2 ** 15),
      ]
      rng_mrg_overflow(sizes, fct, mode, should_raise_error=True)
      # should not raise error
      sizes = [(2 ** 5,), (2 ** 5, 2 ** 5), (2 ** 5, 2 ** 5, 2 ** 5)]
      rng_mrg_overflow(sizes, fct, mode, should_raise_error=False)
      # should support int32 sizes
      sizes = [(np.int32(2 ** 10),), (np.int32(2), np.int32(2 ** 10), np.int32(2 ** 10))]
      rng_mrg_overflow(sizes, fct, mode, should_raise_error=False)
  
  
  def test_validate_input_types_gpuarray_backend():
      with config.change_flags(compute_test_value="raise"):
          rstate = np.zeros((7, 6), dtype="int32")
          rstate = gpuarray_shared_constructor(rstate)
          rng_mrg.mrg_uniform.new(rstate, ndim=None, dtype="float32", size=(3,))
  
  
  def test_f16_nonzero():
      try:
          # To have aesara.shared(x) try to move on the GPU
          aesara.compile.shared_constructor(gpuarray_shared_constructor)
          cpu_f16_nonzero(mode=mode, op_to_check=GPUA_mrg_uniform)
      finally:
          aesara.compile.shared_constructor(gpuarray_shared_constructor, remove=True)
  
  
  def test_cpu_target_with_shared_variable():
      srng = MRG_RandomStream()
      s = np.random.rand(2, 3).astype("float32")
      x = gpuarray_shared_constructor(s, name="x")
      try:
          # To have aesara.shared(x) try to move on the GPU
          aesara.compile.shared_constructor(gpuarray_shared_constructor)
          y = srng.uniform(x.shape, target="cpu")
          y.name = "y"
          z = (x * y).sum()
          z.name = "z"
  
          fz = aesara.function([], z, mode=mode)
  
          nodes = fz.maker.fgraph.toposort()
          assert not any(isinstance(node.op, GPUA_mrg_uniform) for node in nodes)
      finally:
          aesara.compile.shared_constructor(gpuarray_shared_constructor, remove=True)
  