Function zip

tfjs-data/src/dataset.ts:627–659

(datasets: DatasetContainer)

 * @doc {heading: 'Data', subheading: 'Operations', namespace: 'data'}
 */
export function zip<O extends tf.TensorContainer>(datasets: DatasetContainer):
    Dataset<O> {
  // manually type-check the argument for JS users
  if (!isIterable(datasets)) {
    throw new Error('The argument to zip() must be an object or array.');
  }
  let size;
  if (Array.isArray(datasets)) {
    for (let i = 0; i < datasets.length; i++) {
      size = size == null ? (datasets[i] as Dataset<O>).size :
                            Math.min(size, (datasets[i] as Dataset<O>).size);
    }
  } else if (datasets instanceof Object) {
    for (const ds in datasets) {
      size = size == null ? (datasets[ds] as Dataset<O>).size :
                            Math.min(size, (datasets[ds] as Dataset<O>).size);
    }
  }
  return datasetFromIteratorFn<O>(async () => {
    const streams = await deepMapAndAwaitAll(datasets, d => {
      if (d instanceof Dataset) {
        return {value: d.iterator(), recurse: false};
      } else if (isIterable(d)) {
        return {value: null, recurse: true};
      } else {
        throw new Error(
            'Leaves of the structure passed to zip() must be Datasets, ' +
            'not primitives.');
      }
    });
    return iteratorFromZipped<O>(streams, ZipMismatchMode.SHORTEST);
  }, size);
}

/**
 * A zip function for use with deepZip, passed via the columnMajorBatch call.