Method next

tfjs-layers/src/engine/dataset_fakes.ts:115–185

()

  }

  async next(): Promise<IteratorResult<FitDatasetElement>> {
    const done = ++this.batchCount > this.numBatches;
    if (done) {
      return {done, value: null};
    }
    if (this.xTensorsFunc == null) {
      // Generate data randomly.
      return {
        done,
        value: done ? null : {
          xs: generateRandomTensorContainer(this.xBatchShape),
          ys: generateRandomTensorContainer(this.yBatchShape)
        }
      };
    } else {
      // Use preset tensors.
      if ((this.batchCount - 1) % this.numBatches === 0) {
        this.xTensorValues = this.xTensorsFunc();
        this.yTensorValues = this.yTensorsFunc();
        this.tensorIndex = 0;
      }
      const index = this.tensorIndex++;

      let xs: tfc.Tensor|{[name: string]: tfc.Tensor};
      if (Array.isArray(this.xTensorValues)) {
        xs = this.xTensorValues[index];
        tfc.util.assert(
            tfc.util.arraysEqual(xs.shape, this.xBatchShape as Shape),
            () => `Shape mismatch: expected: ${
                      JSON.stringify(this.xBatchShape)}; ` +
                `actual: ${JSON.stringify((xs as tfc.Tensor).shape)}`);
      } else {
        xs = {};
        for (const key in this.xTensorValues) {
          xs[key] = this.xTensorValues[key][index];
          tfc.util.assert(
              tfc.util.arraysEqual(xs[key].shape, this.xBatchShape as Shape),
              () => `Shape mismatch: expected: ${
                        JSON.stringify(this.xBatchShape)}; ` +
                  `actual: ${JSON.stringify((xs as tfc.Tensor).shape)}`);
        }
      }

      let ys: tfc.Tensor|{[name: string]: tfc.Tensor};
      if (Array.isArray(this.yTensorValues)) {
        // Get preset ys tensors for single-output models.
        ys = this.yTensorValues[index];
        tfc.util.assert(
            tfc.util.arraysEqual(ys.shape, this.yBatchShape as Shape),
            () => `Shape mismatch: expected: ${
                      JSON.stringify(this.yBatchShape)}; ` +
                `actual: ${JSON.stringify((ys as tfc.Tensor).shape)}`);
      } else {
        // Get preset ys tensors for multi-output models.
        ys = {};
        this.yBatchShape = this.yBatchShape as {[name: string]: Shape};
        for (const key in this.yTensorValues) {
          ys[key] = this.yTensorValues[key][index];