code: def check_with_device(device, dtype):
      # The thresholds chosen for the tests are too loose. We will rely on the other tests to test the samples from the
      #  generators.
      tol = 0.1
      symbols = [
          {
              "name": "normal",
              "symbol": mx.sym.random.normal,
              "ndop": mx.nd.random.normal,
              "pdfsymbol": mx.sym.random_pdf_normal,
              "pdffunc": ss.norm.pdf,
              "discrete": False,
              "params": {"loc": 10.0, "scale": 0.5},
              "inputs": [
                  ("loc", [[0.0, 2.5], [-9.75, -7.0]]),
                  ("scale", [[1.0, 3.7], [4.2, 1.5]]),
              ],
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(
                          x.astype(np.float64) - params["loc"]
                      ),
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - params["scale"],
                      tol,
                  ),
              ],
          },
          {
              "name": "normal_like",
              "symbol": mx.sym.random.normal_like,
              "ndop": mx.nd.random.normal_like,
              "params": {"loc": 10.0, "scale": 0.5},
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(
                          x.astype(np.float64) - params["loc"]
                      ),
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - params["scale"],
                      tol,
                  ),
              ],
          },
          {
              "name": "randn",
              "symbol": mx.sym.random.randn,
              "ndop": mx.nd.random.randn,
              "params": {"loc": 10.0, "scale": 0.5},
              "inputs": [
                  ("loc", [[0.0, 2.5], [-9.75, -7.0]]),
                  ("scale", [[1.0, 3.7], [4.2, 1.5]]),
              ],
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(
                          x.astype(np.float64) - params["loc"]
                      ),
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - params["scale"],
                      tol,
                  ),
              ],
          },
          {
              "name": "uniform",
              "symbol": mx.sym.random.uniform,
              "ndop": mx.nd.random.uniform,
              "pdfsymbol": mx.sym.random_pdf_uniform,
              "pdffunc": lambda x, low, high: ss.uniform.pdf(
                  x, low, high - low
              ),
              "discrete": False,
              "params": {"low": -1.5, "high": 3.0},
              "inputs": [
                  ("low", [[0.0, 2.5], [-9.75, -1.0]]),
                  ("high", [[1.0, 3.7], [4.2, 10.5]]),
              ],
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(x.astype(np.float64))
                      - (params["low"] + params["high"]) / 2.0,
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - np.sqrt(1.0 / 12.0)
                      * (params["high"] - params["low"]),
                      tol,
                  ),
              ],
          },
          {
              "name": "uniform_like",
              "symbol": mx.sym.random.uniform_like,
              "ndop": mx.nd.random.uniform_like,
              "params": {"low": -1.5, "high": 3.0},
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(x.astype(np.float64))
                      - (params["low"] + params["high"]) / 2.0,
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - np.sqrt(1.0 / 12.0)
                      * (params["high"] - params["low"]),
                      tol,
                  ),
              ],
          },
          {
              "name": "gamma",
              "symbol": mx.sym.random.gamma,
              "ndop": mx.nd.random.gamma,
              "pdfsymbol": mx.sym.random_pdf_gamma,
              "pdffunc": lambda x, alpha, beta: ss.gamma.pdf(
                  x, alpha, 0, 1 / beta
              ),
              "discrete": False,
              "params": {"alpha": 9.0, "beta": 0.5},
              "inputs": [
                  ("alpha", [[0.1, 2.5], [9.75, 11.0]]),
                  ("beta", [[1.0, 0.7], [0.5, 0.3]]),
              ],
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(x.astype(np.float64))
                      - params["alpha"] * params["beta"],
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - np.sqrt(params["alpha"] * params["beta"] ** 2),
                      tol,
                  ),
              ],
          },
          {
              "name": "gamma_like",
              "symbol": mx.sym.random.gamma_like,
              "ndop": mx.nd.random.gamma_like,
              "params": {"alpha": 9.0, "beta": 0.5},
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(x.astype(np.float64))
                      - params["alpha"] * params["beta"],
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - np.sqrt(params["alpha"] * params["beta"] ** 2),
                      tol,
                  ),
              ],
          },
          {
              "name": "exponential",
              "symbol": mx.sym.random.exponential,
              "ndop": mx.nd.random.exponential,
              "pdfsymbol": mx.sym.random_pdf_exponential,
              "pdffunc": lambda x, lam: ss.expon.pdf(x, 0, 1 / lam),
              "discrete": False,
              "params": {"scale": 1.0 / 4.0},
              "inputs": [
                  (
                      "scale",
                      [[1.0 / 1.0, 1.0 / 8.5], [1.0 / 2.7, 1.0 / 0.5]],
                  )
              ],
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(x.astype(np.float64))
                      - params["scale"],
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - params["scale"],
                      tol,
                  ),
              ],
          },
          {
              "name": "exponential_like",
              "symbol": mx.sym.random.exponential_like,
              "ndop": mx.nd.random.exponential_like,
              "params": {"lam": 4.0},
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(x.astype(np.float64))
                      - 1.0 / params["lam"],
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - 1.0 / params["lam"],
                      tol,
                  ),
              ],
          },
          {
              "name": "poisson",
              "symbol": mx.sym.random.poisson,
              "ndop": mx.nd.random.poisson,
              "pdfsymbol": mx.sym.random_pdf_poisson,
              "pdffunc": ss.poisson.pmf,
              "discrete": True,
              "params": {"lam": 4.0},
              "inputs": [("lam", [[25.0, 8.5], [2.7, 0.5]])],
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(x.astype(np.float64))
                      - params["lam"],
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - np.sqrt(params["lam"]),
                      tol,
                  ),
              ],
          },
          {
              "name": "poisson_like",
              "symbol": mx.sym.random.poisson_like,
              "ndop": mx.nd.random.poisson_like,
              "params": {"lam": 4.0},
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(x.astype(np.float64))
                      - params["lam"],
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - np.sqrt(params["lam"]),
                      tol,
                  ),
              ],
          },
          {
              "name": "neg_binomial",
              "symbol": mx.sym.random.negative_binomial,
              "ndop": mx.nd.random.negative_binomial,
              "pdfsymbol": mx.sym.random_pdf_negative_binomial,
              "pdffunc": ss.nbinom.pmf,
              "discrete": True,
              "params": {"k": 3, "p": 0.4},
              "inputs": [
                  ("k", [[3, 4], [5, 6]]),
                  ("p", [[0.4, 0.77], [0.5, 0.84]]),
              ],
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(x.astype(np.float64))
                      - params["k"] * (1.0 - params["p"]) / params["p"],
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - np.sqrt(params["k"] * (1.0 - params["p"]))
                      / params["p"],
                      tol,
                  ),
              ],
          },
          {
              "name": "neg_binomial_like",
              "symbol": mx.sym.random.negative_binomial_like,
              "ndop": mx.nd.random.negative_binomial_like,
              "params": {"k": 3, "p": 0.4},
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(x.astype(np.float64))
                      - params["k"] * (1.0 - params["p"]) / params["p"],
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - np.sqrt(params["k"] * (1.0 - params["p"]))
                      / params["p"],
                      tol,
                  ),
              ],
          },
          {
              "name": "gen_neg_binomial",
              "symbol": mx.sym.random.generalized_negative_binomial,
              "ndop": mx.nd.random.generalized_negative_binomial,
              "pdfsymbol": mx.sym.random_pdf_generalized_negative_binomial,
              "pdffunc": lambda x, mu, alpha: ss.nbinom.pmf(
                  x, 1.0 / alpha, 1.0 / (mu * alpha + 1.0)
              ),
              "discrete": True,
              "params": {"mu": 2.0, "alpha": 0.3},
              "inputs": [
                  ("mu", [[2.0, 2.5], [1.3, 1.9]]),
                  ("alpha", [[1.0, 0.1], [0.2, 0.5]]),
              ],
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(x.astype(np.float64))
                      - params["mu"],
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - np.sqrt(
                          params["mu"]
                          + params["alpha"] * params["mu"] ** 2
                      ),
                      tol,
                  ),
              ],
          },
          {
              "name": "gen_neg_binomial_like",
              "symbol": mx.sym.random.generalized_negative_binomial_like,
              "ndop": mx.nd.random.generalized_negative_binomial_like,
              "params": {"mu": 2.0, "alpha": 0.3},
              "checks": [
                  (
                      "mean",
                      lambda x, params: np.mean(x.astype(np.float64))
                      - params["mu"],
                      tol,
                  ),
                  (
                      "std",
                      lambda x, params: np.std(x.astype(np.float64))
                      - np.sqrt(
                          params["mu"]
                          + params["alpha"] * params["mu"] ** 2
                      ),
                      tol,
                  ),
              ],
          },
      ]
  
      # Create enough samples such that we get a meaningful distribution.
      shape = (500, 500)
      # Test pdf on smaller shapes as backward checks will take too long otherwise.
      # This must be a subshape of the former one.
      pdfshape = (30, 30)
      for symbdic in symbols:
          name = symbdic["name"]
          ndop = symbdic["ndop"]
  
          # check directly
          params = symbdic["params"].copy()
          params.update(shape=shape, dtype=dtype, ctx=device)
          args = ()
          if name == "randn":
              params.pop("shape")  # randn does not accept shape param
              args = shape
          if name.endswith("_like"):
              params["data"] = mx.nd.ones(
                  params.pop("shape"),
                  dtype=params.pop("dtype"),
                  ctx=params.pop("ctx"),
              )
          mx.random.seed(128)
          ret1 = ndop(*args, **params).asnumpy()
          mx.random.seed(128)
          ret2 = ndop(*args, **params).asnumpy()
          assert same(
              ret1, ret2
          ), f"ndarray test: `{name}` should give the same result with the same seed"
  
          for check_name, check_func, tol in symbdic["checks"]:
              assert (
                  np.abs(check_func(ret1, params)) < tol
              ), f"ndarray test: {check_name} check for `{name}` did not pass"
  
          # check multi-distribution sampling
          if "inputs" not in symbdic:
              continue  # randn does not support multi-distribution sampling
  
          params = {"shape": shape, "dtype": dtype, "ctx": device}
          params.update(
              {
                  k: mx.nd.array(v, ctx=device, dtype=dtype)
                  for k, v in symbdic["inputs"]
              }
          )
          if name == "randn":
              params.pop("shape")  # randn does not accept shape param
              args = shape
          mx.random.seed(128)
          ret1 = ndop(*args, **params).asnumpy()
          mx.random.seed(128)
          ret2 = ndop(*args, **params).asnumpy()
          assert same(
              ret1, ret2
          ), f"ndarray test: `{name}` should give the same result with the same seed"
          for i in range(2):
              for j in range(2):
                  stats = {k: v[i][j] for k, v in symbdic["inputs"]}
                  for check_name, check_func, tol in symbdic["checks"]:
                      err = np.abs(check_func(ret2[i, j], stats))
                      assert (
                          err < tol
                      ), f"{err} vs {tol}: symbolic test: {check_name} check for `{name}` did not pass"
  
          # check symbolic
          symbol = symbdic["symbol"]
          X = mx.sym.Variable("X")
          params = symbdic["params"].copy()
          params.update(shape=shape, dtype=dtype)
          if name.endswith("_like") or name == "randn":
              params["data"] = mx.sym.ones(params.pop("shape"))
          Y = symbol(**params) + X
          x = mx.nd.zeros(shape, dtype=dtype, ctx=device)
          xgrad = mx.nd.zeros(shape, dtype=dtype, ctx=device)
          yexec = Y._bind(device, {"X": x}, {"X": xgrad})
          mx.random.seed(128)
          yexec.forward(is_train=True)
          yexec.backward(yexec.outputs[0])
          un1 = (yexec.outputs[0] - x).copyto(device)
          assert same(xgrad.asnumpy(), un1.asnumpy())
          mx.random.seed(128)
          yexec.forward()
          un2 = (yexec.outputs[0] - x).copyto(device)
          assert same(
              un1.asnumpy(), un2.asnumpy()
          ), f"symbolic test: `{name}` should give the same result with the same seed"
  
          ret1 = un1.asnumpy()
          for check_name, check_func, tol in symbdic["checks"]:
              assert (
                  np.abs(check_func(ret1, params)) < tol
              ), f"symbolic test: {check_name} check for `{name}` did not pass"
          if name.endswith("_like"):
              continue
  
          # check multi-distribution sampling
          symbol = symbdic["symbol"]
          params = {"shape": shape, "dtype": dtype}
          single_param = len(symbdic["inputs"]) == 1
          v1 = mx.sym.Variable("v1")
          v2 = mx.sym.Variable("v2")
          if name == "randn":
              params.pop("shape")  # randn does not accept shape param
              args = shape
              Y = (
                  symbol(v1, **params)
                  if single_param
                  else symbol(*args, loc=v1, scale=v2, **params)
              )
          else:
              Y = (
                  symbol(v1, **params)
                  if single_param
                  else symbol(v1, v2, **params)
              )
          bindings = {"v1": mx.nd.array(symbdic["inputs"][0][1])}
          if not single_param:
              bindings.update(
                  {"v2": mx.nd.array(symbdic["inputs"][1][1])}
              )
          yexec = Y._bind(ctx=device, args=bindings)
          yexec.forward()
          un1 = yexec.outputs[0].copyto(device).asnumpy()
          params = {}
          for i, r in enumerate(symbdic["inputs"][0][1]):
              for j, p1 in enumerate(r):
                  params.update({symbdic["inputs"][0][0]: p1})
                  if not single_param:
                      params.update(
                          {
                              symbdic["inputs"][1][0]: symbdic[
                                  "inputs"
                              ][1][1][i][j]
                          }
                      )
                  samples = un1[i, j]
                  for check_name, check_func, tol in symbdic["checks"]:
                      assert (
                          np.abs(check_func(samples, params)) < tol
                      ), f"symbolic test: {check_name} check for `{name}` did not pass"
  
          if "pdfsymbol" not in symbdic:
              continue  # randn not tested for pdf
  
          # check pdfs with only a subset of the generated samples
          un1 = np.resize(
              un1,
              (un1.shape[0], un1.shape[1], pdfshape[0], pdfshape[1]),
          )
          symbol = symbdic["pdfsymbol"]
          pdffunc = symbdic["pdffunc"]
          v0 = mx.sym.Variable("v0")
          v1 = mx.sym.Variable("v1")
          v2 = mx.sym.Variable("v2")
          p1 = np.array(symbdic["inputs"][0][1])
          p2 = (
              None
              if single_param
              else np.array(symbdic["inputs"][1][1])
          )
          # Move samples away from boundaries of support
          if name == "gamma" or name == "exponential":
              un1 = np.maximum(un1, 1e-1)
          if name == "uniform":
              un1 = np.minimum(
                  np.maximum(
                      un1.reshape((un1.shape[0], un1.shape[1], -1)),
                      p1.reshape((p1.shape[0], p1.shape[1], -1)) + 1e-4,
                  ),
                  p2.reshape((p2.shape[0], p2.shape[1], -1)) - 1e-4,
              ).reshape(un1.shape)
          for use_log in [False, True]:
              test_pdf = (
                  symbol(v0, v1, is_log=use_log)
                  if single_param
                  else symbol(v0, v1, v2, is_log=use_log)
              )
              forw_atol = 1e-7 if dtype != np.float16 else 1e-3
              forw_rtol = 1e-4 if dtype != np.float16 else 5e-2
              backw_atol = 1e-3
              backw_rtol = 5e-2
              if single_param:
                  res = pdffunc(
                      un1.reshape((un1.shape[0], un1.shape[1], -1)),
                      p1.reshape((p1.shape[0], p1.shape[1], -1)),
                  ).reshape(un1.shape)
                  if use_log:
                      res = np.log(res)
                  check_symbolic_forward(
                      test_pdf,
                      [un1, p1],
                      [res],
                      atol=forw_atol,
                      rtol=forw_rtol,
                      dtype=dtype,
                  )
                  if dtype == np.float64:
                      grad_nodes = (
                          ["v1"]
                          if symbdic["discrete"]
                          else ["v0", "v1"]
                      )
                      check_numeric_gradient(
                          test_pdf,
                          [un1, p1],
                          grad_nodes=grad_nodes,
                          atol=backw_atol,
                          rtol=backw_rtol,
                          dtype=dtype,
                      )
              else:
                  res = pdffunc(
                      un1.reshape((un1.shape[0], un1.shape[1], -1)),
                      p1.reshape((p1.shape[0], p1.shape[1], -1)),
                      p2.reshape((p2.shape[0], p2.shape[1], -1)),
                  ).reshape(un1.shape)
                  if use_log:
                      res = np.log(res)
                  check_symbolic_forward(
                      test_pdf,
                      [un1, p1, p2],
                      [res],
                      atol=forw_atol,
                      rtol=forw_rtol,
                      dtype=dtype,
                  )
                  if dtype == np.float64:
                      grad_nodes = (
                          ["v1", "v2"]
                          if symbdic["discrete"]
                          else ["v0", "v1", "v2"]
                      )
                      check_numeric_gradient(
                          test_pdf,
                          [un1, p1, p2],
                          grad_nodes=grad_nodes,
                          atol=backw_atol,
                          rtol=backw_rtol,
                          dtype=dtype,
                      )
  
imports: import os
  import math
  import itertools
  import mxnet
  from mxnet.test_utils import verify_generator, gen_buckets_probs_with_ppf, assert_almost_equal
  import numpy
  import random
  from common import retry, random_seed
  import scipy.stats
  import unittest
  import pytest
  from mxnet.test_utils import *
  from mxnet.base import MXNetError
  from common import assertRaises
label_desc: Use image node_id 2 for calling a global function with example usage: check_with_device(device, dtype) without return types