/**
 * @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.
 * =============================================================================
 */

/**
 * Unit tests for wrapper layers.
 */

import {ones, scalar, serialization, Tensor, tensor2d, Tensor3D, tensor3d} from '@tensorflow/tfjs-core';

import {Layer, SymbolicTensor} from '../engine/topology';
import {LayersModel} from '../engine/training';
import * as tfl from '../index';
import {BidirectionalMergeMode, VALID_BIDIRECTIONAL_MERGE_MODES} from '../keras_format/common';
import {convertPythonicToTs} from '../utils/serialization_utils';
import {describeMathCPU, describeMathCPUAndWebGL2, expectTensorsClose} from '../utils/test_utils';

import {Dense, Reshape} from './core';
import {RNN, SimpleRNN} from './recurrent';
import {deserialize} from './serialization';
import {Bidirectional, checkBidirectionalMergeMode, TimeDistributed} from './wrappers';

describeMathCPU('TimeDistributed Layer: Symbolic', () => {
  it('3D input: Dense', () => {
    const input = new tfl.SymbolicTensor('float32', [10, 8, 2], null, [], null);
    const wrapper = tfl.layers.timeDistributed({layer: new Dense({units: 3})});
    const output = wrapper.apply(input) as tfl.SymbolicTensor;
    expect(wrapper.trainable).toEqual(true);
    expect(wrapper.getWeights().length).toEqual(2);  // kernel and bias.
    expect(output.dtype).toEqual(input.dtype);
    expect(output.shape).toEqual([10, 8, 3]);
  });
  it('4D input: Reshape', () => {
    const input =
        new tfl.SymbolicTensor('float32', [10, 8, 2, 3], null, [], null);
    const wrapper =
        tfl.layers.timeDistributed({layer: new Reshape({targetShape: [6]})});
    const output = wrapper.apply(input) as tfl.SymbolicTensor;
    expect(output.dtype).toEqual(input.dtype);
    expect(output.shape).toEqual([10, 8, 6]);
  });
  it('2D input leads to exception', () => {
    const input = new tfl.SymbolicTensor('float32', [10, 2], null, [], null);
    const wrapper = tfl.layers.timeDistributed({layer: new Dense({units: 3})});
    expect(() => wrapper.apply(input))
        .toThrowError(
            /TimeDistributed .*expects an input shape >= 3D, .* \[10,.*2\]/);
  });
  it('getConfig and fromConfig: round trip', () => {
    const wrapper = tfl.layers.timeDistributed({layer: new Dense({units: 3})});
    const config = wrapper.getConfig();
    const wrapperPrime = TimeDistributed.fromConfig(TimeDistributed, config);
    expect(wrapperPrime.getConfig()).toEqual(wrapper.getConfig());
  });
});

