Support for ConvLSTM2D

tfjs-layers/src/exports_layers.ts

@@ -14,6 +14,7 @@ import {input} from './exports';
 import {ELU, ELULayerArgs, LeakyReLU, LeakyReLULayerArgs, PReLU, PReLULayerArgs, ReLU, ReLULayerArgs, Softmax, SoftmaxLayerArgs, ThresholdedReLU, ThresholdedReLULayerArgs} from './layers/advanced_activations';
 import {Conv1D, Conv2D, Conv2DTranspose, Conv3D, ConvLayerArgs, Cropping2D, Cropping2DLayerArgs, SeparableConv2D, SeparableConvLayerArgs, UpSampling2D, UpSampling2DLayerArgs} from './layers/convolutional';
 import {DepthwiseConv2D, DepthwiseConv2DLayerArgs} from './layers/convolutional_depthwise';
+import {ConvLSTM2D, ConvLSTM2DArgs, ConvLSTM2DCell, ConvLSTM2DCellArgs} from './layers/convolutional_recurrent';
 import {Activation, ActivationLayerArgs, Dense, DenseLayerArgs, Dropout, DropoutLayerArgs, Flatten, FlattenLayerArgs, Masking, MaskingArgs, Permute, PermuteLayerArgs, RepeatVector, RepeatVectorLayerArgs, Reshape, ReshapeLayerArgs, SpatialDropout1D, SpatialDropout1DLayerConfig} from './layers/core';
 import {Embedding, EmbeddingLayerArgs} from './layers/embeddings';
 import {Add, Average, Concatenate, ConcatenateLayerArgs, Dot, DotLayerArgs, Maximum, Minimum, Multiply} from './layers/merge';
@@ -1352,6 +1353,74 @@ export function simpleRNNCell(args: SimpleRNNCellLayerArgs): RNNCell {
   return new SimpleRNNCell(args);
 }
 
+/**
+ * Convolutional LSTM layer - Xingjian Shi 2015.
+ *
+ * This is an `ConvRNN2D` layer consisting of one `ConvLSTM2DCell`. However,
+ * unlike the underlying `ConvLSTM2DCell`, the `apply` method of `ConvLSTM2D`
+ * operates on a sequence of inputs. The shape of the input (not including the
+ * first, batch dimension) needs to be 4-D, with the first dimension being time
+ * steps. For example:
+ *
+ * ```js
+ * const filters = 3;
+ * const kernelSize = 3;
+ *
+ * const batchSize = 4;
+ * const sequenceLength = 2;
+ * const size = 5;
+ * const channels = 3;
+ *
+ * const inputShape = [batchSize, sequenceLength, size, size, channels];
+ * const input = tf.ones(inputShape);
+ *
+ * const layer = tf.layers.convLstm2d({filters, kernelSize});
+ *
+ * const output = layer.apply(input);
+ * ```
+ */
+/** @doc {heading: 'Layers', subheading: 'Recurrent', namespace: 'layers'} */
+export function convLstm2d(args: ConvLSTM2DArgs): ConvLSTM2D {
+  return new ConvLSTM2D(args);
+}
+
+/**
+ * Cell class for `ConvLSTM2D`.
+ *
+ * `ConvLSTM2DCell` is distinct from the `ConvRNN2D` subclass `ConvLSTM2D` in
+ * that its `call` method takes the input data of only a single time step and
+ * returns the cell's output at the time step, while `ConvLSTM2D` takes the
+ * input data over a number of time steps. For example:
+ *
+ * ```js
+ * const filters = 3;
+ * const kernelSize = 3;
+ *
+ * const sequenceLength = 1;
+ * const size = 5;
+ * const channels = 3;
+ *
+ * const inputShape = [sequenceLength, size, size, channels];
+ * const input = tf.ones(inputShape);
+ *
+ * const cell = tf.layers.convLstm2dCell({filters, kernelSize});
+ *
+ * cell.build(input.shape);
+ *
+ * const outputSize = size - kernelSize + 1;
+ * const outShape = [sequenceLength, outputSize, outputSize, filters];
+ *
+ * const initialH = tf.zeros(outShape);
+ * const initialC = tf.zeros(outShape);
+ *
+ * const [o, h, c] = cell.call([input, initialH, initialC], {});
+ * ```
+ */
+/** @doc {heading: 'Layers', subheading: 'Recurrent', namespace: 'layers'} */
+export function convLstm2dCell(args: ConvLSTM2DCellArgs): ConvLSTM2DCell {
+  return new ConvLSTM2DCell(args);
+}
+
 /**
  * Base class for recurrent layers.
  *

tfjs-layers/src/initializers.ts

@@ -545,9 +545,8 @@ export class Orthogonal extends Initializer {
 
   apply(shape: Shape, dtype?: DataType): Tensor {
     return tidy(() => {
-      if (shape.length !== 2) {
-        throw new NotImplementedError(
-            'The Orthogonal Initializer does not support non-2D shapes yet.');
+      if (shape.length < 2) {
+        throw new NotImplementedError('Shape must be at least 2D.');
       }
       if (shape[0] * shape[1] > 2000) {
         console.warn(