From 6f9c1a27af14eeeea2152510160993bcae59ed9f Mon Sep 17 00:00:00 2001
From: wj <126436871+WenJing95@users.noreply.github.com>
Date: Wed, 9 Sep 2026 07:28:05 +0900
Subject: [PATCH] Fix sum return types for dimensioned collections

---
 test/typescript-tests/testTypes.ts | 49 ++++++++++++++++++++++++++++++
 types/index.d.ts                   | 34 ++++++++++++++++-----
 2 files changed, 75 insertions(+), 8 deletions(-)

diff --git a/test/typescript-tests/testTypes.ts b/test/typescript-tests/testTypes.ts
index 9a6d066481..5d0d4d68df 100644
--- a/test/typescript-tests/testTypes.ts
+++ b/test/typescript-tests/testTypes.ts
@@ -3061,6 +3061,55 @@ Statistics functions' return types
     number | BigNumber | bigint | Fraction | Complex | Unit
   >()
 
+  const sumArray = [
+    [1, 2, 3],
+    [4, 5, 6]
+  ]
+  const sumMatrix = math.matrix(sumArray)
+  expectTypeOf(math.sum(1, 2, 3)).toMatchTypeOf<number>()
+  expectTypeOf(math.sum([1, 2, 3])).toEqualTypeOf<number>()
+  expectTypeOf(math.sum(sumArray)).toEqualTypeOf<MathScalarType>()
+  expectTypeOf(math.sum(sumMatrix)).toEqualTypeOf<MathScalarType>()
+  expectTypeOf(math.sum([1, 2, 3], 0)).toEqualTypeOf<number>()
+  expectTypeOf(math.sum(sumArray, 0)).toEqualTypeOf<number[]>()
+  expectTypeOf(math.sum(sumArray, math.bignumber(1))).toEqualTypeOf<number[]>()
+  expectTypeOf(math.sum(sumMatrix, 0)).toEqualTypeOf<Matrix>()
+  expectTypeOf(math.sum(sumMatrix, math.bignumber(1))).toEqualTypeOf<Matrix>()
+  assert.deepStrictEqual(math.sum(sumMatrix, 0).toArray(), [5, 7, 9])
+
+  const sumNumberLiterals = math.sum(
+    [
+      [1, 2],
+      [3, 4]
+    ],
+    0
+  )
+  expectTypeOf(sumNumberLiterals).toEqualTypeOf<number[]>()
+  assert.deepStrictEqual(sumNumberLiterals, [4, 6])
+  const sumBigintLiterals = math.sum([[BigInt(1) as 1n, BigInt(2) as 2n]], 1)
+  expectTypeOf(sumBigintLiterals).toEqualTypeOf<bigint[]>()
+  assert.deepStrictEqual(sumBigintLiterals, [BigInt(3)])
+  expectTypeOf(math.sum(math.matrix<1 | 2>([[1, 2]]), 0)).toEqualTypeOf<
+    Matrix<number>
+  >()
+
+  const sumBigNumbers = [[math.bignumber(1), math.bignumber(2)]]
+  expectTypeOf(math.sum(sumBigNumbers, 0)).toEqualTypeOf<BigNumber[]>()
+  expectTypeOf(
+    math.sum(math.matrix<BigNumber>(sumBigNumbers), 0)
+  ).toEqualTypeOf<Matrix<BigNumber>>()
+
+  const sumArray3d = [sumArray, sumArray]
+  expectTypeOf(math.sum(sumArray3d, 0)).toEqualTypeOf<
+    MathScalarType | MathArray<MathScalarType>
+  >()
+  expectTypeOf(math.sum(math.matrix(sumArray3d), 0)).toEqualTypeOf<Matrix>()
+
+  const sumCollection = sumArray as MathCollection
+  expectTypeOf(math.sum(sumCollection, 0)).toEqualTypeOf<
+    MathScalarType | MathCollection
+  >()
+
   expectTypeOf(math.quantileSeq([1, 2, 3], 0.75)).toMatchTypeOf<number>()
   expectTypeOf(math.quantileSeq([1, 2, 3, 4, 5], [0.25, 0.75])).toMatchTypeOf<
     MathArray | MathScalarType
diff --git a/types/index.d.ts b/types/index.d.ts
index d37ff59f22..f017475c4b 100644
--- a/types/index.d.ts
+++ b/types/index.d.ts
@@ -13,6 +13,8 @@ export type NoLiteralType<T> = T extends number
       ? boolean
       : T
 
+type SumScalarType<T> = T extends bigint ? bigint : NoLiteralType<T>
+
 export type MathNumericType = number | BigNumber | bigint | Fraction | Complex
 export type MathScalarType = MathNumericType | Unit
 export type MathGeneric<T extends MathScalarType = MathNumericType> = T
@@ -3254,19 +3256,35 @@ export interface MathJsInstance extends MathJsFactory {
   sum(...args: MathScalarType[]): MathScalarType
   /**
    * @param A A single matrix
-   * @param dimension The sum over the selected dimension
    * @returns The sum of all values
    */
-  sum<T extends MathScalarType>(
-    A: T[] | T[][],
-    dimension?: number | BigNumber
-  ): T
+  sum<T extends MathScalarType>(A: T[] | T[][]): T
+  sum(A: MathCollection): MathScalarType
   /**
    * @param A A single matrix
-   * @param dimension The sum over the selected dimension
-   * @returns The sum of all values
+   * @param dimension The dimension along which to sum
+   * @returns The sums along the selected dimension
    */
-  sum(A: MathCollection, dimension?: number | BigNumber): MathScalarType
+  sum<T extends MathScalarType>(
+    A: T[],
+    dimension: number | BigNumber
+  ): SumScalarType<T>
+  sum<T extends MathScalarType>(
+    A: T[][],
+    dimension: number | BigNumber
+  ): SumScalarType<T>[]
+  sum<T extends MathScalarType>(
+    A: MathArray<T>,
+    dimension: number | BigNumber
+  ): SumScalarType<T> | MathArray<SumScalarType<T>>
+  sum<T extends MathScalarType>(
+    A: Matrix<T>,
+    dimension: number | BigNumber
+  ): Matrix<SumScalarType<T>>
+  sum(
+    A: MathCollection,
+    dimension: number | BigNumber
+  ): MathScalarType | MathCollection
 
   /**
    * Count the number of elements of a matrix, array or string.