describeMathCPUAndWebGL2('TimeDistributed Layer: Tensor', () => {
  it('3D input: Dense', () => {
    const input = tensor3d(
        [
          [[1, 2], [3, 4], [5, 6], [7, 8]],
          [[-1, -2], [-3, -4], [-5, -6], [-7, -8]]
        ],
        [2, 4, 2]);
    // Given an all-ones Dense kernel and no bias, the output at each timestep
    // is expected to be [3, 7, 11, 15], give or take a minus sign.
    const wrapper = tfl.layers.timeDistributed({
      layer: new Dense({units: 1, kernelInitializer: 'ones', useBias: false})
    });
    const output = wrapper.apply(input) as Tensor;
    expectTensorsClose(
        output,
        tensor3d(
            [[[3], [7], [11], [15]], [[-3], [-7], [-11], [-15]]], [2, 4, 1]));
  });

  // Reference Python Keras code:
  // ```py
  // import keras
  // import numpy as np
  //
  // model_as_layer = keras.Sequential(
  //     layers=[keras.layers.Dense(
  //         units=3,
  //         input_shape=[2, 5],
  //         kernel_initializer='ones',
  //         activation='softmax')])
  //
  // td = keras.layers.TimeDistributed(
  //     layer=model_as_layer, input_shape=[2, 5])
  // model = keras.Sequential(layers=[td])
  // model.summary()
  //
  // xs = np.ones([1, 2, 5])
  // ys = model.predict(xs)
  // print(ys)
  // ```
  it('LayersModel as constituent layer', () => {
    const modelAsLayer = tfl.sequential({
      layers: [tfl.layers.dense({
        activation: 'softmax',
        units: 3,
        inputShape: [2, 5],
        kernelInitializer: 'ones'
      })],
    });
    const td =
        tfl.layers.timeDistributed({layer: modelAsLayer, inputShape: [2, 5]});
    const model = tfl.sequential({layers: [td]});
    const ys = model.predict(ones([1, 2, 5])) as Tensor;
    expect(ys.shape).toEqual([1, 2, 3]);
    expectTensorsClose(ys, tensor3d([[
                         [0.33333334, 0.33333334, 0.33333334],
                         [0.33333334, 0.33333334, 0.33333334]
                       ]]));
  });

  // Reference Python code:
  // ```py
  // import keras
  // import numpy as np
  //
  // input1 = keras.Input(shape=[4])
  // input2 = keras.Input(shape=[4])
  // y1 = keras.layers.Embedding(10,
  //                             3,
  //                             input_length=4,
  //                             mask_zero=True,
  //                             embeddings_initializer='ones')(input1)
  // y1 = keras.layers.LSTM(3,
  //                        return_sequences=True,
  //                        recurrent_initializer='ones',
  //                        kernel_initializer='ones',
  //                        bias_initializer='zeros')(y1)
  // z1 = keras.layers.Embedding(10,
  //                             3,
  //                             input_length=4,
  //                             mask_zero=True,
  //                             embeddings_initializer='ones')(input2)
  // z1 = keras.layers.LSTM(3,
  //                        return_sequences=True,
  //                        recurrent_initializer='ones',
  //                        kernel_initializer='ones',
  //                        bias_initializer='zeros')(z1)
  // y2 = keras.layers.Dot(axes=[2, 2])([y1, z1])
  // y3 = keras.layers.Concatenate()([y1, y2])
  // y3 = keras.layers.TimeDistributed(keras.layers.Dense(
  //     1,
  //     kernel_initializer='ones',
  //     bias_initializer='zeros'))(y3)
  // model = keras.Model(inputs=[input1, input2], outputs=y3)
  // model.summary()
  //
  // xs1 = np.array([[0, 0, 0, 0], [1, 0, 0, 0], [1, 2, 0, 0], [1, 2, 3, 0]])
  // xs2 = np.ones([4, 4])
  // ys = model.predict([xs1, xs2])
  // print(ys)
  // ```
  it('With masking', () => {
    const input1 = tfl.input({shape: [4]});
    const input2 = tfl.input({shape: [4]});
    let y1 = tfl.layers
                 .embedding({
                   inputDim: 10,
                   outputDim: 3,
                   inputLength: 4,
                   maskZero: true,
                   embeddingsInitializer: 'ones'
                 })
                 .apply(input1) as SymbolicTensor;
    y1 = tfl.layers
             .lstm({
               units: 3,
               returnSequences: true,
               recurrentInitializer: 'ones',
               kernelInitializer: 'ones',
               biasInitializer: 'zeros'
             })
             .apply(y1) as SymbolicTensor;
    let z1 = tfl.layers
                 .embedding({
                   inputDim: 10,
                   outputDim: 3,
                   inputLength: 4,
                   maskZero: true,
                   embeddingsInitializer: 'ones'
                 })
                 .apply(input2) as SymbolicTensor;
    z1 = tfl.layers
             .lstm({
               units: 3,
               returnSequences: true,
               recurrentInitializer: 'ones',
               kernelInitializer: 'ones',
               biasInitializer: 'zeros'
             })
             .apply(z1) as SymbolicTensor;
    const y2 = tfl.layers.dot({axes: [2, 2]}).apply([y1, z1]) as SymbolicTensor;
    let y3 = tfl.layers.concatenate().apply([y1, y2]) as SymbolicTensor;
    y3 = tfl.layers
             .timeDistributed({
               layer: tfl.layers.dense({
                 units: 1,
                 kernelInitializer: 'ones',
                 biasInitializer: 'zeros'
               })
             })
             .apply(y3) as SymbolicTensor;
    const model = tfl.model({inputs: [input1, input2], outputs: y3});

    const xs1 =
        tensor2d([[0, 0, 0, 0], [1, 0, 0, 0], [1, 2, 0, 0], [1, 2, 3, 0]]);
    const xs2 = ones([4, 4]);
    const ys = model.predict([xs1, xs2]) as Tensor;
    expectTensorsClose(
        ys, tensor3d([
          [[0], [0], [0], [0]],
          [[10.7489796], [10.7489796], [10.7489796], [10.7489796]],
          [[10.7489796], [13.6384077], [13.6384077], [13.6384077]],
          [[10.7489796], [13.6384077], [14.0818386], [14.0818386]]
        ]));
  });
});

