(
    input_shape=gin.REQUIRED, input_dtype=jnp.int32, input_range=(0, 255),
    output_shape=gin.REQUIRED, output_dtype=jnp.int32, output_range=(0, 9))
return Inputs(random_minibatches)
Make random Inputs for debugging.

  Args:
    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
Make random Inputs for debugging.
Make random Inputs for debugging .
def random_inputs(
    input_shape=gin.REQUIRED, input_dtype=jnp.int32, input_range=(0, 255),
    output_shape=gin.REQUIRED, output_dtype=jnp.int32, output_range=(0, 9)):
  """Make random Inputs for debugging.

  Args:
    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_minibatches(n_devices):
    """Generate a stream of random mini-batches."""
    assert input_range[0] % n_devices == 0
    if input_dtype in [jnp.float16, jnp.float32, jnp.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

  return Inputs(random_minibatches)
def random_inputs ( input_shape = gin . REQUIRED , input_dtype = jnp . int32 , input_range = ( 0 , 255 ) , output_shape = gin . REQUIRED , output_dtype = jnp . int32 , output_range = ( 0 , 9 ) ) : def random_minibatches ( n_devices ) : """Generate a stream of random mini-batches.""" assert input_range [ 0 ] % n_devices == 0 if input_dtype in [ jnp . float16 , jnp . float32 , jnp . 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 return Inputs ( random_minibatches )