Make random Inputs for debugging.

  Args:
    num_devices: how many devices to build the inputs for.
    input_shape: the shape of inputs (including batch dimension).
    input_dtype: the type of the inputs (int32 by default).
    input_range: the range of inputs (defaults to (0, 255)).
    output_shape: the shape of outputs (including batch dimension).
    output_dtype: the type of the outputs (int32 by default).
    output_range: the range of outputs (defaults to (0, 9)).

  Returns:
    trax.inputs.Inputs

def random_inputs(
    num_devices,
    input_shape=gin.REQUIRED, input_dtype=np.int32, input_range=(0, 255),
    output_shape=gin.REQUIRED, output_dtype=np.int32, output_range=(0, 9)):
  """Make random Inputs for debugging.

  Args:
    num_devices: how many devices to build the inputs for.
    input_shape: the shape of inputs (including batch dimension).
    input_dtype: the type of the inputs (int32 by default).
    input_range: the range of inputs (defaults to (0, 255)).
    output_shape: the shape of outputs (including batch dimension).
    output_dtype: the type of the outputs (int32 by default).
    output_range: the range of outputs (defaults to (0, 9)).

  Returns:
    trax.inputs.Inputs
  """
  if input_shape[0] % num_devices != 0:
    tf.logging.fatal(
        "num_devices[%d] should divide the first dimension of input_shape[%s]",
        num_devices, input_shape)
  if output_shape[0] % num_devices != 0:
    tf.logging.fatal(
        "num_devices[%d] should divide the first dimension of output_shape[%s]",
        num_devices, output_shape)

  def random_minibatches():
    """Generate a stream of random mini-batches."""
    if input_dtype in [np.float16, np.float32, np.float64]:
      rand = np.random.uniform
    else:
      rand = np.random.random_integers
    while True:
      inp = rand(input_range[0], input_range[1], input_shape)
      inp = inp.astype(input_dtype)
      out = rand(output_range[0], output_range[1], output_shape)
      out = out.astype(output_dtype)
      yield inp, out

  input_shape_without_batch = list(input_shape)[1:]
  return Inputs(train_stream=random_minibatches,
                train_eval_stream=random_minibatches,
                eval_stream=random_minibatches,
                input_shape=input_shape_without_batch)