describeMathCPU('Bidirectional Layer: Symbolic', () => {
  const mergeModes: BidirectionalMergeMode[] = [
    null,
    'concat',
    'ave',
    'mul',
    'sum',
  ];
  const returnStateValues: boolean[] = [false, true];

  for (const mergeMode of mergeModes) {
    for (const returnState of returnStateValues) {
      const testTitle = `3D input: returnSequence=false, ` +
          `mergeMode=${mergeMode}; returnState=${returnState}`;
      it(testTitle, () => {
        const input =
            new tfl.SymbolicTensor('float32', [10, 8, 2], null, [], null);
        const bidi = tfl.layers.bidirectional({
          layer: new SimpleRNN(
              {units: 3, recurrentInitializer: 'glorotNormal', returnState}),
          mergeMode,
        });
        // TODO(cais): Remove recurrentInitializer once Orthogonal initializer
        // is available.
        let outputs = bidi.apply(input);
        expect(bidi.trainable).toEqual(true);
        // {kernel, recurrentKernel, bias} * {forward, backward}.
        expect(bidi.getWeights().length).toEqual(6);
        if (!returnState) {
          if (mergeMode === null) {
            outputs = outputs as tfl.SymbolicTensor[];
            expect(outputs.length).toEqual(2);
            expect(outputs[0].shape).toEqual([10, 3]);
            expect(outputs[1].shape).toEqual([10, 3]);
          } else if (mergeMode === 'concat') {
            outputs = outputs as tfl.SymbolicTensor;
            expect(outputs.shape).toEqual([10, 6]);
          } else {
            outputs = outputs as tfl.SymbolicTensor;
            expect(outputs.shape).toEqual([10, 3]);
          }
        } else {
          if (mergeMode === null) {
            outputs = outputs as tfl.SymbolicTensor[];
            expect(outputs.length).toEqual(4);
            expect(outputs[0].shape).toEqual([10, 3]);
            expect(outputs[1].shape).toEqual([10, 3]);
            expect(outputs[2].shape).toEqual([10, 3]);
            expect(outputs[3].shape).toEqual([10, 3]);
          } else if (mergeMode === 'concat') {
            outputs = outputs as tfl.SymbolicTensor[];
            expect(outputs.length).toEqual(3);
            expect(outputs[0].shape).toEqual([10, 6]);
            expect(outputs[1].shape).toEqual([10, 3]);
            expect(outputs[2].shape).toEqual([10, 3]);
          } else {
            outputs = outputs as tfl.SymbolicTensor[];
            expect(outputs.length).toEqual(3);
            expect(outputs[0].shape).toEqual([10, 3]);
            expect(outputs[1].shape).toEqual([10, 3]);
            expect(outputs[2].shape).toEqual([10, 3]);
          }
        }
      });
    }
  }
  it('returnSequence=true', () => {
    const input = new tfl.SymbolicTensor('float32', [10, 8, 2], null, [], null);
    const bidi = tfl.layers.bidirectional({
      layer: new SimpleRNN({
        units: 3,
        recurrentInitializer: 'glorotNormal',
        returnSequences: true,
        returnState: true
      }),
      mergeMode: 'ave'
    });
    const outputs = bidi.apply(input) as tfl.SymbolicTensor[];
    expect(outputs.length).toEqual(3);
    expect(outputs[0].shape).toEqual([10, 8, 3]);
    expect(outputs[1].shape).toEqual([10, 3]);
    expect(outputs[2].shape).toEqual([10, 3]);
  });
  it('Serialization round trip', () => {
    const layer = tfl.layers.bidirectional({
      layer: new SimpleRNN({units: 3}),
      mergeMode: 'concat',
      inputShape: [4, 4],
    });
    const model = tfl.sequential({layers: [layer]});
    const unused: {} = null;
    const modelJSON = model.toJSON(unused, false);
    const modelPrime = deserialize(
                           convertPythonicToTs(modelJSON) as
                           serialization.ConfigDict) as tfl.Sequential;
    expect((modelPrime.layers[0] as Bidirectional).getConfig().mergeMode)
        .toEqual('concat');
  });
  it('mergeMode defaults to concat', () => {
    const bidi = tfl.layers.bidirectional({layer: new SimpleRNN({units: 3})});
    expect(bidi.mergeMode).toEqual('concat');
    expect(bidi.getConfig().mergeMode).toEqual('concat');
  });
});

