From b1527ae9b53d8c31dd63b4f93e0f685a287ce38d Mon Sep 17 00:00:00 2001
From: Sarthak Tayal <sarthaktayal2@gmail.com>
Date: Tue, 4 Aug 2026 01:51:44 -0400
Subject: [PATCH 01/11]

---
 src/matrix.js | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/src/matrix.js b/src/matrix.js
index bd8e513..5d6f5da 100644
--- a/src/matrix.js
+++ b/src/matrix.js
@@ -181,6 +181,30 @@ export class AbstractMatrix {
     return this;
   }
 
+  applyAlongAxis(callback, by) {
+    if (typeof callback !== 'function') {
+      throw new TypeError('callback must be a function');
+    }
+    switch (by) {
+      case 'row': {
+        const result = new Matrix(this.rows, 1);
+        for (let i = 0; i < this.rows; i++) {
+          result.set(i, 0, callback.call(this, this.getRow(i), i));
+        }
+        return result;
+      }
+      case 'column': {
+        const result = new Matrix(1, this.columns);
+        for (let i = 0; i < this.columns; i++) {
+          result.set(0, i, callback.call(this, this.getColumn(i), i));
+        }
+        return result;
+      }
+      default:
+        throw new Error(`invalid option: ${by}`);
+    }
+  }
+
   to1DArray() {
     let array = [];
     for (let i = 0; i < this.rows; i++) {

From 21a911ff99406a2ba4c020572526edac844681db Mon Sep 17 00:00:00 2001
From: Sarthak Tayal <sarthaktayal2@gmail.com>
Date: Tue, 4 Aug 2026 01:52:05 -0400
Subject: [PATCH 02/11] feat: add applyAlongAxis to reduce a matrix row wise or
 column wise

---
 matrix.d.ts | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/matrix.d.ts b/matrix.d.ts
index a418fe8..31314fb 100644
--- a/matrix.d.ts
+++ b/matrix.d.ts
@@ -294,6 +294,18 @@ export abstract class AbstractMatrix {
    */
   apply(callback: (row: number, column: number) => void): this;
 
+  /**
+   * Applies a callback to each row or each column of the matrix and collects the returned values.
+   * The function is called in the matrix (this) context.
+   * @param callback - Function that will be called with each row or column and its index.
+   * @param by - Iterate by 'row' or 'column'.
+   * @returns - A column vector when iterating by row, a row vector when iterating by column.
+   */
+  applyAlongAxis(
+    callback: (vector: number[], index: number) => number,
+    by: MatrixDimension,
+  ): Matrix;
+
   /**
    * Returns a new 1D array filled row by row with the matrix values.
    */

From 5e1e5f725e5ca1b0870e795c17c584f8c982343e Mon Sep 17 00:00:00 2001
From: Sarthak Tayal <sarthaktayal2@gmail.com>
Date: Tue, 4 Aug 2026 01:52:42 -0400
Subject: [PATCH 03/11]

---
 src/__tests__/matrix/applyAlongAxis.test.js | 126 ++++++++++++++++++++
 1 file changed, 126 insertions(+)
 create mode 100644 src/__tests__/matrix/applyAlongAxis.test.js

diff --git a/src/__tests__/matrix/applyAlongAxis.test.js b/src/__tests__/matrix/applyAlongAxis.test.js
new file mode 100644
index 0000000..4869226
--- /dev/null
+++ b/src/__tests__/matrix/applyAlongAxis.test.js
@@ -0,0 +1,126 @@
+import { describe, it, expect } from 'vitest';
+
+import { Matrix } from '../..';
+
+describe('applyAlongAxis', () => {
+  const matrix = new Matrix([
+    [1, 2, 3],
+    [4, 5, 6],
+  ]);
+
+  function sum(vector) {
+    return vector.reduce((accumulator, value) => accumulator + value, 0);
+  }
+
+  it('by row returns a column vector', () => {
+    const result = matrix.applyAlongAxis(sum, 'row');
+    expect(result.rows).toBe(2);
+    expect(result.columns).toBe(1);
+    expect(result.to2DArray()).toStrictEqual([[6], [15]]);
+  });
+
+  it('by column returns a row vector', () => {
+    const result = matrix.applyAlongAxis(sum, 'column');
+    expect(result.rows).toBe(1);
+    expect(result.columns).toBe(3);
+    expect(result.to2DArray()).toStrictEqual([[5, 7, 9]]);
+  });
+
+  it('the callback gets the vector and its index', () => {
+    const seen = [];
+    matrix.applyAlongAxis((vector, index) => {
+      seen.push([index, vector]);
+      return index;
+    }, 'column');
+    expect(seen).toStrictEqual([
+      [0, [1, 4]],
+      [1, [2, 5]],
+      [2, [3, 6]],
+    ]);
+  });
+
+  it('the source matrix is left untouched', () => {
+    matrix.applyAlongAxis(sum, 'row');
+    expect(matrix.to2DArray()).toStrictEqual([
+      [1, 2, 3],
+      [4, 5, 6],
+    ]);
+  });
+
+  it('works with a callback that is not a reduction', () => {
+    const result = matrix.applyAlongAxis(
+      (vector) => Math.max(...vector),
+      'row',
+    );
+    expect(result.to2DArray()).toStrictEqual([[3], [6]]);
+  });
+});
+
+describe('applyAlongAxis with degenerate matrices', () => {
+  const emptyMatrix = new Matrix(0, 0);
+  const zeroRowMatrix = new Matrix(0, 2);
+  const zeroColumnMatrix = new Matrix(3, 0);
+
+  function count(vector) {
+    return vector.length;
+  }
+
+  it('by row of a 0x0 matrix', () => {
+    const result = emptyMatrix.applyAlongAxis(count, 'row');
+    expect(result.rows).toBe(0);
+    expect(result.columns).toBe(1);
+  });
+
+  it('by column of a 0x0 matrix', () => {
+    const result = emptyMatrix.applyAlongAxis(count, 'column');
+    expect(result.rows).toBe(1);
+    expect(result.columns).toBe(0);
+  });
+
+  it('by column of a 0 row matrix', () => {
+    expect(
+      zeroRowMatrix.applyAlongAxis(count, 'column').to2DArray(),
+    ).toStrictEqual([[0, 0]]);
+  });
+
+  it('by row of a 0 column matrix', () => {
+    expect(
+      zeroColumnMatrix.applyAlongAxis(count, 'row').to2DArray(),
+    ).toStrictEqual([[0], [0], [0]]);
+  });
+});
+
+describe('applyAlongAxis error handling', () => {
+  const matrix = new Matrix([
+    [1, 2],
+    [3, 4],
+  ]);
+
+  function sum(vector) {
+    return vector.reduce((accumulator, value) => accumulator + value, 0);
+  }
+
+  it('throws when the callback is missing', () => {
+    expect(() => matrix.applyAlongAxis(undefined, 'row')).toThrow(
+      /^callback must be a function$/,
+    );
+  });
+
+  it('throws when the callback is not a function', () => {
+    expect(() => matrix.applyAlongAxis(42, 'column')).toThrow(
+      /^callback must be a function$/,
+    );
+  });
+
+  it('throws when the dimension is missing', () => {
+    expect(() => matrix.applyAlongAxis(sum)).toThrow(
+      /^invalid option: undefined$/,
+    );
+  });
+
+  it('throws when the dimension is unknown', () => {
+    expect(() => matrix.applyAlongAxis(sum, 'diagonal')).toThrow(
+      /^invalid option: diagonal$/,
+    );
+  });
+});

From 04473502c455a067d6872ee1ccdbe756206cce85 Mon Sep 17 00:00:00 2001
From: Sarthak Tayal <sarthaktayal2@gmail.com>
Date: Tue, 4 Aug 2026 01:53:26 -0400
Subject: [PATCH 04/11]

---
 src/__tests__/matrix/applyAlongAxis.test.js | 32 ++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/__tests__/matrix/applyAlongAxis.test.js b/src/__tests__/matrix/applyAlongAxis.test.js
index 4869226..4cf3907 100644
--- a/src/__tests__/matrix/applyAlongAxis.test.js
+++ b/src/__tests__/matrix/applyAlongAxis.test.js
@@ -1,6 +1,6 @@
 import { describe, it, expect } from 'vitest';
 
-import { Matrix } from '../..';
+import { Matrix, MatrixTransposeView, SymmetricMatrix } from '../..';
 
 describe('applyAlongAxis', () => {
   const matrix = new Matrix([
@@ -56,6 +56,36 @@ describe('applyAlongAxis', () => {
   });
 });
 
+describe('applyAlongAxis on other matrix kinds', () => {
+  function sum(vector) {
+    return vector.reduce((accumulator, value) => accumulator + value, 0);
+  }
+
+  it('reads a transpose view along the view dimensions', () => {
+    const view = new MatrixTransposeView(
+      new Matrix([
+        [1, 2, 3],
+        [4, 5, 6],
+      ]),
+    );
+    expect(view.applyAlongAxis(sum, 'row').to2DArray()).toStrictEqual([
+      [5],
+      [7],
+      [9],
+    ]);
+  });
+
+  it('works on a symmetric matrix', () => {
+    const symmetric = new SymmetricMatrix([
+      [1, 2],
+      [2, 3],
+    ]);
+    expect(symmetric.applyAlongAxis(sum, 'column').to2DArray()).toStrictEqual([
+      [3, 5],
+    ]);
+  });
+});
+
 describe('applyAlongAxis with degenerate matrices', () => {
   const emptyMatrix = new Matrix(0, 0);
   const zeroRowMatrix = new Matrix(0, 2);

From 54896d04039d6a5e778d556e4e74e996b20eb4a7 Mon Sep 17 00:00:00 2001
From: Sarthak Tayal <sarthaktayal2@gmail.com>
Date: Tue, 4 Aug 2026 01:54:15 -0400
Subject: [PATCH 05/11]

---
 README.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/README.md b/README.md
index f4ea79a..d39935f 100644
--- a/README.md
+++ b/README.md
@@ -122,6 +122,7 @@ var m              = A.mean();           // m = 2.75
 var product        = A.prod();           // product = -10 (product of all values of the matrix)
 var norm           = A.norm();           // norm = 10.14889156509222 (Frobenius norm of the matrix)
 var transpose      = A.transpose();      // transpose = Matrix [[1, 10], [1, -1], rows: 2, columns: 2]
+var rowMax         = A.applyAlongAxis(v => Math.max(...v), 'row'); // rowMax = Matrix [[1], [10], rows: 2, columns: 1]
 ```
 
 #### Instantiation of matrix

From 1b0e14b9b6b5bda666f8ddd0953b541c9bbda981 Mon Sep 17 00:00:00 2001
From: Sarthak Tayal <sarthaktayal2@gmail.com>
Date: Tue, 4 Aug 2026 01:54:26 -0400
Subject: [PATCH 06/11] docs: show applyAlongAxis in the readme examples

---
 README.md | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/README.md b/README.md
index d39935f..d5c397f 100644
--- a/README.md
+++ b/README.md
@@ -125,6 +125,20 @@ var transpose      = A.transpose();      // transpose = Matrix [[1, 10], [1, -1]
 var rowMax         = A.applyAlongAxis(v => Math.max(...v), 'row'); // rowMax = Matrix [[1], [10], rows: 2, columns: 1]
 ```
 
+#### Row and column wise reductions
+```js
+var M = new Matrix([
+  [1, 2, 3],
+  [4, 5, 6],
+]);
+
+var sumOf = (vector) => vector.reduce((total, value) => total + value, 0);
+
+var rowSums    = M.applyAlongAxis(sumOf, 'row');    // rowSums    = Matrix [[6], [15], rows: 2, columns: 1]
+var columnSums = M.applyAlongAxis(sumOf, 'column'); // columnSums = Matrix [[5, 7, 9], rows: 1, columns: 3]
+```
+The callback receives each row or column as a plain array along with its index, so any reduction can be expressed with it.
+
 #### Instantiation of matrix
 ```js
 var z = Matrix.zeros(3, 2); // z = Matrix [[0, 0], [0, 0], [0, 0], rows: 3, columns: 2]

From b9ad495144ee62ff49a71bbfb5a35ad3d565c78d Mon Sep 17 00:00:00 2001
From: Sarthak Tayal <sarthaktayal2@gmail.com>
Date: Tue, 4 Aug 2026 01:54:39 -0400
Subject: [PATCH 07/11]

---
 README.md | 1 -
 1 file changed, 1 deletion(-)

diff --git a/README.md b/README.md
index d5c397f..cbfb3ac 100644
--- a/README.md
+++ b/README.md
@@ -122,7 +122,6 @@ var m              = A.mean();           // m = 2.75
 var product        = A.prod();           // product = -10 (product of all values of the matrix)
 var norm           = A.norm();           // norm = 10.14889156509222 (Frobenius norm of the matrix)
 var transpose      = A.transpose();      // transpose = Matrix [[1, 10], [1, -1], rows: 2, columns: 2]
-var rowMax         = A.applyAlongAxis(v => Math.max(...v), 'row'); // rowMax = Matrix [[1], [10], rows: 2, columns: 1]
 ```
 
 #### Row and column wise reductions

From eb04c169cdb328cb7b6b5f754ae4730521cea928 Mon Sep 17 00:00:00 2001
From: Sarthak Tayal <sarthaktayal2@gmail.com>
Date: Tue, 4 Aug 2026 01:55:29 -0400
Subject: [PATCH 08/11]

---
 src/__tests__/matrix/applyAlongAxis.test.js | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/__tests__/matrix/applyAlongAxis.test.js b/src/__tests__/matrix/applyAlongAxis.test.js
index 4cf3907..c626c0b 100644
--- a/src/__tests__/matrix/applyAlongAxis.test.js
+++ b/src/__tests__/matrix/applyAlongAxis.test.js
@@ -84,6 +84,14 @@ describe('applyAlongAxis on other matrix kinds', () => {
       [3, 5],
     ]);
   });
+
+  it('always returns a plain matrix', () => {
+    const symmetric = new SymmetricMatrix([
+      [1, 2],
+      [2, 3],
+    ]);
+    expect(symmetric.applyAlongAxis(sum, 'row')).toBeInstanceOf(Matrix);
+  });
 });
 
 describe('applyAlongAxis with degenerate matrices', () => {

From 35f71e1fc9a2e149fafae6fe0543f7b18ebbd6ab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C3=ABl=20Zasso?= <targos@protonmail.com>
Date: Wed, 5 Aug 2026 13:59:43 +0200
Subject: [PATCH 09/11] fix(types): declare `this` of callback

---
 matrix.d.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/matrix.d.ts b/matrix.d.ts
index 31314fb..722610a 100644
--- a/matrix.d.ts
+++ b/matrix.d.ts
@@ -292,7 +292,7 @@ export abstract class AbstractMatrix {
    * Applies a callback for each element of the matrix. The function is called in the matrix (this) context.
    * @param callback - Function that will be called for each element in the matrix.
    */
-  apply(callback: (row: number, column: number) => void): this;
+  apply(callback: (this: this, row: number, column: number) => void): this;
 
   /**
    * Applies a callback to each row or each column of the matrix and collects the returned values.
@@ -302,7 +302,7 @@ export abstract class AbstractMatrix {
    * @returns - A column vector when iterating by row, a row vector when iterating by column.
    */
   applyAlongAxis(
-    callback: (vector: number[], index: number) => number,
+    callback: (this: this, vector: number[], index: number) => number,
     by: MatrixDimension,
   ): Matrix;
 

From 0c79eba84a85fb2872e606ace5aebadecbea3195 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C3=ABl=20Zasso?= <targos@protonmail.com>
Date: Wed, 5 Aug 2026 14:15:05 +0200
Subject: [PATCH 10/11] fix: return an array like other reduction methods

---
 README.md                                   |  4 +-
 matrix.d.ts                                 |  4 +-
 src/__tests__/matrix/applyAlongAxis.test.js | 56 ++++++++-------------
 src/matrix.js                               | 52 +++++++++----------
 4 files changed, 52 insertions(+), 64 deletions(-)

diff --git a/README.md b/README.md
index cbfb3ac..466e56c 100644
--- a/README.md
+++ b/README.md
@@ -133,8 +133,8 @@ var M = new Matrix([
 
 var sumOf = (vector) => vector.reduce((total, value) => total + value, 0);
 
-var rowSums    = M.applyAlongAxis(sumOf, 'row');    // rowSums    = Matrix [[6], [15], rows: 2, columns: 1]
-var columnSums = M.applyAlongAxis(sumOf, 'column'); // columnSums = Matrix [[5, 7, 9], rows: 1, columns: 3]
+var rowSums    = M.applyAlongAxis(sumOf, 'row');    // rowSums    = [6, 15]
+var columnSums = M.applyAlongAxis(sumOf, 'column'); // columnSums = [5, 7, 9]
 ```
 The callback receives each row or column as a plain array along with its index, so any reduction can be expressed with it.
 
diff --git a/matrix.d.ts b/matrix.d.ts
index 722610a..e53497f 100644
--- a/matrix.d.ts
+++ b/matrix.d.ts
@@ -299,12 +299,12 @@ export abstract class AbstractMatrix {
    * The function is called in the matrix (this) context.
    * @param callback - Function that will be called with each row or column and its index.
    * @param by - Iterate by 'row' or 'column'.
-   * @returns - A column vector when iterating by row, a row vector when iterating by column.
+   * @returns - An array with the reduced column or row.
    */
   applyAlongAxis(
     callback: (this: this, vector: number[], index: number) => number,
     by: MatrixDimension,
-  ): Matrix;
+  ): number[];
 
   /**
    * Returns a new 1D array filled row by row with the matrix values.
diff --git a/src/__tests__/matrix/applyAlongAxis.test.js b/src/__tests__/matrix/applyAlongAxis.test.js
index c626c0b..7f05393 100644
--- a/src/__tests__/matrix/applyAlongAxis.test.js
+++ b/src/__tests__/matrix/applyAlongAxis.test.js
@@ -1,4 +1,4 @@
-import { describe, it, expect } from 'vitest';
+import { describe, expect, it } from 'vitest';
 
 import { Matrix, MatrixTransposeView, SymmetricMatrix } from '../..';
 
@@ -14,16 +14,12 @@ describe('applyAlongAxis', () => {
 
   it('by row returns a column vector', () => {
     const result = matrix.applyAlongAxis(sum, 'row');
-    expect(result.rows).toBe(2);
-    expect(result.columns).toBe(1);
-    expect(result.to2DArray()).toStrictEqual([[6], [15]]);
+    expect(result).toStrictEqual([6, 15]);
   });
 
   it('by column returns a row vector', () => {
     const result = matrix.applyAlongAxis(sum, 'column');
-    expect(result.rows).toBe(1);
-    expect(result.columns).toBe(3);
-    expect(result.to2DArray()).toStrictEqual([[5, 7, 9]]);
+    expect(result).toStrictEqual([5, 7, 9]);
   });
 
   it('the callback gets the vector and its index', () => {
@@ -52,7 +48,17 @@ describe('applyAlongAxis', () => {
       (vector) => Math.max(...vector),
       'row',
     );
-    expect(result.to2DArray()).toStrictEqual([[3], [6]]);
+    expect(result).toStrictEqual([3, 6]);
+  });
+
+  it('receives the matrix as `this`', () => {
+    const matrix = Matrix.zeros(1, 1);
+    let that;
+    matrix.applyAlongAxis(function cb() {
+      // eslint-disable-next-line no-invalid-this
+      that = this;
+    }, 'row');
+    expect(that).toBe(matrix);
   });
 });
 
@@ -68,11 +74,7 @@ describe('applyAlongAxis on other matrix kinds', () => {
         [4, 5, 6],
       ]),
     );
-    expect(view.applyAlongAxis(sum, 'row').to2DArray()).toStrictEqual([
-      [5],
-      [7],
-      [9],
-    ]);
+    expect(view.applyAlongAxis(sum, 'row')).toStrictEqual([5, 7, 9]);
   });
 
   it('works on a symmetric matrix', () => {
@@ -80,17 +82,7 @@ describe('applyAlongAxis on other matrix kinds', () => {
       [1, 2],
       [2, 3],
     ]);
-    expect(symmetric.applyAlongAxis(sum, 'column').to2DArray()).toStrictEqual([
-      [3, 5],
-    ]);
-  });
-
-  it('always returns a plain matrix', () => {
-    const symmetric = new SymmetricMatrix([
-      [1, 2],
-      [2, 3],
-    ]);
-    expect(symmetric.applyAlongAxis(sum, 'row')).toBeInstanceOf(Matrix);
+    expect(symmetric.applyAlongAxis(sum, 'column')).toStrictEqual([3, 5]);
   });
 });
 
@@ -105,26 +97,22 @@ describe('applyAlongAxis with degenerate matrices', () => {
 
   it('by row of a 0x0 matrix', () => {
     const result = emptyMatrix.applyAlongAxis(count, 'row');
-    expect(result.rows).toBe(0);
-    expect(result.columns).toBe(1);
+    expect(result).toStrictEqual([]);
   });
 
   it('by column of a 0x0 matrix', () => {
     const result = emptyMatrix.applyAlongAxis(count, 'column');
-    expect(result.rows).toBe(1);
-    expect(result.columns).toBe(0);
+    expect(result).toStrictEqual([]);
   });
 
   it('by column of a 0 row matrix', () => {
-    expect(
-      zeroRowMatrix.applyAlongAxis(count, 'column').to2DArray(),
-    ).toStrictEqual([[0, 0]]);
+    expect(zeroRowMatrix.applyAlongAxis(count, 'column')).toStrictEqual([0, 0]);
   });
 
   it('by row of a 0 column matrix', () => {
-    expect(
-      zeroColumnMatrix.applyAlongAxis(count, 'row').to2DArray(),
-    ).toStrictEqual([[0], [0], [0]]);
+    expect(zeroColumnMatrix.applyAlongAxis(count, 'row')).toStrictEqual([
+      0, 0, 0,
+    ]);
   });
 });
 
diff --git a/src/matrix.js b/src/matrix.js
index 5d6f5da..e81c569 100644
--- a/src/matrix.js
+++ b/src/matrix.js
@@ -4,34 +4,34 @@ import rescale from 'ml-array-rescale';
 import { inspectMatrix, inspectMatrixWithOptions } from './inspect';
 import { installMathOperations } from './mathOperations';
 import {
-  sumByRow,
-  sumByColumn,
-  sumAll,
-  productByRow,
-  productByColumn,
-  productAll,
-  varianceByRow,
-  varianceByColumn,
-  varianceAll,
-  centerByRow,
-  centerByColumn,
   centerAll,
-  scaleByRow,
-  scaleByColumn,
-  scaleAll,
-  getScaleByRow,
-  getScaleByColumn,
+  centerByColumn,
+  centerByRow,
   getScaleAll,
+  getScaleByColumn,
+  getScaleByRow,
+  productAll,
+  productByColumn,
+  productByRow,
+  scaleAll,
+  scaleByColumn,
+  scaleByRow,
+  sumAll,
+  sumByColumn,
+  sumByRow,
+  varianceAll,
+  varianceByColumn,
+  varianceByRow,
 } from './stat';
 import {
-  checkRowVector,
-  checkRowIndex,
   checkColumnIndex,
+  checkColumnIndices,
   checkColumnVector,
-  checkRange,
   checkNonEmpty,
+  checkRange,
+  checkRowIndex,
   checkRowIndices,
-  checkColumnIndices,
+  checkRowVector,
 } from './util';
 
 export class AbstractMatrix {
@@ -185,24 +185,24 @@ export class AbstractMatrix {
     if (typeof callback !== 'function') {
       throw new TypeError('callback must be a function');
     }
+    const result = [];
     switch (by) {
       case 'row': {
-        const result = new Matrix(this.rows, 1);
         for (let i = 0; i < this.rows; i++) {
-          result.set(i, 0, callback.call(this, this.getRow(i), i));
+          result.push(callback.call(this, this.getRow(i), i));
         }
-        return result;
+        break;
       }
       case 'column': {
-        const result = new Matrix(1, this.columns);
         for (let i = 0; i < this.columns; i++) {
-          result.set(0, i, callback.call(this, this.getColumn(i), i));
+          result.push(callback.call(this, this.getColumn(i), i));
         }
-        return result;
+        break;
       }
       default:
         throw new Error(`invalid option: ${by}`);
     }
+    return result;
   }
 
   to1DArray() {

From c10201a80efe6dc4841e41a28a1d41372745a1b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C3=ABl=20Zasso?= <targos@protonmail.com>
Date: Wed, 5 Aug 2026 14:17:33 +0200
Subject: [PATCH 11/11] fix: allow any return type from the callback

---
 matrix.d.ts                                 | 6 +++---
 src/__tests__/matrix/applyAlongAxis.test.js | 5 +++++
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/matrix.d.ts b/matrix.d.ts
index e53497f..3634409 100644
--- a/matrix.d.ts
+++ b/matrix.d.ts
@@ -301,10 +301,10 @@ export abstract class AbstractMatrix {
    * @param by - Iterate by 'row' or 'column'.
    * @returns - An array with the reduced column or row.
    */
-  applyAlongAxis(
-    callback: (this: this, vector: number[], index: number) => number,
+  applyAlongAxis<ReturnType>(
+    callback: (this: this, vector: number[], index: number) => ReturnType,
     by: MatrixDimension,
-  ): number[];
+  ): ReturnType[];
 
   /**
    * Returns a new 1D array filled row by row with the matrix values.
diff --git a/src/__tests__/matrix/applyAlongAxis.test.js b/src/__tests__/matrix/applyAlongAxis.test.js
index 7f05393..19950ac 100644
--- a/src/__tests__/matrix/applyAlongAxis.test.js
+++ b/src/__tests__/matrix/applyAlongAxis.test.js
@@ -60,6 +60,11 @@ describe('applyAlongAxis', () => {
     }, 'row');
     expect(that).toBe(matrix);
   });
+
+  it('can return any type', () => {
+    const rows = matrix.applyAlongAxis((row) => row, 'row');
+    expect(rows).toStrictEqual(matrix.to2DArray());
+  });
 });
 
 describe('applyAlongAxis on other matrix kinds', () => {