Merge pull request #64 from BruceDai/add_triangular

Implement triangular

src/lib/validate-input.js

@@ -444,3 +444,13 @@ export function validateGatherParams(input, indices, {axis = 0} = {}) {
     }
   }
 }
+
+export function validateTriangularParams(input, {diagonal = 0} = {}) {
+  const inputRank = input.rank;
+  if (inputRank < 2) {
+    throw new Error('The input should be at least a 2-D tensor.');
+  }
+  if (!Number.isInteger(diagonal)) {
+    throw new Error(`The diagonal should be an integer.`);
+  }
+}

src/triangular.js

@@ -0,0 +1,39 @@
+'use strict';
+
+import {Tensor, sizeOfShape} from './lib/tensor.js';
+import {validateTriangularParams} from './lib/validate-input.js';
+
+/**
+ * Get retained boolean flag.
+ * @param {Array} location
+ * @param {Boolean} upper
+ * @param {Number} diagonal
+ * @return {Boolean}
+ */
+function isRetainedValue(location, upper, diagonal) {
+  const i = location[location.length - 2];
+  const j = location[location.length - 1];
+  return upper ? j >= i + diagonal : j <= i + diagonal;
+}
+
+/**
+ * Given a 2-D tensor (matrix), return a 2-D tensor containing either the upper or lower triangular
+ * part of the input tensor.
+ * @param {Tensor} input
+ * @param {MLTriangularOptions} [options]
+ * @return {Tensor}
+ */
+export function triangular(input, {upper = true, diagonal = 0} = {}) {
+  validateTriangularParams(...arguments);
+  const shapeOutput = input.shape.slice();
+  const output = new Tensor(shapeOutput);
+
+  for (let outputIndex = 0; outputIndex < sizeOfShape(shapeOutput); ++outputIndex) {
+    const outputLoc = output.locationFromIndex(outputIndex);
+    const retainedFlag = isRetainedValue(outputLoc, upper, diagonal);
+    const inputValue = retainedFlag ? input.getValueByLocation(outputLoc) : 0;
+    output.setValueByLocation(outputLoc, inputValue);
+  }
+
+  return output;
+}