Function take

tfjs-data/src/dataset.ts:473–490

* Creates a `Dataset` with at most `count` initial elements from this * dataset. * * ```js * const a = tf.data.array([1, 2, 3, 4, 5, 6]).take(3); * await a.forEachAsync(e => console.log(e)); * ``` * * @param count: The number of elements of this dataset that should be taken

(count: number)

   * @doc {heading: 'Data', subheading: 'Classes'}
   */
  take(count: number): Dataset<T> {
    const base = this;
    let size;
    if (this.size != null && this.size > count) {
      // If the size of this dataset is greater than count, the new dataset's
      // size is count.
      size = count;
    } else if (this.size != null && this.size <= count) {
      // If the size of this dataset is equal or smaller than count, the new
      // dataset's size is the size of this dataset.
      size = this.size;
    } else {
      // If the size of this dataset is null, the new dataset's size is null.
      size = null;
    }
    return datasetFromIteratorFn(
        async () => (await base.iterator()).take(count), size);
  }

  /**
   * Collect all elements of this dataset into an array.