/**

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

* =============================================================================

*/

/**

* TensorFlow.js Layers: Basic Layers.

*/

import {any, notEqual, serialization, Tensor, tidy, transpose, util} from '@tensorflow/tfjs-core';

import {Activation as ActivationFn, getActivation, serializeActivation} from '../activations';

import * as K from '../backend/tfjs_backend';

import {Constraint, ConstraintIdentifier, getConstraint, serializeConstraint} from '../constraints';

import {DisposeResult, InputSpec, Layer, LayerArgs} from '../engine/topology';

import {ValueError} from '../errors';

import {getInitializer, Initializer, InitializerIdentifier, serializeInitializer} from '../initializers';

import {ActivationIdentifier} from '../keras_format/activation_config';

import {DataFormat, Shape} from '../keras_format/common';

import {LayerConfig} from '../keras_format/topology_config';

import {getRegularizer, Regularizer, RegularizerIdentifier, serializeRegularizer} from '../regularizers';

import {Kwargs} from '../types';

import {assertPositiveInteger, mapActivationToFusedKernel} from '../utils/generic_utils';

import {arrayProd, range} from '../utils/math_utils';

import {getExactlyOneShape, getExactlyOneTensor} from '../utils/types_utils';

import {LayerVariable} from '../variables';

export declare interface DropoutLayerArgs extends LayerArgs {

/** Float between 0 and 1. Fraction of the input units to drop. */

rate: number;

/**

* Integer array representing the shape of the binary dropout mask that will

* be multiplied with the input.

*

* For instance, if your inputs have shape `(batchSize, timesteps, features)`

* and you want the dropout mask to be the same for all timesteps, you can use

* `noise_shape=(batch_size, 1, features)`.

*/

noiseShape?: number[];

/** An integer to use as random seed. */

seed?: number;

}

export class Dropout extends Layer {

/** @nocollapse */

static className = 'Dropout';

private readonly rate: number;

private readonly noiseShape: number[];

private readonly seed: number;

constructor(args: DropoutLayerArgs) {

super(args);

this.rate = Math.max(Math.min(args.rate, 1), 0);

// So that the scalar doesn't get tidied up between executions.

this.noiseShape = args.noiseShape;

this.seed = args.seed;

this.supportsMasking = true;

}

protected getNoiseShape(input: Tensor): Shape {

if (this.noiseShape == null) {

return this.noiseShape;

}

const inputShape = input.shape;

const noiseShape: Shape = [];

for (let i = 0; i < this.noiseShape.length; ++i) {

noiseShape.push(

this.noiseShape[i] == null ? inputShape[i] : this.noiseShape[i]);

}

return noiseShape;

}

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

return tidy(() => {

this.invokeCallHook(inputs, kwargs);

const input = getExactlyOneTensor(inputs);

if (0 < this.rate && this.rate < 1) {

const training =

kwargs['training'] == null ? false : kwargs['training'];

const noiseShape = this.getNoiseShape(input);

const output = K.inTrainPhase(

() => K.dropout(input, this.rate, noiseShape, this.seed),

() => input, training);

return output;

}

return inputs;

});

}

getConfig(): serialization.ConfigDict {

const config = {

rate: this.rate,

noiseShape: this.noiseShape,

seed: this.seed,

};

const baseConfig = super.getConfig();

Object.assign(config, baseConfig);

return config;

}

dispose(): DisposeResult {

return super.dispose();

}

}

serialization.registerClass(Dropout);

