Function batch

tfjs-data/src/dataset.ts:136–159

* Groups elements into batches. * * It is assumed that each of the incoming dataset elements has the same * structure -- i.e. the same set of keys at each location in an object * hierarchy. For each key, the resulting `Dataset` provides a batched * element collecting all of the incom

(batchSize: number, smallLastBatch = true)

   * @doc {heading: 'Data', subheading: 'Classes'}
   */
  batch(batchSize: number, smallLastBatch = true): Dataset<tf.TensorContainer> {
    const base = this;
    tf.util.assert(
        batchSize > 0, () => `batchSize needs to be positive, but it is
      ${batchSize}`);
    let size;
    if (this.size === Infinity || this.size == null) {
      // If the size of this dataset is infinity or null, the new size keeps the
      // same.
      size = this.size;
    } else if (smallLastBatch) {
      // If the size of this dataset is known and include small last batch, the
      // new size is full batch count plus last batch.
      size = Math.ceil(this.size / batchSize);
    } else {
      // If the size of this dataset is known and not include small last batch,
      // the new size is full batch count.
      size = Math.floor(this.size / batchSize);
    }
    return datasetFromIteratorFn(async () => {
      return (await base.iterator())
          .columnMajorBatch(batchSize, smallLastBatch, deepBatchConcat);
    }, size);
  }

  /**
   * Concatenates this `Dataset` with another.