def rand_some_inf_and_nan():
  """Return a random sampler that produces infinities in floating types."""
  rng = npr.RandomState(1)
  base_rand = rand_default()

  """
  TODO: Complex numbers are not correctly tested
  If blocks should be switched in order, and relevant tests should be fixed
  """
  def rand(shape, dtype):
    """The random sampler function."""
    if not onp.issubdtype(dtype, onp.floating):
      # only float types have inf
      return base_rand(shape, dtype)

    if onp.issubdtype(dtype, onp.complexfloating):
      base_dtype = onp.real(onp.array(0, dtype=dtype)).dtype
      out = (rand(shape, base_dtype) +
             onp.array(1j, dtype) * rand(shape, base_dtype))
      return _cast_to_shape(out, shape, dtype)

    dims = _dims_of_shape(shape)
    posinf_flips = rng.rand(*dims) < 0.1
    neginf_flips = rng.rand(*dims) < 0.1
    nan_flips = rng.rand(*dims) < 0.1

    vals = base_rand(shape, dtype)
    vals = onp.where(posinf_flips, onp.array(onp.inf, dtype=dtype), vals)
    vals = onp.where(neginf_flips, onp.array(-onp.inf, dtype=dtype), vals)
    vals = onp.where(nan_flips, onp.array(onp.nan, dtype=dtype), vals)

    return _cast_to_shape(onp.asarray(vals, dtype=dtype), shape, dtype)

  return rand
Return a random sampler that produces infinities in floating types.
trax/tf_numpy/jax_tests/test_util.py