Support MatrixBandPart and allow numLower to be a tensor

tfjs-converter/docs/supported_ops.md

@@ -248,6 +248,7 @@
 |Einsum|Einsum|
 |MatMul|matMul|
 |Transpose|transpose|
+|MatrixBandPart|MatrixBandPart|
 |Not mapped|dot|
 |Not mapped|norm|
 |Not mapped|outerProduct|

tfjs-converter/python/tensorflowjs/op_list/matrices.json

@@ -225,5 +225,26 @@
         "type": "dtype"
       }
     ]
+  },
+  {
+    "tfOpName": "MatrixBandPart",
+    "category": "matrices",
+    "inputs": [
+      {
+        "start": 0,
+        "name": "a",
+        "type": "tensor"
+      },
+      {
+        "start": 1,
+        "name": "numLower",
+        "type": "tensor"
+      },
+      {
+        "start": 1,
+        "name": "numUpper",
+        "type": "tensor"
+      }
+    ]
   }
 ]

tfjs-converter/src/operations/executors/matrices_executor.ts

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {Tensor, Tensor2D} from '@tensorflow/tfjs-core';
+import {Scalar, Tensor, Tensor2D} from '@tensorflow/tfjs-core';
 // tslint:disable-next-line: no-imports-from-dist
 import * as tfOps from '@tensorflow/tfjs-core/dist/ops/ops_for_converter';
 
@@ -26,8 +26,8 @@ import {InternalOpExecutor, Node} from '../types';
 import {getParamValue} from './utils';
 
 export const executeOp: InternalOpExecutor =
-    (node: Node, tensorMap: NamedTensorsMap,
-     context: ExecutionContext, ops = tfOps): Tensor[] => {
+    (node: Node, tensorMap: NamedTensorsMap, context: ExecutionContext,
+     ops = tfOps): Tensor[] => {
       switch (node.op) {
         case 'BatchMatMul':
         case 'BatchMatMulV2':
@@ -89,6 +89,12 @@ export const executeOp: InternalOpExecutor =
             leakyreluAlpha
           })];
 
+        case 'MatrixBandPart':
+          return [ops.linalg.bandPart(
+              getParamValue('a', node, tensorMap, context) as Tensor2D,
+              getParamValue('numLower', node, tensorMap, context) as Scalar,
+              getParamValue('numUpper', node, tensorMap, context) as Scalar)];
+
         default:
           throw TypeError(`Node type ${node.op} is not implemented`);
       }

tfjs-converter/src/operations/executors/matrices_executor_test.ts

@@ -18,9 +18,9 @@
 import {Tensor} from '@tensorflow/tfjs-core';
 // tslint:disable-next-line: no-imports-from-dist
 import * as tfOps from '@tensorflow/tfjs-core/dist/ops/ops_for_converter';
-import * as matrices from '../op_list/matrices';
 
 import {ExecutionContext} from '../../executor/execution_context';
+import * as matrices from '../op_list/matrices';
 import {Node} from '../types';
 
 import {executeOp} from './matrices_executor';
@@ -141,7 +141,7 @@ describe('matrices', () => {
         node.op = '_FusedMatMul';
 
         node.attrParams['fusedOps'] =
-          createStrArrayAttr(['biasadd', 'leakyrelu']);
+            createStrArrayAttr(['biasadd', 'leakyrelu']);
         node.attrParams['numArgs'] = createNumberAttr(1);
         node.attrParams.transposeA = createBoolAttr(true);
         node.attrParams.transposeB = createBoolAttr(false);
@@ -199,5 +199,20 @@ describe('matrices', () => {
         expect(spyOps.transpose).toHaveBeenCalledWith(input1[0], [1, 2]);
       });
     });
+    describe('MatrixBandPart', () => {
+      it('should call tfOps.linalg.bandPart', () => {
+        node.op = 'MatrixBandPart';
+        node.inputNames = ['input1', 'input2', 'input3'];
+        node.inputParams.a = createTensorAttr(0);
+        node.inputParams.numLower = createTensorAttr(1);
+        node.inputParams.numUpper = createTensorAttr(2);
+        const input3 = [tfOps.scalar(3.0)];
+        spyOps.linalg.bandPart.and.returnValue({});
+        executeOp(node, {input1, input2, input3}, context, spyOpsAsTfOps);
+
+        expect(spyOps.linalg.bandPart)
+            .toHaveBeenCalledWith(input1[0], input2[0], input3[0]);
+      });
+    });
   });
 });