describe('checkBidirectionalMergeMode', () => {
  it('Valid values', () => {
    const extendedValues =
        VALID_BIDIRECTIONAL_MERGE_MODES.concat([undefined, null]);
    for (const validValue of extendedValues) {
      // Using implicit "expect().toNotThrow()" for valid values
      checkBidirectionalMergeMode(validValue);
    }
  });
  it('Invalid values', () => {
    // Test invalid values are rejected, and reported in the error.
    expect(() => checkBidirectionalMergeMode('foo')).toThrowError(/foo/);
    try {
      checkBidirectionalMergeMode('bad');
    } catch (e) {
      expect(e).toMatch('BidirectionalMergeMode');
      // Test that the error message contains the list of valid values.
      for (const validValue of VALID_BIDIRECTIONAL_MERGE_MODES) {
        expect(e).toMatch(validValue);
      }
    }
  });
});

describeMathCPUAndWebGL2('Bidirectional Layer: Tensor', () => {
  // The golden tensor values used in the tests below can be obtained with
  // PyKeras code such as the following:
  // ```python
  // import keras
  // import numpy as np
  //
  // bidi = keras.layers.Bidirectional(
  //     keras.layers.SimpleRNN(
  //         3,
  //         kernel_initializer='ones',
  //         recurrent_initializer='ones',
  //         return_state=True),
  //     merge_mode='ave')
  //
  // time_steps = 4
  // input_size = 2
  // inputs = keras.Input([time_steps, input_size])
  // outputs = bidi(inputs)
  // model = keras.Model(inputs, outputs)
  //
  // x = np.array(
  //     [[[0.05, 0.05],
  //       [-0.05, -0.05],
  //       [0.1, 0.1],
  //       [-0.1, -0.1]]])
  // print(model.predict(x))
  // ```

  // TODO(bileschi): This should be tfl.layers.Layer.
  let bidi: Layer;
  let x: Tensor3D;
  function createLayerAndData(
      mergeMode: BidirectionalMergeMode, returnState: boolean) {
    const units = 3;
    bidi = tfl.layers.bidirectional({
      layer: new SimpleRNN({
        units,
        kernelInitializer: 'ones',
        recurrentInitializer: 'ones',
        useBias: false,
        returnState
      }),
      mergeMode,
    });
    const timeSteps = 4;
    const inputSize = 2;
    x = tensor3d(
        [[[0.05, 0.05], [-0.05, -0.05], [0.1, 0.1], [-0.1, -0.1]]],
        [1, timeSteps, inputSize]);
  }

  const mergeModes: BidirectionalMergeMode[] =
      [null, undefined, 'concat', 'mul'];
  for (const mergeMode of mergeModes) {
    it(`No returnState, mergeMode=${mergeMode}`, () => {
      createLayerAndData(mergeMode, false);
      let y = bidi.apply(x);
      if (mergeMode === null) {
        y = y as Tensor[];
        expect(y.length).toEqual(2);
        expectTensorsClose(
            y[0], tensor2d([[0.9440416, 0.9440416, 0.9440416]], [1, 3]));
        expectTensorsClose(
            y[1], tensor2d([[-0.9842659, -0.9842659, -0.9842659]], [1, 3]));
      } else if (mergeMode === undefined || mergeMode === 'concat') {
        y = y as Tensor;
        expectTensorsClose(
            y,
            tensor2d(
                [[
                  0.9440416, 0.9440416, 0.9440416, -0.9842659, -0.9842659,
                  -0.9842659
                ]],
                [1, 6]));
      } else if (mergeMode === 'mul') {
        y = y as Tensor;
        expectTensorsClose(
            y, tensor2d([[-0.929188, -0.929188, -0.929188]], [1, 3]));
      }
    });
  }
  it('returnState', () => {
    createLayerAndData('ave', true);
    const y = bidi.apply(x) as Tensor[];
    expect(y.length).toEqual(3);
    expectTensorsClose(
        y[0], tensor2d([[-0.02011216, -0.02011216, -0.02011216]], [1, 3]));
    expectTensorsClose(
        y[1], tensor2d([[0.9440416, 0.9440416, 0.9440416]], [1, 3]));
    expectTensorsClose(
        y[2], tensor2d([[-0.9842659, -0.9842659, -0.9842659]], [1, 3]));
  });

  // The golden values in the test below can be obtained with the following
  // Python Keras code.
  //
  // ```python
  // import keras
  // import numpy
  //
  // rnn = keras.layers.LSTM(
  //     1,
  //     kernel_initializer='ones',
  //     recurrent_initializer='ones',
  //     bias_initializer='ones',
  //     go_backwards=True)
  // bidi = keras.layers.Bidirectional(
  //     rnn, merge_mode='concat', input_shape=[2, 2])
  // model = keras.Sequential([bidi])
  // model.compile(loss='mean_squared_error', optimizer='sgd')
  //
  // x = np.array([[[0.1, 0.2],
  //               [-0.1, 0.1]]])
  // y = np.array([[0.3, 0.5]])
  // print(model.predict(x))
  //
  // history = model.fit(x, y)
  // print(history.history)
  // ```
  it('Backwards LSTM: predict and fit: concat', async () => {
    const lstm = tfl.layers.lstm({
      units: 1,
      kernelInitializer: 'ones',
      recurrentInitializer: 'ones',
      biasInitializer: 'ones',
      goBackwards: true
    }) as RNN;
    const bidi = tfl.layers.bidirectional(
        {layer: lstm, inputShape: [2, 2], mergeMode: 'concat'});
    const model = tfl.sequential({layers: [bidi]});
    model.compile({loss: 'meanSquaredError', optimizer: 'sgd'});

    const x = tensor3d([[[0.1, 0.2], [-0.1, 0.1]]]);
    const y = tensor2d([[0.3, 0.5]]);
    expectTensorsClose(
        model.predict(x) as Tensor, tensor2d([[0.69299805, 0.66088]]));
    const history = await model.fit(x, y);
    expect(history.history.loss[0]).toBeCloseTo(0.0901649);
    expectTensorsClose(
        model.predict(x) as Tensor, tensor2d([[0.6927189, 0.66077083]]));
  });

  // The golden values in the test below can be obtained with the following
  // Python Keras code.
  //
  // ```python
  // import keras
  // import numpy
  //
  // rnn = keras.layers.LSTM(
  //     1,
  //     kernel_initializer='ones',
  //     recurrent_initializer='ones',
  //     bias_initializer='ones',
  //     go_backwards=True)
  // bidi = keras.layers.Bidirectional(
  //     rnn, merge_mode='ave', input_shape=[2, 2])
  // model = keras.Sequential([bidi])
  // model.compile(loss='mean_squared_error', optimizer='sgd')
  //
  // x = np.array([[[0.1, 0.2],
  //               [-0.1, 0.1]]])
  // y = np.array([[0.4]])
  // print(model.predict(x))
  //
  // history = model.fit(x, y)
  // print(history.history)
  // ```
  it('Backwards LSTM: predict and fit: ave', done => {
    const lstm = tfl.layers.lstm({
      units: 1,
      kernelInitializer: 'ones',
      recurrentInitializer: 'ones',
      biasInitializer: 'ones',
      goBackwards: true
    }) as RNN;
    const bidi = tfl.layers.bidirectional(
        {layer: lstm, inputShape: [2, 2], mergeMode: 'ave'});
    const model = tfl.sequential({layers: [bidi]});
    model.compile({loss: 'meanSquaredError', optimizer: 'sgd'});

    const x = tensor3d([[[0.1, 0.2], [-0.1, 0.1]]]);
    const y = tensor2d([[0.4]]);
    expectTensorsClose(model.predict(x) as Tensor, tensor2d([[0.676939]]));
    model.fit(x, y)
        .then(history => {
          expect(history.history.loss[0]).toBeCloseTo(0.0766952);
          expectTensorsClose(
              model.predict(x) as Tensor, tensor2d([[0.6767467]]));
          done();
        })
        .catch(err => {
          done.fail(err.stack);
        });
  });
});

