Function computeOutputShape

tfjs-layers/src/engine/container.ts:783–859

* Computes the output shape of the layer. * * Assumes that the layer will be built to match that input shape provided. * * @param inputShape A shape (tuple of integers) or a list of shape tuples * (one per output tensor of the layer). Shape tuples can include null for * free di

(inputShape: Shape|Shape[])

   *   free dimensions, instead of an integer.
   */
  override computeOutputShape(inputShape: Shape|Shape[]): Shape|Shape[] {
    const inputShapes = types_utils.normalizeShapeList(inputShape);
    if (inputShapes.length !== this.inputLayers.length) {
      throw new ValueError(
          `Invalid inputShape argument ${inputShape}: ` +
          `model has ${this.inputLayers.length} tensor inputs.`);
    }

    // TODO(michaelterry): Add caching
    const layersToOutputShapes: {[shapeKey: string]: Shape} = {};
    for (let i = 0; i < inputShapes.length; i++) {
      const layer = this.inputLayers[i];
      const inputShape = inputShapes[i];
      // It's an input layer: computeOutputShape is identity,
      // and there is only one node and one tensor output.
      const shapeKey = layer.name + '_0_0';
      layersToOutputShapes[shapeKey] = inputShape;
    }

    const depthKeys = Object.keys(this.nodesByDepth)
                          .map(x => parseInt(x, 10))
                          .sort(generic_utils.reverseNumberCompare);
    // Iterate over nodes, by depth level.
    if (depthKeys.length > 1) {
      for (const depth of depthKeys) {
        const nodes = this.nodesByDepth[depth];
        for (const node of nodes) {
          // This is always a single layer, never a list.
          const layer = node.outboundLayer;
          if (this.inputLayers.map(x => x.id).indexOf(layer.id) !== -1) {
            // We've already covered the input layers a few lines above.
            continue;
          }
          // Potentially redundant list, same size of node.inputTensors.
          const inputShapes: Shape[] = [];
          for (let j = 0; j < node.inboundLayers.length; j++) {
            const inboundLayer = node.inboundLayers[j];
            const nodeIndex = node.nodeIndices[j];
            const tensorIndex = node.tensorIndices[j];
            const shapeKey = `${inboundLayer.name}_${nodeIndex}_${tensorIndex}`;
            const inputShape = layersToOutputShapes[shapeKey];
            inputShapes.push(inputShape);
          }

          const outputShape = layer.computeOutputShape(
              generic_utils.singletonOrArray(inputShapes));

          const outputShapes = types_utils.normalizeShapeList(outputShape);
          const nodeIndex = layer.inboundNodes.indexOf(node);
          for (let j = 0; j < outputShapes.length; j++) {
            const shapeKey = `${layer.name}_${nodeIndex}_${j}`;
            layersToOutputShapes[shapeKey] = outputShapes[j];
          }
        }
      }
    }

    // Read final output shapes from layersToOutputShapes.