catniff 0.8.22 → 0.9.0

package/dist/core.d.ts

@@ -70,6 +70,7 @@ export declare class Tensor {
     chunk(chunks: number, dim?: number): Tensor[];
     expand(newShape: number[]): Tensor;
     unfold(dim: number, size: number, step: number): Tensor;
+    pad(pad: number[], mode?: string, value?: number): Tensor;
     cat(other: Tensor | TensorValue, dim?: number): Tensor;
     stack(others: (Tensor | TensorValue)[], dim?: number): Tensor;
     squeeze(dims?: number[] | number): Tensor;
@@ -218,6 +219,7 @@ export declare class Tensor {
     mv(other: TensorValue | Tensor): Tensor;
     matmul(other: TensorValue | Tensor): Tensor;
     tensordot(other: TensorValue | Tensor, axes?: number | [number, number] | [number[], number[]]): Tensor;
+    conv2d(weight: Tensor | TensorValue, bias?: Tensor | TensorValue, stride?: number | [number, number], padding?: number | [number, number], dilation?: number | [number, number], groups?: number): Tensor;
     dropout(rate: number): Tensor;
     triu(diagonal?: number): Tensor;
     tril(diagonal?: number): Tensor;

package/dist/core.js

@@ -328,8 +328,14 @@ class Tensor {
         }
         const reducedGrad = accumGrad.sum(axesToReduce, true);
         const squeezedGrad = reducedGrad.squeeze(axesToSqueeze);
+        // Enforce 0-offset contiguous grads and correct dtype
         if (typeof tensor.grad === "undefined") {
-            tensor.grad = squeezedGrad;
+            let grad = squeezedGrad;
+            // Handle potentially contiguous tensors with non zero offset
+            if (grad.offset !== 0) {
+                grad = grad.clone();
+            }
+            tensor.grad = grad.contiguous().cast(tensor.dtype);
         }
         else {
             tensor.grad = tensor.grad.add(squeezedGrad.cast(tensor.dtype));
@@ -795,12 +801,12 @@ class Tensor {
                 const outGrad = out.grad;
                 const grad = Tensor.zerosLike(this);
                 for (let i = 0; i < out.numel; i++) {
-                    const coords = Tensor.indexToCoords(i, newStrides);
+                    const coords = Tensor.indexToCoords(i, Tensor.getStrides(outGrad.shape));
                     const windowIdx = coords[dim];
                     const withinWindow = coords[coords.length - 1];
                     coords[dim] = windowIdx * step + withinWindow;
                     coords.pop();
-                    const sourceIdx = Tensor.coordsToIndex(coords, this.strides);
+                    const sourceIdx = Tensor.coordsToIndex(coords, Tensor.getStrides(grad.shape));
                     grad.value[sourceIdx] += outGrad.value[i];
                 }
                 Tensor.addGrad(this, grad);
@@ -808,6 +814,70 @@ class Tensor {
         }
         return out;
     }
+    // Tensor padding
+    pad(pad, mode = "constant", value = 0) {
+        const original = this.clone().contiguous(); // This is needed for index padding to work
+        const outputShape = [...original.shape];
+        const paddingPerDim = [];
+        for (let i = 0; i < original.shape.length; i++) {
+            const left = pad[(original.shape.length - 1 - i) * 2] || 0;
+            const right = pad[(original.shape.length - 1 - i) * 2 + 1] || 0;
+            paddingPerDim[i] = { left, right };
+            outputShape[i] += left + right;
+        }
+        const outputSize = Tensor.shapeToSize(outputShape);
+        if (mode === "constant") {
+            const outputValue = new dtype_1.TypedArray[original.dtype](outputSize).fill(value);
+            const outputStrides = Tensor.getStrides(outputShape);
+            for (let index = 0; index < original.numel; index++) {
+                const coords = Tensor.indexToCoords(index, original.strides);
+                let paddedIndex = 0;
+                // Pad each coord
+                for (let j = 0; j < original.shape.length; j++) {
+                    const shiftedCoord = coords[j] + paddingPerDim[j].left;
+                    paddedIndex += shiftedCoord * outputStrides[j];
+                }
+                outputValue[paddedIndex] = original.value[index];
+            }
+            const out = new Tensor(outputValue, {
+                shape: outputShape,
+                strides: outputStrides,
+                offset: 0,
+                dtype: original.dtype,
+                device: original.device
+            });
+            if (original.requiresGrad) {
+                out.requiresGrad = true;
+                out.children.push(original);
+                out.gradFn = () => {
+                    const outGrad = out.grad;
+                    const gradValue = new dtype_1.TypedArray[original.dtype](original.numel);
+                    const gradStrides = Tensor.getStrides(original.shape);
+                    for (let index = 0; index < gradValue.length; index++) {
+                        const coords = Tensor.indexToCoords(index, gradStrides);
+                        let paddedIndex = 0;
+                        // Pad each coord
+                        for (let j = 0; j < original.shape.length; j++) {
+                            const shiftedCoord = coords[j] + paddingPerDim[j].left;
+                            paddedIndex += shiftedCoord * outputStrides[j];
+                        }
+                        gradValue[index] = outGrad.value[paddedIndex];
+                    }
+                    Tensor.addGrad(original, new Tensor(gradValue, {
+                        shape: original.shape,
+                        strides: gradStrides,
+                        offset: 0,
+                        dtype: original.dtype,
+                        device: original.device
+                    }));
+                };
+            }
+            return out;
+        }
+        else {
+            throw new Error(`Padding mode not supported: "${mode}"`);
+        }
+    }
     // Tensor concatentation
     cat(other, dim = 0) {
         other = this.handleOther(other);
@@ -2113,6 +2183,74 @@ class Tensor {
         ];
         return result2D.reshape(finalShape);
     }
+    // 2D convolution
+    conv2d(weight, bias, stride = 1, padding = 0, dilation = 1, groups = 1) {
+        weight = this.handleOther(weight);
+        const [sH, sW] = Array.isArray(stride) ? stride : [stride, stride];
+        const [pH, pW] = Array.isArray(padding) ? padding : [padding, padding];
+        const [dH, dW] = Array.isArray(dilation) ? dilation : [dilation, dilation];
+        const [N, Cin, H, W] = this.shape;
+        const [Cout, CinPerGroup, kH, kW] = weight.shape;
+        // Pad input
+        let x = (pH > 0 || pW > 0) ? this.pad([pW, pW, pH, pH]) : this;
+        const Hp = H + 2 * pH;
+        const Wp = W + 2 * pW;
+        const Hout = Math.floor((Hp - dH * (kH - 1) - 1) / sH + 1);
+        const Wout = Math.floor((Wp - dW * (kW - 1) - 1) / sW + 1);
+        // Unfold H with a window large enough to cover the dilated kernel extent,
+        // then slice every dH-th position to realise the dilation holes.
+        // x: [N, Cin, Hp, Wp]
+        //   -> unfold(2, dH*(kH-1)+1, sH)
+        //   -> [N, Cin, Hout, Wp, dH*(kH-1)+1]
+        //   -> slice step dH on last dim
+        //   -> [N, Cin, Hout, Wp, kH]
+        const dilKH = dH * (kH - 1) + 1;
+        x = x.unfold(2, dilKH, sH);
+        if (dH > 1)
+            x = x.slice([[0, N], [0, Cin], [0, Hout], [0, Wp], [0, dilKH, dH]]);
+        // Unfold W
+        // x: [N, Cin, Hout, Wp, kH]
+        //   -> unfold(3, dW*(kW-1)+1, sW)
+        //   -> [N, Cin, Hout, Wout, kH, dW*(kW-1)+1]
+        //   -> slice step dW on last dim
+        //   -> [N, Cin, Hout, Wout, kH, kW]
+        const dilKW = dW * (kW - 1) + 1;
+        x = x.unfold(3, dilKW, sW);
+        if (dW > 1)
+            x = x.slice([[0, N], [0, Cin], [0, Hout], [0, Wout], [0, kH], [0, dilKW, dW]]);
+        // Reshape patches to [N, Hout*Wout, Cin*kH*kW]
+        // permute [0,2,3,1,4,5] -> [N, Hout, Wout, Cin, kH, kW]
+        // then reshape merges the spatial and channel-kernel dims.
+        // reshape() forces contiguity internally so no explicit .contiguous() needed.
+        x = x.permute([0, 2, 3, 1, 4, 5]).reshape([N, Hout * Wout, Cin * kH * kW]);
+        // Matmul with weight
+        // weight: [Cout, CinPerGroup, kH, kW] -> [Cout, CinPerGroup*kH*kW]
+        const w = weight.reshape([Cout, CinPerGroup * kH * kW]);
+        let out;
+        if (groups === 1) {
+            // x: [N, Hout*Wout, Cin*kH*kW]  @  w.t(): [Cin*kH*kW, Cout]
+            // -> [N, Hout*Wout, Cout]
+            out = x.matmul(w.t());
+        }
+        else {
+            // Each group handles Cin/groups input channels and Cout/groups output channels.
+            // chunk(groups, 2) splits the Cin*kH*kW axis into groups equal slices,
+            // each of size CinPerGroup*kH*kW — valid because reshape laid Cin outermost.
+            const patchChunks = x.chunk(groups, 2); // Tensor[groups], each [N, Hout*Wout, CinPerGroup*kH*kW]
+            const weightChunks = w.chunk(groups, 0); // Tensor[groups], each [Cout/groups, CinPerGroup*kH*kW]
+            const groupOuts = patchChunks.map((patch, i) => patch.matmul(weightChunks[i].t()) // [N, Hout*Wout, Cout/groups]
+            );
+            // Cat all group outputs along the channel axis (dim 2)
+            out = groupOuts.reduce((acc, g) => acc.cat(g, 2)); // [N, Hout*Wout, Cout]
+        }
+        // Restore [N, Cout, Hout, Wout]
+        out = out.permute([0, 2, 1]).reshape([N, Cout, Hout, Wout]);
+        // Bias
+        if (bias) {
+            out = out.add(this.handleOther(bias).reshape([1, Cout, 1, 1]));
+        }
+        return out;
+    }
     // Dropout
     dropout(rate) {
         if (!Tensor.training || rate === 0)

package/dist/nn.d.ts

@@ -55,6 +55,16 @@ export declare class LSTMCell {
     constructor(inputSize: number, hiddenSize: number, bias?: boolean, device?: string, dtype?: dtype);
     forward(input: Tensor, hidden: Tensor, cell: Tensor): [Tensor, Tensor];
 }
+export declare class Conv2d {
+    weight: Tensor;
+    bias?: Tensor;
+    stride: number | [number, number];
+    padding: number | [number, number];
+    dilation: number | [number, number];
+    groups: number;
+    constructor(inChannels: number, outChannels: number, kernelSize: number, stride?: number | [number, number], padding?: number | [number, number], dilation?: number | [number, number], groups?: number, bias?: boolean, device?: string, dtype?: dtype);
+    forward(input: Tensor): Tensor;
+}
 export declare class BatchNorm {
     weight?: Tensor;
     bias?: Tensor;
@@ -127,6 +137,7 @@ export declare const nn: {
     RNNCell: typeof RNNCell;
     GRUCell: typeof GRUCell;
     LSTMCell: typeof LSTMCell;
+    Conv2d: typeof Conv2d;
     BatchNorm: typeof BatchNorm;
     InstanceNorm: typeof InstanceNorm;
     GroupNorm: typeof GroupNorm;

package/dist/nn.js

@@ -1,6 +1,6 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.nn = exports.MultiheadAttention = exports.Embedding = exports.RMSNorm = exports.LayerNorm = exports.GroupNorm = exports.InstanceNorm = exports.BatchNorm = exports.LSTMCell = exports.GRUCell = exports.RNNCell = exports.Sequential = exports.Linear = void 0;
+exports.nn = exports.MultiheadAttention = exports.Embedding = exports.RMSNorm = exports.LayerNorm = exports.GroupNorm = exports.InstanceNorm = exports.BatchNorm = exports.Conv2d = exports.LSTMCell = exports.GRUCell = exports.RNNCell = exports.Sequential = exports.Linear = void 0;
 const core_1 = require("./core");
 class Linear {
     weight;
@@ -136,6 +136,30 @@ class LSTMCell {
     }
 }
 exports.LSTMCell = LSTMCell;
+class Conv2d {
+    weight;
+    bias;
+    stride;
+    padding;
+    dilation;
+    groups;
+    constructor(inChannels, outChannels, kernelSize, stride = 1, padding = 0, dilation = 1, groups = 1, bias = true, device, dtype) {
+        this.stride = stride;
+        this.padding = padding;
+        this.dilation = dilation;
+        this.groups = groups;
+        const fanIn = (inChannels / groups) * kernelSize * kernelSize;
+        const bound = Math.sqrt(1 / fanIn);
+        this.weight = core_1.Tensor.uniform([outChannels, inChannels / groups, kernelSize, kernelSize], -bound, bound, { requiresGrad: true, device, dtype });
+        if (bias) {
+            this.bias = core_1.Tensor.uniform([outChannels], -bound, bound, { requiresGrad: true, device, dtype });
+        }
+    }
+    forward(input) {
+        return input.conv2d(this.weight, this.bias, this.stride, this.padding, this.dilation, this.groups);
+    }
+}
+exports.Conv2d = Conv2d;
 class BatchNorm {
     weight;
     bias;
@@ -436,6 +460,7 @@ exports.nn = {
     RNNCell,
     GRUCell,
     LSTMCell,
+    Conv2d,
     BatchNorm,
     InstanceNorm,
     GroupNorm,

package/package.json

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