text: /**
   * @license
   * Copyright 2018 Google LLC
   *
   * Use of this source code is governed by an MIT-style
   * license that can be found in the LICENSE file or at
   * https://opensource.org/licenses/MIT.
   * =============================================================================
   */
  /**
   * Layers that augment the functionality of a base layer.
   */
  import * as tfc from '@tensorflow/tfjs-core';
  import { serialization, tidy } from '@tensorflow/tfjs-core';
  import * as K from '../backend/tfjs_backend';
  import { nameScope } from '../common';
  import { InputSpec, Layer, SymbolicTensor } from '../engine/topology';
  import { NotImplementedError, ValueError } from '../errors';
  import { VALID_BIDIRECTIONAL_MERGE_MODES } from '../keras_format/common';
  import * as generic_utils from '../utils/generic_utils';
  import { getExactlyOneShape, getExactlyOneTensor } from '../utils/types_utils';
  import { rnn, standardizeArgs } from './recurrent';
  import { deserialize } from './serialization';
  /**
   * Abstract wrapper base class.
   *
   * Wrappers take another layer and augment it in various ways.
   * Do not use this class as a layer, it is only an abstract base class.
   * Two usable wrappers are the `TimeDistributed` and `Bidirectional` wrappers.
   */
  export class Wrapper extends Layer {
      constructor(args) {
          // Porting Note: In PyKeras, `self.layer` is set prior to the calling
          //   `super()`. But we can't do that here due to TypeScript's restriction.
          //   See: https://github.com/Microsoft/TypeScript/issues/8277
          //   As a result, we have to add checks in `get trainable()` and
          //   `set trainable()` below in order to prevent using `this.layer` when
          //   its value is `undefined`. The super constructor does use the getter
          //   and the setter of `this.layer`.
          super(args);
          this.layer = args.layer;
      }
      build(inputShape) {
          this.built = true;
      }
      // TODO(cais): Implement activityRegularizer getter.
      get trainable() {
          // Porting Note: the check of `this.layer` here is necessary due to the
          //   way the `constructor` of this class is written (see Porting Note
          //   above).
          if (this.layer != null) {
              return this.layer.trainable;
          }
          else {
              return false;
          }
      }
      set trainable(value) {
          // Porting Note: the check of `this.layer` here is necessary due to the
          //   way the `constructor` of this class is written (see Porting Note
          //   above).
          if (this.layer != null) {
              this.layer.trainable = value;
          }
      }
      get trainableWeights() {
          return this.layer.trainableWeights;
      }
      // TODO(cais): Implement setter for trainableWeights.
      get nonTrainableWeights() {
          return this.layer.nonTrainableWeights;
      }
      // TODO(cais): Implement setter for nonTrainableWeights.
      get updates() {
          // tslint:disable-next-line:no-any
          return this.layer._updates;
      }
      // TODO(cais): Implement getUpdatesFor().
      get losses() {
          return this.layer.losses;
      }
      // TODO(cais): Implement getLossesFor().
      getWeights() {
          return this.layer.getWeights();
      }
      setWeights(weights) {
          this.layer.setWeights(weights);
      }
      getConfig() {
          const config = {
              'layer': {
                  'className': this.layer.getClassName(),
                  'config': this.layer.getConfig(),
              }
          };
          const baseConfig = super.getConfig();
          Object.assign(config, baseConfig);
          return config;
      }
      setFastWeightInitDuringBuild(value) {
          super.setFastWeightInitDuringBuild(value);
          if (this.layer != null) {
              this.layer.setFastWeightInitDuringBuild(value);
          }
      }
      /** @nocollapse */
      static fromConfig(cls, config, customObjects = {}) {
          const layerConfig = config['layer'];
          const layer = deserialize(layerConfig, customObjects);
          delete config['layer'];
          const newConfig = { layer };
          Object.assign(newConfig, config);
          return new cls(newConfig);
      }
  }
  export class TimeDistributed extends Wrapper {
      constructor(args) {
          super(args);
          this.supportsMasking = true;
      }
      build(inputShape) {
          inputShape = getExactlyOneShape(inputShape);
          if (inputShape.length < 3) {
              throw new ValueError(`TimeDistributed layer expects an input shape >= 3D, but received ` +
                  `input shape ${JSON.stringify(inputShape)}`);
          }
          this.inputSpec = [{ shape: inputShape }];
          const childInputShape = [inputShape[0]].concat(inputShape.slice(2));
          if (!this.layer.built) {
              this.layer.build(childInputShape);
              this.layer.built = true;
          }
          super.build(inputShape);
      }
      computeOutputShape(inputShape) {
          inputShape = getExactlyOneShape(inputShape);
          const childInputShape = [inputShape[0]].concat(inputShape.slice(2));
          const childOutputShape = this.layer.computeOutputShape(childInputShape);
          const timesteps = inputShape[1];
          return [childOutputShape[0], timesteps].concat(childOutputShape.slice(1));
      }
      call(inputs, kwargs) {
          return tidy(() => {
              // TODO(cais): Add 'training' and 'useLearningPhase' to kwargs.
              inputs = getExactlyOneTensor(inputs);
              // Porting Note: In tfjs-layers, `inputs` are always concrete tensor
              // values. Hence the inputs can't have an undetermined first (batch)
              // dimension, which is why we always use the K.rnn approach here.
              const step = (inputs, states) => {
                  // TODO(cais): Add useLearningPhase.
                  // NOTE(cais): `layer.call` may return a length-1 array of Tensor in
                  //   some cases (e.g., `layer` is a `Sequential` instance), which is
                  //   why `getExactlyOneTensor` is used below.
                  const output = getExactlyOneTensor(this.layer.call(inputs, kwargs));
                  return [output, []];
              };
              const rnnOutputs = rnn(step, inputs, [], false /* goBackwards */, null /* mask */, null /* constants */, false /* unroll */, true /* needPerStepOutputs */);
              const y = rnnOutputs[1];
              // TODO(cais): Add activity regularization.
              // TODO(cais): Add useLearningPhase.
              return y;
          });
      }
  }
  /** @nocollapse */
  TimeDistributed.className = 'TimeDistributed';
  serialization.registerClass(TimeDistributed);
  export function checkBidirectionalMergeMode(value) {
      generic_utils.checkStringTypeUnionValue(VALID_BIDIRECTIONAL_MERGE_MODES, 'BidirectionalMergeMode', value);
  }
  const DEFAULT_BIDIRECTIONAL_MERGE_MODE = 'concat';
  export class Bidirectional extends Wrapper {
      constructor(args) {
          super(args);
          // Note: When creating `this.forwardLayer`, the original Layer object
          //   (`config.layer`) ought to be cloned. This is why we call
          //   `getConfig()` followed by `deserialize()`. Without this cloning,
          //   the layer names saved during serialization will incorrectly contain
          //   the 'forward_' prefix. In Python Keras, this is done using
          //   `copy.copy` (shallow copy), which does not have a simple equivalent
          //   in JavaScript. JavaScript's `Object.assign()` does not copy
          //   methods.
          const layerConfig = args.layer.getConfig();
          const forwDict = {};
          forwDict['className'] = args.layer.getClassName();
          forwDict['config'] = layerConfig;
          this.forwardLayer = deserialize(forwDict);
          layerConfig['goBackwards'] =
              layerConfig['goBackwards'] === true ? false : true;
          const backDict = {};
          backDict['className'] = args.layer.getClassName();
          backDict['config'] = layerConfig;
          this.backwardLayer = deserialize(backDict);
          this.forwardLayer.name = 'forward_' + this.forwardLayer.name;
          this.backwardLayer.name = 'backward_' + this.backwardLayer.name;
          this.mergeMode = args.mergeMode === undefined ?
              DEFAULT_BIDIRECTIONAL_MERGE_MODE :
              args.mergeMode;
          checkBidirectionalMergeMode(this.mergeMode);
          if (args.weights) {
              throw new NotImplementedError('weights support is not implemented for Bidirectional layer yet.');
          }
          this._stateful = args.layer.stateful;
          this.returnSequences = args.layer.returnSequences;
          this.returnState = args.layer.returnState;
          this.supportsMasking = true;
          this._trainable = true;
          this.inputSpec = args.layer.inputSpec;
          this.numConstants = null;
      }
      get trainable() {
          return this._trainable;
      }
      set trainable(value) {
          // Porting Note: the check of `this.layer` here is necessary due to the
          //   way the `constructor` of this class is written (see Porting Note
          //   above).
          this._trainable = value;
          if (this.forwardLayer != null) {
              this.forwardLayer.trainable = value;
          }
          if (this.backwardLayer != null) {
              this.backwardLayer.trainable = value;
          }
      }
      getWeights() {
          return this.forwardLayer.getWeights().concat(this.backwardLayer.getWeights());
      }
      setWeights(weights) {
          const numWeights = weights.length;
          const numeightsOver2 = Math.floor(numWeights / 2);
          this.forwardLayer.setWeights(weights.slice(0, numeightsOver2));
          this.backwardLayer.setWeights(weights.slice(numeightsOver2));
      }
      computeOutputShape(inputShape) {
          let layerShapes = this.forwardLayer.computeOutputShape(inputShape);
          if (!(Array.isArray(layerShapes) && Array.isArray(layerShapes[0]))) {
              layerShapes = [layerShapes];
          }
          layerShapes = layerShapes;
          let outputShape;
          let outputShapes;
          let stateShape;
          if (this.returnState) {
              stateShape = layerShapes.slice(1);
              outputShape = layerShapes[0];
          }
          else {
              outputShape = layerShapes[0];
          }
          outputShape = outputShape;
          if (this.mergeMode === 'concat') {
              outputShape[outputShape.length - 1] *= 2;
              outputShapes = [outputShape];
          }
          else if (this.mergeMode == null) {
              outputShapes = [outputShape, outputShape.slice()];
          }
          else {
              outputShapes = [outputShape];
          }
          if (this.returnState) {
              if (this.mergeMode == null) {
                  return outputShapes.concat(stateShape).concat(stateShape.slice());
              }
              return [outputShape].concat(stateShape).concat(stateShape.slice());
          }
          return generic_utils.singletonOrArray(outputShapes);
      }
      apply(inputs, kwargs) {
          let initialState = kwargs == null ? null : kwargs['initialState'];
          let constants = kwargs == null ? null : kwargs['constants'];
          if (kwargs == null) {
              kwargs = {};
          }
          const standardized = standardizeArgs(inputs, initialState, constants, this.numConstants);
          inputs = standardized.inputs;
          initialState = standardized.initialState;
          constants = standardized.constants;
          if (Array.isArray(inputs)) {
              initialState = inputs.slice(1);
              inputs = inputs[0];
          }
          if ((initialState == null || initialState.length === 0) &&
              constants == null) {
              return super.apply(inputs, kwargs);
          }
          const additionalInputs = [];
          const additionalSpecs = [];
          if (initialState != null) {
              const numStates = initialState.length;
              if (numStates % 2 > 0) {
                  throw new ValueError('When passing `initialState` to a Bidrectional RNN, ' +
                      'the state should be an Array containing the states of ' +
                      'the underlying RNNs.');
              }
              kwargs['initialState'] = initialState;
              additionalInputs.push(...initialState);
              const stateSpecs = initialState
                  .map(state => new InputSpec({ shape: state.shape }));
              this.forwardLayer.stateSpec = stateSpecs.slice(0, numStates / 2);
              this.backwardLayer.stateSpec = stateSpecs.slice(numStates / 2);
              additionalSpecs.push(...stateSpecs);
          }
          if (constants != null) {
              throw new NotImplementedError('Support for constants in Bidirectional layers is not ' +
                  'implemented yet.');
          }
          const isSymbolicTensor = additionalInputs[0] instanceof SymbolicTensor;
          for (const tensor of additionalInputs) {
              if (tensor instanceof SymbolicTensor !== isSymbolicTensor) {
                  throw new ValueError('The initial state of a Bidirectional layer cannot be ' +
                      'specified as a mix of symbolic and non-symbolic tensors');
              }
          }
          if (isSymbolicTensor) {
              // Compute the full input and specs, including the states.
              const fullInput = [inputs].concat(additionalInputs);
              const fullInputSpec = this.inputSpec.concat(additionalSpecs);
              // Perform the call temporarily and replace inputSpec.
              // Note: with initial states symbolic calls and non-symbolic calls to
              // this method differ in how the initial states are passed. For
              // symbolic calls, the initial states are passed in the first arg, as
              // an Array of SymbolicTensors; for non-symbolic calls, they are
              // passed in the second arg as a part of the kwargs. Hence the need to
              // temporarily modify inputSpec here.
              // TODO(cais): Make refactoring so that this hacky code below is no
              // longer needed.
              const originalInputSpec = this.inputSpec;
              this.inputSpec = fullInputSpec;
              const output = super.apply(fullInput, kwargs);
              this.inputSpec = originalInputSpec;
              return output;
          }
          else {
              return super.apply(inputs, kwargs);
          }
      }
      call(inputs, kwargs) {
          return tidy(() => {
              const initialState = kwargs['initialState'];
              let y;
              let yRev;
              if (initialState == null) {
                  y = this.forwardLayer.call(inputs, kwargs);
                  yRev = this.backwardLayer.call(inputs, kwargs);
              }
              else {
                  const forwardState = initialState.slice(0, initialState.length / 2);
                  const backwardState = initialState.slice(initialState.length / 2);
                  y = this.forwardLayer.call(inputs, Object.assign(kwargs, { initialState: forwardState }));
                  yRev = this.backwardLayer.call(inputs, Object.assign(kwargs, { initialState: backwardState }));
              }
              let states;
              if (this.returnState) {
                  if (Array.isArray(y)) {
                      states = y.slice(1).concat(yRev.slice(1));
                  }
                  else {
                  }
                  y = y[0];
                  yRev = yRev[0];
              }
              if (this.returnSequences) {
                  yRev = tfc.reverse(yRev, 1);
              }
              let output;
              if (this.mergeMode === 'concat') {
                  output = K.concatenate([y, yRev]);
              }
              else if (this.mergeMode === 'sum') {
                  output = tfc.add(y, yRev);
              }
              else if (this.mergeMode === 'ave') {
                  output = tfc.mul(.5, tfc.add(y, yRev));
              }
              else if (this.mergeMode === 'mul') {
                  output = tfc.mul(y, yRev);
              }
              else if (this.mergeMode == null) {
                  output = [y, yRev];
              }
              // TODO(cais): Properly set learning phase.
              if (this.returnState) {
                  if (this.mergeMode == null) {
                      return output.concat(states);
                  }
                  return [output].concat(states);
              }
              return output;
          });
      }
      resetStates(states) {
          this.forwardLayer.resetStates();
          this.backwardLayer.resetStates();
      }
      build(inputShape) {
          nameScope(this.forwardLayer.name, () => {
              this.forwardLayer.build(inputShape);
          });
          nameScope(this.backwardLayer.name, () => {
              this.backwardLayer.build(inputShape);
          });
          this.built = true;
      }
      computeMask(inputs, mask) {
          if (Array.isArray(mask)) {
              mask = mask[0];
          }
          let outputMask;
          if (this.returnSequences) {
              if (this.mergeMode == null) {
                  outputMask = [mask, mask];
              }
              else {
                  outputMask = mask;
              }
          }
          else {
              if (this.mergeMode == null) {
                  outputMask = [null, null];
              }
              else {
                  outputMask = null;
              }
          }
          if (this.returnState) {
              const states = this.forwardLayer.states;
              const stateMask = states.map(state => null);
              if (Array.isArray(outputMask)) {
                  return outputMask.concat(stateMask).concat(stateMask);
              }
              else {
                  return [outputMask].concat(stateMask).concat(stateMask);
              }
          }
          else {
              return outputMask;
          }
      }
      get trainableWeights() {
          return this.forwardLayer.trainableWeights.concat(this.backwardLayer.trainableWeights);
      }
      get nonTrainableWeights() {
          return this.forwardLayer.nonTrainableWeights.concat(this.backwardLayer.nonTrainableWeights);
      }
      // TODO(cais): Implement constraints().
      setFastWeightInitDuringBuild(value) {
          super.setFastWeightInitDuringBuild(value);
          if (this.forwardLayer != null) {
              this.forwardLayer.setFastWeightInitDuringBuild(value);
          }
          if (this.backwardLayer != null) {
              this.backwardLayer.setFastWeightInitDuringBuild(value);
          }
      }
      getConfig() {
          const config = {
              'mergeMode': this.mergeMode,
          };
          // TODO(cais): Add logic for `numConstants` once the property is added.
          const baseConfig = super.getConfig();
          Object.assign(config, baseConfig);
          return config;
      }
      /** @nocollapse */
      static fromConfig(cls, config) {
          const rnnLayer = deserialize(config['layer']);
          delete config['layer'];
          // TODO(cais): Add logic for `numConstants` once the property is added.
          if (config['numConstants'] != null) {
              throw new NotImplementedError(`Deserialization of a Bidirectional layer with numConstants ` +
                  `present is not supported yet.`);
          }
          // tslint:disable-next-line:no-any
          const newConfig = config;
          newConfig['layer'] = rnnLayer;
          return new cls(newConfig);
      }
  }
  /** @nocollapse */
  Bidirectional.className = 'Bidirectional';
  serialization.registerClass(Bidirectional);
  //# sourceMappingURL=wrappers.js.map