export declare interface DenseLayerArgs extends LayerArgs {

/** Positive integer, dimensionality of the output space. */

units: number;

/**

* Activation function to use.

*

* If unspecified, no activation is applied.

*/

activation?: ActivationIdentifier;

/** Whether to apply a bias. */

useBias?: boolean;

/**

* Initializer for the dense kernel weights matrix.

*/

kernelInitializer?: InitializerIdentifier|Initializer;

/**

* Initializer for the bias vector.

*/

biasInitializer?: InitializerIdentifier|Initializer;

/**

* If specified, defines inputShape as `[inputDim]`.

*/

inputDim?: number;

/**

* Constraint for the kernel weights.

*/

kernelConstraint?: ConstraintIdentifier|Constraint;

/**

* Constraint for the bias vector.

*/

biasConstraint?: ConstraintIdentifier|Constraint;

/**

* Regularizer function applied to the dense kernel weights matrix.

*/

kernelRegularizer?: RegularizerIdentifier|Regularizer;

/**

* Regularizer function applied to the bias vector.

*/

biasRegularizer?: RegularizerIdentifier|Regularizer;

/**

* Regularizer function applied to the activation.

*/

activityRegularizer?: RegularizerIdentifier|Regularizer;

}

export interface SpatialDropout1DLayerConfig extends LayerConfig {

/** Float between 0 and 1. Fraction of the input units to drop. */

rate: number;

/** An integer to use as random seed. */

seed?: number;

}

export class SpatialDropout1D extends Dropout {

/** @nocollapse */

static className = 'SpatialDropout1D';

constructor(args: SpatialDropout1DLayerConfig) {

super(args);

this.inputSpec = [{ndim: 3}];

}

protected getNoiseShape(input: Tensor): Shape {

const inputShape = input.shape;

return [inputShape[0], 1, inputShape[2]];

}

}

serialization.registerClass(SpatialDropout1D);

export class Dense extends Layer {

/** @nocollapse */

static className = 'Dense';

private units: number;

// Default activation: Linear (none).

private activation: ActivationFn = null;

private useBias = true;

private kernelInitializer: Initializer;

private biasInitializer: Initializer;

private kernel: LayerVariable = null;

private bias: LayerVariable = null;

readonly DEFAULT_KERNEL_INITIALIZER: InitializerIdentifier = 'glorotNormal';

readonly DEFAULT_BIAS_INITIALIZER: InitializerIdentifier = 'zeros';

private readonly kernelConstraint?: Constraint;

private readonly biasConstraint?: Constraint;

private readonly kernelRegularizer?: Regularizer;

private readonly biasRegularizer?: Regularizer;

constructor(args: DenseLayerArgs) {

super(args);

if (args.batchInputShape == null && args.inputShape == null &&

args.inputDim != null) {

// This logic is copied from Layer's constructor, since we can't

// do exactly what the Python constructor does for Dense().

let batchSize: number = null;

if (args.batchSize != null) {

batchSize = args.batchSize;

}

this.batchInputShape = [batchSize, args.inputDim];

}

this.units = args.units;

assertPositiveInteger(this.units, 'units');

this.activation = getActivation(args.activation);

if (args.useBias != null) {

this.useBias = args.useBias;

}

this.kernelInitializer = getInitializer(

args.kernelInitializer || this.DEFAULT_KERNEL_INITIALIZER);

this.biasInitializer =

getInitializer(args.biasInitializer || this.DEFAULT_BIAS_INITIALIZER);

this.kernelConstraint = getConstraint(args.kernelConstraint);

this.biasConstraint = getConstraint(args.biasConstraint);

this.kernelRegularizer = getRegularizer(args.kernelRegularizer);

this.biasRegularizer = getRegularizer(args.biasRegularizer);

this.activityRegularizer = getRegularizer(args.activityRegularizer);

this.supportsMasking = true;

this.inputSpec = [{minNDim: 2}];

}

public build(inputShape: Shape|Shape[]): void {

inputShape = getExactlyOneShape(inputShape);

const inputLastDim = inputShape[inputShape.length - 1];

if (this.kernel == null) {

this.kernel = this.addWeight(

'kernel', [inputLastDim, this.units], null, this.kernelInitializer,

this.kernelRegularizer, true, this.kernelConstraint);

if (this.useBias) {

this.bias = this.addWeight(

'bias', [this.units], null, this.biasInitializer,

this.biasRegularizer, true, this.biasConstraint);

}

}

this.inputSpec = [{minNDim: 2, axes: {[-1]: inputLastDim}}];

this.built = true;

}

computeOutputShape(inputShape: Shape|Shape[]): Shape|Shape[] {

inputShape = getExactlyOneShape(inputShape);

const outputShape = inputShape.slice();

outputShape[outputShape.length - 1] = this.units;

return outputShape;

}

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

return tidy(() => {

this.invokeCallHook(inputs, kwargs);

// Dense layer accepts only a single input.

const input = getExactlyOneTensor(inputs);

const fusedActivationName =

mapActivationToFusedKernel(this.activation.getClassName());

let output: Tensor;

if (fusedActivationName != null) {

output = K.dot(

input, this.kernel.read(), fusedActivationName,

this.bias ? this.bias.read() : null);

} else {

output = K.dot(input, this.kernel.read());

if (this.bias != null) {

output = K.biasAdd(output, this.bias.read());

}

if (this.activation != null) {

output = this.activation.apply(output);

}

}

return output;

});

}

