catniff 0.8.17 → 0.8.19

package/dist/core.d.ts

@@ -44,7 +44,7 @@ export declare class Tensor {
     static coordsToIndex(coords: number[], strides: number[]): number;
     static shapeToSize(shape: number[]): number;
     static getResultDtype(type1: dtype, type2: dtype): dtype;
-    handleOther(other: Tensor | TensorValue): Tensor;
+    handleOther(other: Tensor | TensorValue, forceSameDevice?: boolean): Tensor;
     static elementWiseAB(tA: Tensor, tB: Tensor, op: (tA: number, tB: number) => number): Tensor;
     static elementWiseSelf(tA: Tensor, op: (tA: number) => number): Tensor;
     elementWiseABDAG(other: TensorValue | Tensor, op: (a: number, b: number) => number, thisGrad?: (self: Tensor, other: Tensor, outGrad: Tensor) => Tensor, otherGrad?: (self: Tensor, other: Tensor, outGrad: Tensor) => Tensor): Tensor;
@@ -67,6 +67,7 @@ export declare class Tensor {
     chunk(chunks: number, dim?: number): Tensor[];
     expand(newShape: number[]): Tensor;
     cat(other: Tensor | TensorValue, dim?: number): Tensor;
+    stack(others: (Tensor | TensorValue)[], dim?: number): Tensor;
     squeeze(dims?: number[] | number): Tensor;
     unsqueeze(dim: number): Tensor;
     sort(dim?: number, descending?: boolean): Tensor;
@@ -88,6 +89,12 @@ export declare class Tensor {
             outIndex: number;
         }) => number;
     }): Tensor;
+    static reduceArg(tensor: Tensor, dim: number, keepDim: boolean, config: {
+        identity: number;
+        isBetter: (accumulator: number, value: number) => boolean;
+    }): Tensor;
+    argmax(dim: number, keepDim?: boolean): Tensor;
+    argmin(dim: number, keepDim?: boolean): Tensor;
     sum(dims?: number[] | number, keepDims?: boolean): Tensor;
     prod(dims?: number[] | number, keepDims?: boolean): Tensor;
     mean(dims?: number[] | number, keepDims?: boolean): Tensor;

package/dist/core.js

@@ -165,10 +165,10 @@ class Tensor {
         }
         return type2;
     }
-    // Utility to handle other tensor if an op needs a second operand
-    handleOther(other) {
+    // Utility to handle other tensor if an op needs other operands
+    handleOther(other, forceSameDevice = true) {
         if (other instanceof Tensor) {
-            if (this.device !== other.device) {
+            if (forceSameDevice && this.device !== other.device) {
                 throw new Error("Can not operate on tensors that are not on the same device");
             }
             return other;
@@ -602,7 +602,7 @@ class Tensor {
     }
     // Tensor indexing
     index(indices) {
-        const tensorIndices = this.handleOther(indices).clone();
+        const tensorIndices = this.handleOther(indices, false).clone();
         if (tensorIndices.shape.length === 0) {
             return this.indexWithArray([tensorIndices.value[0]]).squeeze(0);
         }
@@ -843,6 +843,15 @@ class Tensor {
         }
         return out;
     }
+    // Tensor stack
+    stack(others, dim = 0) {
+        let out = this.unsqueeze(dim);
+        for (let index = 0; index < others.length; index++) {
+            const other = this.handleOther(others[index]).unsqueeze(dim);
+            out = out.cat(other, dim);
+        }
+        return out;
+    }
     // Tensor squeeze
     squeeze(dims) {
         if (this.shape.length === 0)
@@ -1146,6 +1155,67 @@ class Tensor {
         }
         return keepDims ? out : out.squeeze(dims);
     }
+    // Generic arg reduction operation handler
+    static reduceArg(tensor, dim, keepDim, config) {
+        if (tensor.shape.length === 0)
+            return tensor;
+        // Handle negative indexing
+        if (dim < 0) {
+            dim += tensor.shape.length;
+        }
+        // If dimension out of bound, throw error
+        if (dim >= tensor.shape.length || dim < 0) {
+            throw new Error("Dimension does not exist to apply arg reduction");
+        }
+        const dimSize = tensor.shape[dim];
+        const outputShape = tensor.shape.map((d, i) => dim === i ? 1 : d);
+        const outputStrides = Tensor.getStrides(outputShape);
+        const outputSize = tensor.numel / dimSize;
+        const bestValues = new dtype_1.TypedArray[tensor.dtype](outputSize).fill(config.identity);
+        const bestIndices = new Int32Array(outputSize).fill(0);
+        const linearStrides = Tensor.getStrides(tensor.shape);
+        // Forward pass
+        for (let flatIndex = 0; flatIndex < tensor.numel; flatIndex++) {
+            // Convert linear index to coordinates using contiguous strides
+            const coords = Tensor.indexToCoords(flatIndex, linearStrides);
+            // Coordinate in current dim
+            const dimCoord = coords[dim];
+            // Convert coordinates to actual strided index
+            const realFlatIndex = Tensor.coordsToIndex(coords, tensor.strides) + tensor.offset;
+            // Convert coords to reduced index
+            coords[dim] = 0;
+            const outFlatIndex = Tensor.coordsToIndex(coords, outputStrides);
+            // Check if current value is better to swap
+            const val = tensor.value[realFlatIndex];
+            if (config.isBetter(val, bestValues[outFlatIndex])) {
+                bestValues[outFlatIndex] = val;
+                bestIndices[outFlatIndex] = dimCoord;
+            }
+        }
+        const out = new Tensor(bestIndices, {
+            shape: outputShape,
+            strides: outputStrides,
+            offset: 0,
+            numel: outputSize,
+            device: tensor.device,
+            dtype: "int32"
+        });
+        return keepDim ? out : out.squeeze(dim);
+    }
+    // Tensor argmax
+    argmax(dim, keepDim = false) {
+        return Tensor.reduceArg(this, dim, keepDim, {
+            identity: -Infinity,
+            isBetter: (a, b) => a > b
+        });
+    }
+    // Tensor argmin
+    argmin(dim, keepDim = false) {
+        return Tensor.reduceArg(this, dim, keepDim, {
+            identity: Infinity,
+            isBetter: (a, b) => a < b
+        });
+    }
     // Simplified reduction operations
     sum(dims, keepDims = false) {
         return Tensor.reduce(this, dims, keepDims, {
@@ -2543,11 +2613,15 @@ class Tensor {
     // Returns the nicely Pytorch-like formatted string form
     toString() {
         const val = this.val();
-        // Format a single number (integers get trailing dot)
+        // Format a single number
         const formatNum = (n) => {
-            if (Number.isInteger(n) && Math.abs(n) < 1e8) {
+            // For ints with int dtype
+            if (this.dtype.includes("int"))
+                return n.toFixed(0);
+            // For ints with float dtype
+            if (Number.isInteger(n) && Math.abs(n) < 1e8)
                 return n.toFixed(0) + ".";
-            }
+            // For floats
             return n.toString();
         };
         // Handle scalar

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "catniff",
-  "version": "0.8.17",
+  "version": "0.8.19",
   "description": "Torch-like deep learning framework for Javascript",
   "main": "./dist/index.js",
   "scripts": {