Method call

tfjs-layers/src/layers/normalization.ts:571–624

(inputs: Tensor|Tensor[], kwargs: Kwargs)

  }

  override call(inputs: Tensor|Tensor[], kwargs: Kwargs): Tensor|Tensor[] {
    const input = getExactlyOneTensor(inputs);
    const inputShape = input.shape;
    const nDims = inputShape.length;

    return tidy(() => {
      const keepDims = true;
      let {mean, variance} = moments(input, this.axis, keepDims);
      const broadcastShape = generic_utils.pyListRepeat(1, nDims);
      for (const dim of this.axis as number[]) {
        broadcastShape[dim] = inputShape[dim];
      }

      const broadcast = (v: Tensor) => {
        if (v != null && v.shape.length !== nDims) {
          return tfc.reshape(v, broadcastShape);
        } else {
          return v;
        }
      };

      let scale = this.scale ? broadcast(this.gamma.read()) : null;
      let offset = this.center ? broadcast(this.beta.read()) : null;

      // TODO(https://github.com/tensorflow/tfjs/issues/2120): The tiling below
      // is a workaround for the limitation of core's batchNormalization?d don't
      // support broadcasting in their gradients. In addition, the tiling is
      // necessary to ensure correctness on the browser CPU backend regardless
      // of forward or backward computation. Remove this workaround once the
      // limitation is addressed. See .
      const momentsTiling: number[] = [];
      const scaleOffsetTiling: number[] = [];
      for (let i = 0; i < nDims; ++i) {
        if ((this.axis as number[]).indexOf(i) !== -1) {
          momentsTiling.push(inputShape[i]);
          scaleOffsetTiling.push(1);
        } else {
          momentsTiling.push(1);
          scaleOffsetTiling.push(inputShape[i]);
        }
      }
      mean = tfc.tile(mean, momentsTiling);
      variance = tfc.tile(variance, momentsTiling);
      if (scale != null) {
        scale = tfc.tile(scale, scaleOffsetTiling);
      }
      if (offset != null) {
        offset = tfc.tile(offset, scaleOffsetTiling);
      }

      return batchNormalization(
          input, mean, variance, offset, scale, this.epsilon);
    });
  }

  override getConfig(): serialization.ConfigDict {
    const config: serialization.ConfigDict = {