getConfig(): serialization.ConfigDict {

const config: serialization.ConfigDict = {

units: this.units,

activation: serializeActivation(this.activation),

useBias: this.useBias,

kernelInitializer: serializeInitializer(this.kernelInitializer),

biasInitializer: serializeInitializer(this.biasInitializer),

kernelRegularizer: serializeRegularizer(this.kernelRegularizer),

biasRegularizer: serializeRegularizer(this.biasRegularizer),

activityRegularizer: serializeRegularizer(this.activityRegularizer),

kernelConstraint: serializeConstraint(this.kernelConstraint),

biasConstraint: serializeConstraint(this.biasConstraint)

};

const baseConfig = super.getConfig();

Object.assign(config, baseConfig);

return config;

}

}

serialization.registerClass(Dense);

export declare interface FlattenLayerArgs extends LayerArgs {

/** Image data format: channeLast (default) or channelFirst. */

dataFormat?: DataFormat;

}

export class Flatten extends Layer {

private dataFormat: DataFormat;

/** @nocollapse */

static className = 'Flatten';

constructor(args?: FlattenLayerArgs) {

args = args || {};

super(args);

this.inputSpec = [{minNDim: 3}];

this.dataFormat = args.dataFormat;

}

computeOutputShape(inputShape: Shape|Shape[]): Shape|Shape[] {

inputShape = getExactlyOneShape(inputShape);

for (const dim of inputShape.slice(1)) {

if (dim == null) {

throw new ValueError(

`The shape of the input to "Flatten" is not fully defined ` +

`(got ${inputShape.slice(1)}). Make sure to pass a complete ` +

`"input_shape" or "batch_input_shape" argument to the first ` +

`layer in your model.`);

}

}

return [inputShape[0], arrayProd(inputShape, 1)];

}

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

return tidy(() => {

this.invokeCallHook(inputs, kwargs);

let input = getExactlyOneTensor(inputs);

if (this.dataFormat === 'channelsFirst' && input.rank > 1) {

const permutation: number[] = [0];

for (let i = 2; i < input.rank; ++i) {

permutation.push(i);

}

permutation.push(1);

input = input.transpose(permutation);

}

return K.batchFlatten(input);

});

}

getConfig(): serialization.ConfigDict {

const config: serialization.ConfigDict = {};

if (this.dataFormat != null) {

config['dataFormat'] = this.dataFormat;

}

const baseConfig = super.getConfig();

Object.assign(config, baseConfig);

return config;

}

}

serialization.registerClass(Flatten);

export declare interface ActivationLayerArgs extends LayerArgs {

/**

* Name of the activation function to use.

*/

activation: ActivationIdentifier;

}

