modularise linalg ops

tfjs-core/scripts/touch_modular_op_files.ts

@@ -63,6 +63,11 @@ async function main() {
   let command = `touch ${filePath}`;
   execSync(command);
 
+  // create a test file
+  filePath = `./src/ops/${args.op}_test.ts`;
+  command = `touch ${filePath}`;
+  execSync(command);
+
   if (args.chained) {
     filePath = `./src/public/chained_ops/${args.op}.ts`;
     command = `touch ${filePath}`;

tfjs-core/src/ops/band_part.ts

@@ -0,0 +1,122 @@
+/**
+ * @license
+ * Copyright 2020 Google Inc. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {Tensor} from '../tensor';
+import {convertToTensor} from '../tensor_util_env';
+import {TensorLike} from '../types';
+import {assert} from '../util';
+
+import {reshape, stack, unstack} from './array_ops';
+import {greaterEqual} from './greater_equal';
+import {lessEqual} from './less_equal';
+import {logicalAnd, where} from './logical_ops';
+import {op} from './operation';
+import {sub} from './sub';
+import {range, scalar, zeros} from './tensor_ops';
+
+/**
+ * Copy a tensor setting everything outside a central band in each innermost
+ * matrix to zero.
+ *
+ * The band part is computed as follows: Assume input has `k` dimensions
+ * `[I, J, K, ..., M, N]`, then the output is a tensor with the same shape where
+ * `band[i, j, k, ..., m, n] = in_band(m, n) * input[i, j, k, ..., m, n]`.
+ * The indicator function
+ * `in_band(m, n) = (num_lower < 0 || (m-n) <= num_lower))`
+ * `&& (num_upper < 0 || (n-m) <= num_upper)`
+ *
+ * ```js
+ * const x = tf.tensor2d([[ 0,  1,  2, 3],
+ *                        [-1,  0,  1, 2],
+ *                        [-2, -1,  0, 1],
+ *                        [-3, -2, -1, 0]]);
+ * let y = tf.linalg.bandPart(x, 1, -1);
+ * y.print(); // [[ 0,  1,  2, 3],
+ *            //  [-1,  0,  1, 2],
+ *            //  [ 0, -1,  0, 1],
+ *            //  [ 0, 0 , -1, 0]]
+ * let z = tf.linalg.bandPart(x, 2, 1);
+ * z.print(); // [[ 0,  1,  0, 0],
+ *            //  [-1,  0,  1, 0],
+ *            //  [-2, -1,  0, 1],
+ *            //  [ 0, -2, -1, 0]]
+ * ```
+ *
+ * @param x Rank `k` tensor
+ * @param numLower Number of subdiagonals to keep.
+ *   If negative, keep entire lower triangle.
+ * @param numUpper Number of subdiagonals to keep.
+ *   If negative, keep entire upper triangle.
+ * @returns Rank `k` tensor of the same shape as input.
+ *   The extracted banded tensor.
+ */
+/**
+ * @doc {heading:'Operations', subheading:'Linear Algebra', namespace:'linalg'}
+ */
+function bandPart_<T extends Tensor>(
+    a: T|TensorLike, numLower: number, numUpper: number): T {
+  assert(
+      numLower % 1 === 0,
+      () => `bandPart(): numLower must be an integer, got ${numLower}.`);
+  assert(
+      numUpper % 1 === 0,
+      () => `bandPart(): numUpper must be an integer, got ${numUpper}.`);
+
+  const $a = convertToTensor(a, 'a', 'bandPart');
+
+  assert(
+      $a.rank >= 2,
+      () => `bandPart(): Rank must be at least 2, got ${$a.rank}.`);
+
+  const shape = $a.shape;
+  const [M, N] = $a.shape.slice(-2);
+
+  if (!(numLower <= M)) {
+    throw new Error(
+        `bandPart(): numLower (${numLower})` +
+        ` must not be greater than the number of rows (${M}).`);
+  }
+  if (!(numUpper <= N)) {
+    throw new Error(
+        `bandPart(): numUpper (${numUpper})` +
+        ` must not be greater than the number of columns (${N}).`);
+  }
+
+  if (numLower < 0) {
+    numLower = M;
+  }
+  if (numUpper < 0) {
+    numUpper = N;
+  }
+
+  const i = reshape(range(0, M, 1, 'int32'), [-1, 1]);
+  const j = range(0, N, 1, 'int32');
+  const ij = sub(i, j);
+
+  const inBand = logicalAnd(
+      lessEqual(ij, scalar(+numLower, 'int32')),
+      greaterEqual(ij, scalar(-numUpper, 'int32')));
+
+  const zero = zeros([M, N], $a.dtype);
+
+  return reshape(
+             stack(unstack(reshape($a, [-1, M, N]))
+                       .map(mat => where(inBand, mat, zero))),
+             shape) as T;
+}
+
+export const bandPart = op({bandPart_});

tfjs-core/src/ops/band_part_test.ts

@@ -0,0 +1,181 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import * as tf from '../index';
+import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
+import {Tensor2D, Tensor3D} from '../tensor';
+import {expectArraysClose} from '../test_util';
+
+import {scalar, tensor1d, tensor2d, tensor3d} from './ops';
+
+describeWithFlags('bandPart', ALL_ENVS, () => {
+  it('keeps tensor unchanged', async () => {
+    const x: Tensor2D = tensor2d([1, 1, 1, 1, 1, 1, 1, 1, 1], [3, 3]);
+    expectArraysClose(
+        await tf.linalg.bandPart(x, -1, -1).array(),
+        [[1, 1, 1], [1, 1, 1], [1, 1, 1]]);
+  });
+
+  it('upper triangular matrix', async () => {
+    const x: Tensor2D = tensor2d([1, 1, 1, 1, 1, 1, 1, 1, 1], [3, 3]);
+    expectArraysClose(
+        await tf.linalg.bandPart(x, 0, -1).array(),
+        [[1, 1, 1], [0, 1, 1], [0, 0, 1]]);
+  });
+
+  it('lower triangular matrix', async () => {
+    const x: Tensor2D = tensor2d([1, 1, 1, 1, 1, 1, 1, 1, 1], [3, 3]);
+    expectArraysClose(
+        await tf.linalg.bandPart(x, -1, 0).array(),
+        [[1, 0, 0], [1, 1, 0], [1, 1, 1]]);
+  });
+
+  it('diagonal elements', async () => {
+    const x: Tensor2D = tensor2d([1, 1, 1, 1, 1, 1, 1, 1, 1], [3, 3]);
+    expectArraysClose(
+        await tf.linalg.bandPart(x, 0, 0).array(),
+        [[1, 0, 0], [0, 1, 0], [0, 0, 1]]);
+  });
+
+  it('lower triangular elements', async () => {
+    const x: Tensor2D = tensor2d([1, 1, 1, 1, 1, 1, 1, 1, 1], [3, 3]);
+    expectArraysClose(
+        await tf.linalg.bandPart(x, 1, 0).array(),
+        [[1, 0, 0], [1, 1, 0], [0, 1, 1]]);
+  });
+
+  it('upper triangular elements', async () => {
+    const x: Tensor2D = tensor2d([1, 1, 1, 1, 1, 1, 1, 1, 1], [3, 3]);
+    expectArraysClose(
+        await tf.linalg.bandPart(x, 0, 1).array(),
+        [[1, 1, 0], [0, 1, 1], [0, 0, 1]]);
+  });
+
+  it('4X4 matrix - tensorflow python examples', async () => {
+    const x: Tensor2D = tensor2d(
+        [[0, 1, 2, 3], [-1, 0, 1, 2], [-2, -1, 0, 1], [-3, -2, -1, 0]]);
+    expectArraysClose(
+        await tf.linalg.bandPart(x, 1, -1).array(),
+        [[0, 1, 2, 3], [-1, 0, 1, 2], [0, -1, 0, 1], [0, 0, -1, 0]]);
+    expectArraysClose(
+        await tf.linalg.bandPart(x, 2, 1).array(),
+        [[0, 1, 0, 0], [-1, 0, 1, 0], [-2, -1, 0, 1], [0, -2, -1, 0]]);
+  });
+
+  it('3 dimensional matrix', async () => {
+    const x: Tensor3D = tensor3d([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]);
+    expectArraysClose(
+        await tf.linalg.bandPart(x, 0, 0).array(),
+        [[[1, 0], [0, 1]], [[1, 0], [0, 1]]]);
+  });
+
+  it('2X3X3 tensor', async () => {
+    const x: Tensor3D = tensor3d(
+        [[[1, 1, 1], [1, 1, 1], [1, 1, 1]], [[1, 1, 1], [1, 1, 1], [1, 1, 1]]]);
+    expectArraysClose(
+        await tf.linalg.bandPart(x, 1, 2).array(),
+        [[[1, 1, 1], [1, 1, 1], [0, 1, 1]], [[1, 1, 1], [1, 1, 1], [0, 1, 1]]]);
+  });
+
+  const la = tf.linalg;
+
+  it('fails for scalar', async () => {
+    const x = scalar(1);
+    expect(() => la.bandPart(x, 1, 2)).toThrowError(/bandPart.*rank/i);
+  });
+
+  it('fails for 1D tensor', async () => {
+    const x = tensor1d([1, 2, 3, 4, 5]);
+    expect(() => la.bandPart(x, 1, 2)).toThrowError(/bandPart.*rank/i);
+  });
+
+  it('fails if numLower or numUpper too large', async () => {
+    const a = tf.tensor2d([[1, 2, 3], [4, 5, 6]]);
+
+    for (const numLower of [3, 5, 8, 13]) {
+      for (const numUpper of [-1, 0, 1, 2]) {
+        expect(() => tf.linalg.bandPart(a, numLower, numUpper))
+            .toThrowError(/bandPart.*numLower/i);
+      }
+    }
+
+    for (const numLower of [-1, 0, 1]) {
+      for (const numUpper of [4, 5, 9]) {
+        expect(() => tf.linalg.bandPart(a, numLower, numUpper))
+            .toThrowError(/bandPart.*numUpper/i);
+      }
+    }
+
+    for (const numLower of [3, 5, 8, 13]) {
+      for (const numUpper of [4, 5, 9]) {
+        expect(() => tf.linalg.bandPart(a, numLower, numUpper))
+            .toThrowError(/bandPart.*(numLower|numUpper)/i);
+      }
+    }
+  });
+
+  it('works for 3x4 example', async () => {
+    const a = tf.tensor2d([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]);
+
+    expectArraysClose(
+        await la.bandPart(a, 0, 0).array(),
+        [[1, 0, 0, 0], [0, 6, 0, 0], [0, 0, 11, 0]]);
+    expectArraysClose(
+        await la.bandPart(a, 0, 1).array(),
+        [[1, 2, 0, 0], [0, 6, 7, 0], [0, 0, 11, 12]]);
+    expectArraysClose(
+        await la.bandPart(a, 0, 2).array(),
+        [[1, 2, 3, 0], [0, 6, 7, 8], [0, 0, 11, 12]]);
+    for (const numUpper of [3, 4, -1, -2]) {
+      expectArraysClose(
+          await la.bandPart(a, 0, numUpper).array(),
+          [[1, 2, 3, 4], [0, 6, 7, 8], [0, 0, 11, 12]]);
+    }
+
+    expectArraysClose(
+        await la.bandPart(a, 1, 0).array(),
+        [[1, 0, 0, 0], [5, 6, 0, 0], [0, 10, 11, 0]]);
+    expectArraysClose(
+        await la.bandPart(a, 1, 1).array(),
+        [[1, 2, 0, 0], [5, 6, 7, 0], [0, 10, 11, 12]]);
+    expectArraysClose(
+        await la.bandPart(a, 1, 2).array(),
+        [[1, 2, 3, 0], [5, 6, 7, 8], [0, 10, 11, 12]]);
+    for (const numUpper of [3, 4, -1, -2]) {
+      expectArraysClose(
+          await la.bandPart(a, 1, numUpper).array(),
+          [[1, 2, 3, 4], [5, 6, 7, 8], [0, 10, 11, 12]]);
+    }
+
+    for (const numLower of [2, 3, -1, -2]) {
+      expectArraysClose(
+          await la.bandPart(a, numLower, 0).array(),
+          [[1, 0, 0, 0], [5, 6, 0, 0], [9, 10, 11, 0]]);
+      expectArraysClose(
+          await la.bandPart(a, numLower, 1).array(),
+          [[1, 2, 0, 0], [5, 6, 7, 0], [9, 10, 11, 12]]);
+      expectArraysClose(
+          await la.bandPart(a, numLower, 2).array(),
+          [[1, 2, 3, 0], [5, 6, 7, 8], [9, 10, 11, 12]]);
+      for (const numUpper of [3, 4, -1, -2]) {
+        expectArraysClose(
+            await la.bandPart(a, numLower, numUpper).array(),
+            [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]);
+      }
+    }
+  });
+});