(self, train_stream, eval_stream=None, train_eval_stream=None)
Initialize a new set of inputs.

    Args:
      train_stream: a function taking n_devices (an int) and returning
        a python generator of training batches.
      eval_stream: a function taking n_devices (an int) and returning
        a python generator of validation batches;
        if None, then the training generator will be used for evaluation.
      train_eval_stream: a function taking n_devices (an int) and returning
        a python generator of batches from
        the training set used for evaluation (if None, use train_stream).
Initialize a new set of inputs.
Initialize a new set of inputs .
def __init__(self, train_stream, eval_stream=None, train_eval_stream=None):
    """Initialize a new set of inputs.

    Args:
      train_stream: a function taking n_devices (an int) and returning
        a python generator of training batches.
      eval_stream: a function taking n_devices (an int) and returning
        a python generator of validation batches;
        if None, then the training generator will be used for evaluation.
      train_eval_stream: a function taking n_devices (an int) and returning
        a python generator of batches from
        the training set used for evaluation (if None, use train_stream).
    """
    if not callable(train_stream):
      raise ValueError('Trax Inputs should be initialized with a function. '
                       'Did you forget the n_devices argument? If your inputs '
                       'do not use it, try lambda _: [your-inputs].')

    self._train_stream = train_stream
    self._eval_stream = eval_stream or self._train_stream

    # TODO(lukaszkaiser): should we get rid of this one day?
    self._train_eval_stream = train_eval_stream or self._train_stream

    # Peek into the train stream to get an example shape.
    example_train_batch = next(train_stream(1))
    self._input_shape = tuple(example_train_batch[0].shape)[1:]
    self._input_dtype = example_train_batch[0].dtype
    self._target_shape = tuple(example_train_batch[-1].shape)[1:]
    self._target_dtype = example_train_batch[-1].dtype
    self._example_shape = [x.shape for x in example_train_batch]
    self._example_dtype = [x.dtype for x in example_train_batch]
def __init__ ( self , train_stream , eval_stream = None , train_eval_stream = None ) : if not callable ( train_stream ) : raise ValueError ( 'Trax Inputs should be initialized with a function. ' 'Did you forget the n_devices argument? If your inputs ' 'do not use it, try lambda _: [your-inputs].' ) self . _train_stream = train_stream self . _eval_stream = eval_stream or self . _train_stream # TODO(lukaszkaiser): should we get rid of this one day? self . _train_eval_stream = train_eval_stream or self . _train_stream # Peek into the train stream to get an example shape. example_train_batch = next ( train_stream ( 1 ) ) self . _input_shape = tuple ( example_train_batch [ 0 ] . shape ) [ 1 : ] self . _input_dtype = example_train_batch [ 0 ] . dtype self . _target_shape = tuple ( example_train_batch [ - 1 ] . shape ) [ 1 : ] self . _target_dtype = example_train_batch [ - 1 ] . dtype self . _example_shape = [ x . shape for x in example_train_batch ] self . _example_dtype = [ x . dtype for x in example_train_batch ]