export class Activation extends Layer {

/** @nocollapse */

static className = 'Activation';

activation: ActivationFn;

constructor(args: ActivationLayerArgs) {

super(args);

this.supportsMasking = true;

this.activation = getActivation(args.activation);

}

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

return tidy(() => {

this.invokeCallHook(inputs, kwargs);

const input = getExactlyOneTensor(inputs);

return this.activation.apply(input);

});

}

getConfig(): serialization.ConfigDict {

const config = {activation: serializeActivation(this.activation)};

const baseConfig = super.getConfig();

Object.assign(config, baseConfig);

return config;

}

}

serialization.registerClass(Activation);

export declare interface ReshapeLayerArgs extends LayerArgs {

/** The target shape. Does not include the batch axis. */

targetShape: Shape;

}

export declare interface RepeatVectorLayerArgs extends LayerArgs {

/**

* The integer number of times to repeat the input.

*/

n: number;

}

export class RepeatVector extends Layer {

/** @nocollapse */

static className = 'RepeatVector';

readonly n: number;

constructor(args: RepeatVectorLayerArgs) {

super(args);

this.n = args.n;

this.inputSpec = [{ndim: 2}];

}

computeOutputShape(inputShape: Shape): Shape {

return [inputShape[0], this.n, inputShape[1]];

}

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

return tidy(() => {

inputs = getExactlyOneTensor(inputs);

return K.repeat(inputs, this.n);

});

}

getConfig(): serialization.ConfigDict {

const config = {

n: this.n,

};

const baseConfig = super.getConfig();

Object.assign(config, baseConfig);

return config;

}

}

serialization.registerClass(RepeatVector);

export class Reshape extends Layer {

/** @nocollapse */

static className = 'Reshape';

private targetShape: Shape;

constructor(args: ReshapeLayerArgs) {

super(args);

this.targetShape = args.targetShape;

// Make sure that all unknown dimensions are represented as `null`.

for (let i = 0; i < this.targetShape.length; ++i) {

if (this.isUnknown(this.targetShape[i])) {

this.targetShape[i] = null;

}

}

}

private isUnknown(dim: number): boolean {

return dim < 0 || dim == null;

}

/**

* Finds and replaces a missing dimension in output shape.

*

* This is a near direct port of the internal Numpy function

* `_fix_unknown_dimension` in `numpy/core/src/multiarray/shape.c`.

*

* @param inputShape: Original shape of array begin reshape.

* @param outputShape: Target shape of the array, with at most a single

* `null` or negative number, which indicates an underdetermined dimension

* that should be derived from `inputShape` and the known dimensions of

* `outputShape`.

* @returns: The output shape with `null` replaced with its computed value.

* @throws: ValueError: If `inputShape` and `outputShape` do not match.

*/

private fixUnknownDimension(inputShape: Shape, outputShape: Shape): Shape {

const errorMsg = 'Total size of new array must be unchanged.';

const finalShape = outputShape.slice();

let known = 1;

let unknown = null;

for (let i = 0; i < finalShape.length; ++i) {

const dim = finalShape[i];

if (this.isUnknown(dim)) {

if (unknown === null) {

unknown = i;

} else {

throw new ValueError('Can only specifiy one unknown dimension.');

}

} else {

known *= dim;

}

}

const originalSize = arrayProd(inputShape);

if (unknown !== null) {

if (known === 0 || originalSize % known !== 0) {

throw new ValueError(errorMsg);

}

finalShape[unknown] = originalSize / known;

} else if (originalSize !== known) {

throw new ValueError(errorMsg);

}

return finalShape;

}

computeOutputShape(inputShape: Shape): Shape {

let anyUnknownDims = false;

for (let i = 0; i < inputShape.length; ++i) {

if (this.isUnknown(inputShape[i])) {

anyUnknownDims = true;

break;

}

}

if (anyUnknownDims) {

return inputShape.slice(0, 1).concat(this.targetShape);

} else {

return inputShape.slice(0, 1).concat(

this.fixUnknownDimension(inputShape.slice(1), this.targetShape));

}

}

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

return tidy(() => {

this.invokeCallHook(inputs, kwargs);

const input = getExactlyOneTensor(inputs);

const inputShape = input.shape;

const outputShape = inputShape.slice(0, 1).concat(

this.fixUnknownDimension(inputShape.slice(1), this.targetShape));

return input.reshape(outputShape);

});

}