describeMathCPUAndWebGL2('Bidirectional with initial state', () => {
  const sequenceLength = 4;
  const recurrentUnits = 3;
  const inputDim = 2;

  function createLayerAndTensors(): {
    bidiLayer: Bidirectional,
    initState1: SymbolicTensor,
    initState2: SymbolicTensor,
    inputTensor: SymbolicTensor,
    outputTensors: SymbolicTensor[],
  } {
    const initState1 = tfl.input({shape: [recurrentUnits]});
    const initState2 = tfl.input({shape: [recurrentUnits]});
    const inputTensor = tfl.input({shape: [sequenceLength, inputDim]});
    const bidiLayer = tfl.layers.bidirectional({
      layer: tfl.layers.gru({
        units: recurrentUnits,
        kernelInitializer: 'zeros',
        recurrentInitializer: 'zeros',
        biasInitializer: 'ones'
      }) as RNN,
      mergeMode: null,
    });
    const outputTensors =
        bidiLayer.apply(
            inputTensor, {initialState: [initState1, initState2]}) as
        SymbolicTensor[];
    return {bidiLayer, initState1, initState2, inputTensor, outputTensors};
  }

  it('Correct shapes', () => {
    const layerAndTensors = createLayerAndTensors();
    expect(layerAndTensors.outputTensors.length).toEqual(2);
    expect(layerAndTensors.outputTensors[0].shape).toEqual([
      null, recurrentUnits
    ]);
    expect(layerAndTensors.outputTensors[1].shape).toEqual([
      null, recurrentUnits
    ]);
  });

  it('apply() with concrete tensors', () => {
    const layerAndTensors = createLayerAndTensors();
    const xVal = ones([1, sequenceLength, inputDim]);
    const initState1Val = ones([1, recurrentUnits]).mul(scalar(-1));
    const initState2Val = ones([1, recurrentUnits]);
    const yVals =
        layerAndTensors.bidiLayer.apply(
            xVal, {initialState: [initState1Val, initState2Val]}) as Tensor[];
    expect(yVals.length).toEqual(2);
    expectTensorsClose(
        yVals[0], tensor2d([[0.33863544, 0.33863544, 0.33863544]]));
    expectTensorsClose(yVals[1], tensor2d([[0.8188354, 0.8188354, 0.8188354]]));
  });

  it('LayersModel predict', () => {
    const layerAndTensors = createLayerAndTensors();
    const model = tfl.model({
      inputs: [
        layerAndTensors.inputTensor, layerAndTensors.initState1,
        layerAndTensors.initState2
      ],
      outputs: layerAndTensors.outputTensors
    });
    const xVal = ones([1, sequenceLength, inputDim]);
    const initState1Val = ones([1, recurrentUnits]).mul(scalar(-1));
    const initState2Val = ones([1, recurrentUnits]);
    const yVals =
        model.predict([xVal, initState1Val, initState2Val]) as Tensor[];
    expect(yVals.length).toEqual(2);
    expectTensorsClose(
        yVals[0], tensor2d([[0.33863544, 0.33863544, 0.33863544]]));
    expectTensorsClose(yVals[1], tensor2d([[0.8188354, 0.8188354, 0.8188354]]));
  });

  it('LayersModel serialization round trip', () => {
    // Disable the console warning about the unserialization SymbolicTensor in
    // the CallArgs.
    spyOn(console, 'warn');
    const layerAndTensors = createLayerAndTensors();
    const model = tfl.model({
      inputs: [
        layerAndTensors.inputTensor, layerAndTensors.initState1,
        layerAndTensors.initState2
      ],
      outputs: layerAndTensors.outputTensors
    });
    const json1 = model.toJSON(null, false);
    const model2 =
        deserialize(convertPythonicToTs(json1) as serialization.ConfigDict) as
        LayersModel;
    const json2 = model2.toJSON(null, false);
    expect(json2).toEqual(json1);
  });

  it('Incorrect number of initial-state tensors leads to error', () => {
    const initState1 = tfl.input({shape: [recurrentUnits]});
    const x = tfl.input({shape: [sequenceLength, inputDim]});
    const bidi = tfl.layers.bidirectional({
      layer: tfl.layers.gru({
        units: recurrentUnits,
      }) as RNN,
      mergeMode: null,
    });
    expect(() => bidi.apply(x, {
      initialState: [initState1]
    })).toThrowError(/the state should be .*RNNs/);
  });
});