tfjs-core/src/kernel_names.ts

@@ -489,6 +489,11 @@ export interface LRNGradAttrs {
   beta: number;
 }
 
+export const MatrixBandPart = 'MatrixBandPart';
+export type MatrixBandPartInputs =
+    Pick<NamedTensorInfoMap, 'input'|'numLower'|'numUpper'>;
+export interface MatrixBandPartAttrs {}
+
 export const Max = 'Max';
 export type MaxInputs = Pick<NamedTensorInfoMap, 'x'>;
 export interface MaxAttrs {

tfjs-core/src/ops/linalg/band_part.ts

@@ -15,18 +15,20 @@
  * =============================================================================
  */
 
-import {Tensor} from '../../tensor';
+import {Scalar, Tensor} from '../../tensor';
 import {convertToTensor} from '../../tensor_util_env';
 import {TensorLike} from '../../types';
 import {assert} from '../../util';
 
 import {greaterEqual} from '../greater_equal';
+import {less} from '../less';
 import {lessEqual} from '../less_equal';
 import {logicalAnd} from '../logical_and';
+import {minimum} from '../minimum';
+import {neg} from '../neg';
 import {op} from '../operation';
 import {range} from '../range';
 import {reshape} from '../reshape';
-import {scalar} from '../scalar';
 import {stack} from '../stack';
 import {sub} from '../sub';
 import {unstack} from '../unstack';
@@ -72,48 +74,61 @@ import {zeros} from '../zeros';
  * @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}.`);
-
+    a: T|TensorLike, numLower: number|Scalar, numUpper: number|Scalar): T {
   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}).`);
+  let $numLower: Scalar;
+  let $numUpper: Scalar;
+  if (typeof numLower === 'number') {
+    assert(
+        numLower % 1 === 0,
+        () => `bandPart(): numLower must be an integer, got ${numLower}.`);
+    assert(
+        numLower <= M,
+        () => `bandPart(): numLower (${numLower})` +
+            ` must not be greater than the number of rows (${M}).`);
+    $numLower =
+        convertToTensor(numLower < 0 ? M : numLower, 'numLower', 'bandPart') as
+        Scalar;
+  } else {
+    assert(
+        numLower.dtype === 'int32',
+        () => `bandPart(): numLower's dtype must be an int32.`);
+    // If numLower is a Scalar, checking `numLower <= M` could hurt performance,
+    // but minimum(numLower, M) could avoid unexpected results.
+    $numLower = where(less(numLower, 0), M, minimum(numLower, M)) as Scalar;
   }
 
-  if (numLower < 0) {
-    numLower = M;
-  }
-  if (numUpper < 0) {
-    numUpper = N;
+  if (typeof numUpper === 'number') {
+    assert(
+        numUpper % 1 === 0,
+        () => `bandPart(): numUpper must be an integer, got ${numUpper}.`);
+    assert(
+        numUpper <= N,
+        () => `bandPart(): numUpper (${numUpper})` +
+            ` must not be greater than the number of columns (${N}).`);
+    $numUpper =
+        convertToTensor(numUpper < 0 ? N : numUpper, 'numUpper', 'bandPart') as
+        Scalar;
+  } else {
+    assert(
+        numUpper.dtype === 'int32',
+        () => `bandPart(): numUpper's dtype must be an int32.`);
+    $numUpper = where(less(numUpper, 0), N, minimum(numUpper, N)) as Scalar;
   }
 
   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 inBand =
+      logicalAnd(lessEqual(ij, $numLower), greaterEqual(ij, neg($numUpper)));
 
   const zero = zeros([M, N], $a.dtype);
 

tfjs-core/src/ops/linalg/band_part_test.ts

@@ -178,4 +178,13 @@ describeWithFlags('bandPart', ALL_ENVS, () => {
       }
     }
   });
+
+  it('works for tensor numLower and tensor numUpper', async () => {
+    const x: Tensor2D = tensor2d([1, 1, 1, 1, 1, 1, 1, 1, 1], [3, 3]);
+    expectArraysClose(
+        await tf.linalg
+            .bandPart(x, tf.scalar(-1, 'int32'), tf.scalar(0, 'int32'))
+            .array(),
+        [[1, 0, 0], [1, 1, 0], [1, 1, 1]]);
+  });
 });