getConfig(): serialization.ConfigDict {

const config = {

targetShape: this.targetShape,

};

const baseConfig = super.getConfig();

Object.assign(config, baseConfig);

return config;

}

}

serialization.registerClass(Reshape);

export declare interface PermuteLayerArgs extends LayerArgs {

/**

* Array of integers. Permutation pattern. Does not include the

* sample (batch) dimension. Index starts at 1.

* For instance, `[2, 1]` permutes the first and second dimensions

* of the input.

*/

dims: number[];

}

export class Permute extends Layer {

/** @nocollapse */

static className = 'Permute';

readonly dims: number[];

private readonly dimsIncludingBatch: number[];

constructor(args: PermuteLayerArgs) {

super(args);

if (args.dims == null) {

throw new Error(

'Required configuration field `dims` is missing during Permute ' +

'constructor call.');

}

if (!Array.isArray(args.dims)) {

throw new Error(

'Permute constructor requires `dims` to be an Array, but received ' +

`${args.dims} instead.`);

}

// Check the validity of the permutation indices.

const expectedSortedIndices = range(1, args.dims.length + 1);

if (!util.arraysEqual(args.dims.slice().sort(), expectedSortedIndices)) {

throw new Error(

'Invalid permutation `dims`: ' + JSON.stringify(args.dims) +

' `dims` must contain consecutive integers starting from 1.');

}

this.dims = args.dims;

this.dimsIncludingBatch = [0].concat(this.dims);

this.inputSpec = [new InputSpec({ndim: this.dims.length + 1})];

}

computeOutputShape(inputShape: Shape|Shape[]): Shape|Shape[] {

inputShape = getExactlyOneShape(inputShape);

const outputShape = inputShape.slice();

this.dims.forEach((dim: number, i: number) => {

outputShape[i + 1] = (inputShape as Shape)[dim];

});

return outputShape;

}

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

return transpose(getExactlyOneTensor(inputs), this.dimsIncludingBatch);

}

getConfig(): serialization.ConfigDict {

const config = {

dims: this.dims,

};

const baseConfig = super.getConfig();

Object.assign(config, baseConfig);

return config;

}

}

serialization.registerClass(Permute);

export declare interface MaskingArgs extends LayerArgs {

/**

* Masking Value. Defaults to `0.0`.

*/

maskValue?: number;

}

export class Masking extends Layer {

/** @nocollapse */

static className = 'Masking';

maskValue: number;

constructor(args?: MaskingArgs) {

super(args == null ? {} : args);

this.supportsMasking = true;

if (args != null) {

this.maskValue = args.maskValue == null ? 0 : args.maskValue;

} else {

this.maskValue = 0;

}

}

computeOutputShape(inputShape: Shape|Shape[]): Shape|Shape[] {

return inputShape;

}

getConfig() {

const baseConfig = super.getConfig();

const config = {maskValue: this.maskValue};

Object.assign(config, baseConfig);

return config;

}

computeMask(inputs: Tensor|Tensor[], mask?: Tensor|Tensor[]): Tensor {

const input = getExactlyOneTensor(inputs);

const axis = -1;

return any(notEqual(input, this.maskValue), axis);

}

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

return tidy(() => {

this.invokeCallHook(inputs, kwargs);

const input = getExactlyOneTensor(inputs);

const axis = -1;

const keepDims = true;

const booleanMask = any(notEqual(input, this.maskValue), axis, keepDims);

const output = input.mul(booleanMask.asType(input.dtype));

return output;

});

}

}

serialization.registerClass(Masking);