catniff 0.8.10 → 0.8.12

package/dist/core.d.ts

@@ -206,6 +206,7 @@ export declare class Tensor {
     bmm(other: TensorValue | Tensor): Tensor;
     mv(other: TensorValue | Tensor): Tensor;
     matmul(other: TensorValue | Tensor): Tensor;
+    tensordot(other: TensorValue | Tensor, axes?: number | [number, number] | [number[], number[]]): Tensor;
     dropout(rate: number): Tensor;
     triu(diagonal?: number): Tensor;
     tril(diagonal?: number): Tensor;

package/dist/core.js

@@ -1688,23 +1688,22 @@ class Tensor {
         if (matACols !== matBRows)
             throw new Error("Invalid matrices shape for multiplication");
         const matCDtype = Tensor.getResultDtype(this.dtype, other.dtype);
-        const matCShape = [matARows, matBCols];
-        const matCStrides = Tensor.getStrides(matCShape);
-        const matCSize = Tensor.shapeToSize(matCShape);
+        const matCSize = matARows * matBCols;
         const matC = new dtype_1.TypedArray[matCDtype](matCSize).fill(0);
         for (let i = 0; i < matARows; i++) {
-            for (let j = 0; j < matBCols; j++) {
-                for (let k = 0; k < matACols; k++) {
-                    // Tensor values are 1D arrays so we have to get real index using strides
-                    matC[i * matCStrides[0] + j * matCStrides[1]] +=
-                        matA[i * matAStrides[0] + k * matAStrides[1] + this.offset] *
-                            matB[k * matBStrides[0] + j * matBStrides[1] + other.offset];
+            const aRowOffset = i * matAStrides[0] + this.offset;
+            const cRowOffset = i * matBCols;
+            for (let k = 0; k < matACols; k++) {
+                const aVal = matA[aRowOffset + k * matAStrides[1]];
+                const bRowOffset = k * matBStrides[0] + other.offset;
+                for (let j = 0; j < matBCols; j++) {
+                    matC[cRowOffset + j] += aVal * matB[bRowOffset + j * matBStrides[1]];
                 }
             }
         }
         const out = new Tensor(matC, {
-            shape: matCShape,
-            strides: matCStrides,
+            shape: [matARows, matBCols],
+            strides: [matBCols, 1],
             offset: 0,
             numel: matCSize,
             device: this.device,
@@ -1751,25 +1750,28 @@ class Tensor {
         if (batchACols !== batchBRows)
             throw new Error("Invalid matrices shape for multiplication");
         const batchCDtype = Tensor.getResultDtype(this.dtype, other.dtype);
-        const batchCShape = [batchSize, batchARows, batchBCols];
-        const batchCStrides = Tensor.getStrides(batchCShape);
-        const batchCSize = Tensor.shapeToSize(batchCShape);
+        const matrixSize = batchARows * batchBCols;
+        const batchCSize = batchSize * matrixSize;
         const batchC = new dtype_1.TypedArray[batchCDtype](batchCSize).fill(0);
         for (let q = 0; q < batchSize; q++) {
+            const aQOffset = q * batchAStrides[0] + this.offset;
+            const bQOffset = q * batchBStrides[0] + other.offset;
+            const cQOffset = q * matrixSize;
             for (let i = 0; i < batchARows; i++) {
-                for (let j = 0; j < batchBCols; j++) {
-                    for (let k = 0; k < batchACols; k++) {
-                        // Tensor values are 1D arrays so we have to get real index using strides
-                        batchC[q * batchCStrides[0] + i * batchCStrides[1] + j * batchCStrides[2]] +=
-                            batchA[q * batchAStrides[0] + i * batchAStrides[1] + k * batchAStrides[2] + this.offset] *
-                                batchB[q * batchBStrides[0] + k * batchBStrides[1] + j * batchBStrides[2] + other.offset];
+                const aRowOffset = aQOffset + i * batchAStrides[1];
+                const cRowOffset = cQOffset + i * batchBCols;
+                for (let k = 0; k < batchACols; k++) {
+                    const aVal = batchA[aRowOffset + k * batchAStrides[2]];
+                    const bRowOffset = bQOffset + k * batchBStrides[1];
+                    for (let j = 0; j < batchBCols; j++) {
+                        batchC[cRowOffset + j] += aVal * batchB[bRowOffset + j * batchBStrides[2]];
                     }
                 }
             }
         }
         const out = new Tensor(batchC, {
-            shape: batchCShape,
-            strides: batchCStrides,
+            shape: [batchSize, batchARows, batchBCols],
+            strides: [matrixSize, batchBCols, 1],
             offset: 0,
             numel: batchCSize,
             device: this.device,
@@ -1858,18 +1860,26 @@ class Tensor {
             const outputValue = new dtype_1.TypedArray[outputDtype](outputSize).fill(0);
             const outputOffsetStrides = outputStrides.slice(0, -2);
             // Loop through outer dims and do matmul on two outer-most dims
+            const outputRowStride = outputStrides[lastDim - 1];
+            const outputColStride = outputStrides[lastDim];
+            const selfRowStride = selfStrides[lastDim - 1];
+            const selfColStride = selfStrides[lastDim];
+            const otherRowStride = otherStrides[lastDim - 1];
+            const otherColStride = otherStrides[lastDim];
             for (let index = 0; index < offsetSize; index++) {
                 const coords = Tensor.indexToCoords(index, offsetStrides);
-                const offset = Tensor.coordsToIndex(coords, outputOffsetStrides);
-                const selfOffset = Tensor.coordsToUnbroadcastedIndex(coords, selfOffsetShape, selfOffsetStrides);
-                const otherOffset = Tensor.coordsToUnbroadcastedIndex(coords, otherOffsetShape, otherOffsetStrides);
+                const outBatchOffset = Tensor.coordsToIndex(coords, outputOffsetStrides);
+                const selfBatchOffset = Tensor.coordsToUnbroadcastedIndex(coords, selfOffsetShape, selfOffsetStrides) + this.offset;
+                const otherBatchOffset = Tensor.coordsToUnbroadcastedIndex(coords, otherOffsetShape, otherOffsetStrides) + other.offset;
                 for (let i = 0; i < batchARows; i++) {
-                    for (let j = 0; j < batchBCols; j++) {
-                        for (let k = 0; k < batchACols; k++) {
-                            const outputIdx = offset + i * outputStrides[lastDim - 1] + j * outputStrides[lastDim];
-                            const selfIdx = selfOffset + i * selfStrides[lastDim - 1] + k * selfStrides[lastDim];
-                            const otherIdx = otherOffset + k * otherStrides[lastDim - 1] + j * otherStrides[lastDim];
-                            outputValue[outputIdx] += batchA[selfIdx + this.offset] * batchB[otherIdx + other.offset];
+                    const selfRowOffset = selfBatchOffset + i * selfRowStride;
+                    const outRowOffset = outBatchOffset + i * outputRowStride;
+                    for (let k = 0; k < batchACols; k++) {
+                        const aVal = batchA[selfRowOffset + k * selfColStride];
+                        const otherRowOffset = otherBatchOffset + k * otherRowStride;
+                        for (let j = 0; j < batchBCols; j++) {
+                            outputValue[outRowOffset + j * outputColStride] +=
+                                aVal * batchB[otherRowOffset + j * otherColStride];
                         }
                     }
                 }
@@ -1906,6 +1916,80 @@ class Tensor {
         }
         throw new Error(`Shapes [${this.shape}] and [${other.shape}] are not supported`);
     }
+    // General tensor dot product
+    tensordot(other, axes = 2) {
+        other = this.handleOther(other);
+        let axesA, axesB;
+        // If axes is a number
+        if (typeof axes === "number") {
+            axesA = new Array(axes);
+            axesB = new Array(axes);
+            for (let i = 0; i < axes; i++) {
+                axesA[i] = this.shape.length - axes + i;
+                axesB[i] = i;
+            }
+        }
+        // If axes is a pair of numbers or a pair of axes
+        else {
+            // axes is [axesA, axesB]
+            [axesA, axesB] = axes;
+            // Convert single numbers to arrays
+            if (typeof axesA === "number")
+                axesA = [axesA];
+            if (typeof axesB === "number")
+                axesB = [axesB];
+        }
+        // Normalize axes
+        axesA = Tensor.normalizeDims(axesA, this.shape.length);
+        axesB = Tensor.normalizeDims(axesB, other.shape.length);
+        // Validate axes
+        if (axesA.length !== axesB.length) {
+            throw new Error("Number of axes to contract must be the same for both tensors");
+        }
+        // Validate dimensions match
+        for (let i = 0; i < axesA.length; i++) {
+            if (this.shape[axesA[i]] !== other.shape[axesB[i]]) {
+                throw new Error(`Dimension mismatch: a.shape[${axesA[i]}]=${this.shape[axesA[i]]} ` +
+                    `!= b.shape[${axesB[i]}]=${other.shape[axesB[i]]}`);
+            }
+        }
+        // Identify free (non-contracted) axes
+        const freeA = [];
+        for (let i = 0; i < this.shape.length; i++) {
+            if (!axesA.includes(i)) {
+                freeA.push(i);
+            }
+        }
+        const freeB = [];
+        for (let i = 0; i < other.shape.length; i++) {
+            if (!axesB.includes(i)) {
+                freeB.push(i);
+            }
+        }
+        // Permute a to move contracted axes to the end: [free_a, contract_a]
+        const aPermuted = this.permute([...freeA, ...axesA]);
+        // Permute b to move contracted axes to the beginning: [contract_b, free_b]
+        const bPermuted = other.permute([...axesB, ...freeB]);
+        // Reshape a to 2D matrix
+        // Shape: [product of free_a dims, product of contract_a dims]
+        const freeASize = freeA.reduce((prod, axis) => prod * this.shape[axis], 1);
+        const contractASize = axesA.reduce((prod, axis) => prod * this.shape[axis], 1);
+        const aReshaped = aPermuted.reshape([freeASize, contractASize]);
+        // Reshape b to 2D matrix
+        // Shape: [product of contract_b dims, product of free_b dims]
+        const contractBSize = axesB.reduce((prod, axis) => prod * other.shape[axis], 1);
+        const freeBSize = freeB.reduce((prod, axis) => prod * other.shape[axis], 1);
+        const bReshaped = bPermuted.reshape([contractBSize, freeBSize]);
+        // Use normal matmul, result shape: [freeASize, freeBSize]
+        const result2D = aReshaped.matmul(bReshaped);
+        // Reshape result back to proper n-dimensional form
+        // Final shape: [free_a_dims..., free_b_dims...]
+        const finalShape = [
+            ...freeA.map(i => this.shape[i]),
+            ...freeB.map(i => other.shape[i])
+        ];
+        return result2D.reshape(finalShape);
+    }
     // Dropout
     dropout(rate) {
         if (!Tensor.training || rate === 0)

package/package.json

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