describeMathCPUAndWebGL2('Bidirectional with masking', () => {
  for (const rnnLayerType of [tfl.layers.simpleRNN, tfl.layers.lstm]) {
    // Reference Python TensorFlow code (v1.15.0):
    // ```py
    // import numpy as np
    // import tensorflow as tf
    //
    // tf.enable_v2_behavior()
    // tf.enable_eager_execution()
    //
    // model = tf.keras.Sequential()
    // model.add(tf.keras.layers.Embedding(
    //     3, 3, input_length=2, embeddings_initializer="ones", mask_zero=True))
    // model.add(tf.keras.layers.Bidirectional(tf.keras.layers.SimpleRNN(
    //     4,
    //     kernel_initializer="ones",
    //     recurrent_initializer="ones")))
    // model.add(tf.keras.layers.Dense(1, kernel_initializer="ones"))
    //
    // xs = np.array([[0, 1], [0, 2], [0, 0], [1, 2]], dtype=np.int32)
    // print(model.predict(xs))
    // ```
    it(`Inference ${rnnLayerType.name}`, () => {
      const model = tfl.sequential();
      model.add(tfl.layers.embedding({
        inputLength: 2,
        inputDim: 3,
        outputDim: 3,
        embeddingsInitializer: 'ones',
        maskZero: true
      }));
      model.add(tfl.layers.bidirectional({
        layer: rnnLayerType({
                 units: 4,
                 kernelInitializer: 'ones',
                 recurrentInitializer: 'ones'
               }) as RNN
      }));
      model.add(tfl.layers.dense({
        units: 1,
        kernelInitializer: 'ones',
      }));
      const xs = tensor2d([[0, 1], [0, 2], [0, 0], [1, 2]]);
      const ys = model.predict(xs) as Tensor;
      if (rnnLayerType.name === 'lstm') {
        expectTensorsClose(
            ys, tensor2d([[6.076076], [6.076076], [0], [7.709405]]));
      } else if (rnnLayerType.name === 'simpleRNN') {
        expectTensorsClose(
            ys, tensor2d([[7.960438], [7.960438], [0], [7.9999857]]));
      } else {
        throw new Error(`Unexpected rnnLayerType name: ${rnnLayerType.name}`);
      }
    });

    // Reference Python TensorFlow code (v1.15.0):
    // ```py
    // import numpy as np
    // import tensorflow as tf
    //
    // tf.enable_v2_behavior()
    // tf.enable_eager_execution()
    //
    // model = tf.keras.Sequential()
    // model.add(tf.keras.layers.Embedding(
    //     3, 3, input_length=2, embeddings_initializer="ones", mask_zero=True))
    // model.add(tf.keras.layers.Bidirectional(tf.keras.layers.SimpleRNN(
    //     4,
    //     kernel_initializer="ones",
    //     recurrent_initializer="ones")))
    // model.add(tf.keras.layers.Dense(1, kernel_initializer="ones"))
    // model.compile(optimizer="sgd", loss="mse")
    //
    // xs = np.array([[0, 1], [0, 2], [0, 0], [1, 2]], dtype=np.int32)
    // ys = np.array([[0], [0], [0], [0]], dtype=np.float32)
    // history = model.fit(xs, ys, epochs=5)
    // print(history.history)
    // ```
    it(`Training ${rnnLayerType.name}`, async () => {
      const model = tfl.sequential();
      model.add(tfl.layers.embedding({
        inputLength: 2,
        inputDim: 3,
        outputDim: 3,
        embeddingsInitializer: 'ones',
        maskZero: true
      }));
      model.add(tfl.layers.bidirectional({
        layer: rnnLayerType({
                 units: 4,
                 kernelInitializer: 'ones',
                 recurrentInitializer: 'ones'
               }) as RNN
      }));
      model.add(tfl.layers.dense({
        units: 1,
        kernelInitializer: 'ones',
      }));
      model.compile({loss: 'meanSquaredError', optimizer: 'sgd'});
      const xs = tensor2d([[0, 1], [0, 2], [0, 0], [1, 2]]);
      const ys = tensor2d([[0], [0], [0], [0]]);
      const history = await model.fit(xs, ys, {epochs: 5});
      if (rnnLayerType.name === 'lstm') {
        expectTensorsClose(
            history.history['loss'] as number[],
            [33.318, 27.104, 22.058, 17.961, 14.634]);
      } else if (rnnLayerType.name === 'simpleRNN') {
        expectTensorsClose(
            history.history['loss'] as number[],
            [47.684, 35.741, 26.804, 20.115, 15.108]);
      } else {
        throw new Error(`Unexpected rnnLayerType name:
          ${rnnLayerType.name}`);
      }
    });
  }
});
