From 2790a19af3208ae08f6f3644ae56464bf4b983da Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Wed, 1 Jul 2026 17:32:55 +0200 Subject: [PATCH 1/6] add important linear algebra operations This commit wraps important linear algebra operations, such as qr, cholesky, svd, solve and eigh. --- .../scala/dimwit/linalg/LinearAlgebra.scala | 203 ++++++++++++++++++ .../dimwit/linalg/LinearAlgebraTests.scala | 169 +++++++++++++++ 2 files changed, 372 insertions(+) create mode 100644 core/src/main/scala/dimwit/linalg/LinearAlgebra.scala create mode 100644 core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala diff --git a/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala b/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala new file mode 100644 index 0000000..5360dc9 --- /dev/null +++ b/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala @@ -0,0 +1,203 @@ +package dimwit.linalg + +import dimwit.jax.Jax +import dimwit.tensor.Axis +import dimwit.tensor.Label +import dimwit.tensor.Labels +import dimwit.tensor.ShapeTypeHelpers.AxesRemover +import dimwit.tensor.Tensor +import dimwit.tensor.Tensor0 +import dimwit.tensor.Tensor1 +import dimwit.tensor.Tensor2 +import dimwit.tensor.TensorOps.IsFloating +import dimwit.tensor.TensorOps.IsNumber +import me.shadaj.scalapy.py +import me.shadaj.scalapy.py.SeqConverters + +/** Common linear algebra operations. + */ +object LinearAlgebra: + + enum VectorNormType: + case L1 + case L2 + case Ord(p: Double) + case Inf + + enum MatrixNormType: + case Frobenius + case Nuclear + case Spectral + case One + case Inf + + enum QRMode: + case Reduced + case Complete + + /** Computes the determinant of the tensor `t` along the specified axes (L1, L2) + * + * @param t The input tensor from which to compute the determinant. + * @param axis1 The first axis along which to compute the determinant. + * @param axis2 The second axis along which to compute the determinant. + * @return A new tensor with the determinant computed, where the two specified axes are removed + */ + def det[T <: Tuple: Labels, L1: Label, L2: Label, V: IsFloating](t: Tensor[T, V], axis1: Axis[L1], axis2: Axis[L2])(using + ev: AxesRemover[T, (L1, L2)], + labels: Labels[ev.RemainingAxes] + ): Tensor[ev.RemainingAxes, V] = + // JAX det only works on the last two axes (-2, -1). We must move the user's selected axes to the end. + val moved = Jax.jnp.moveaxis( + t.jaxValue, + source = ev.indices.toPythonProxy, + destination = Seq(-2, -1).toPythonProxy + ) + Tensor(Jax.jnp.linalg.det(moved)) + + /** Extracts the diagonal along the given two axes (with optional offset), + * replacing them by a new 1D axis labeled L1. + * + * @param t The input tensor from which to extract the diagonal. + * @param axis1 The first axis along which to extract the diagonal. + * @param axis2 The second axis along which to extract the diagonal. + * @param offset The offset of the diagonal from the main diagonal. Positive values indicate diagonals above the main diagonal, while negative values indicate diagonals below it. + * @return A new tensor with the diagonal extracted, where the two specified axes are replaced by a new 1D axis labeled L1. + */ + def diagonal[T <: Tuple, L1: Label, L2: Label, V](t: Tensor[T, V], axis1: Axis[L1], axis2: Axis[L2], offset: Int = 0)(using + ev: AxesRemover[T, (L1, L2)], + labels: Labels[ev.RemainingAxes] + ): Tensor[ev.RemainingAxes *: L1 *: EmptyTuple, V] = + Tensor(Jax.jnp.diagonal(t.jaxValue, offset = offset, axis1 = ev.indices(0), axis2 = ev.indices(1))) + + /** Computes the inverse of the tensor t along the last two axes. + * The first axes of the tensors are preserved, while the last + * two axes are replaced by their inverses. + * The tensor must be square along the last two axes. + * + * @return a new tensor with the same shape as t, but with the last two axes replaced by their inverses. + */ + def inv[T <: Tuple: Labels, V: IsFloating](t: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.linalg.inv(t.jaxValue)) + + /** Computes the trace of the tensor `t` along the specified axes (L1, L2) with an optional offset. + * The resulting tensor has the two specified axes removed, and the remaining axes are preserved. + * + * @param t The input tensor from which to compute the trace. + * @param axis1 The first axis along which to compute the trace. + * @param axis2 The second axis along which to compute the trace. + * @param offset The offset of the diagonal from the main diagonal. Positive values indicate diagonals above the main diagonal, while negative values indicate diagonals below it. + * + * @return A new tensor with the trace computed, where the two specified axes are removed, and the remaining axes are preserved. + */ + def trace[T <: Tuple, L1: Label, L2: Label, V: IsNumber](t: Tensor[T, V], axis1: Axis[L1], axis2: Axis[L2], offset: Int = 0)(using + ev: AxesRemover[T, (L1, L2)], + labels: Labels[ev.RemainingAxes] + ): Tensor[ev.RemainingAxes, V] = Tensor(Jax.jnp.trace(t.jaxValue, offset = offset, axis1 = ev.indices(0), axis2 = ev.indices(1))) + + /** Computes the vector norm of the tensor `t` based on the specified `normType`. + * + * @param t The input tensor for which to compute the norm. + * @param normType The type of norm to compute (L1, L2, Ord(p), or Inf). + * @return A new 0-D tensor containing the computed norm of the input tensor. + */ + def norm[L1: Label, V: IsFloating](t: Tensor1[L1, V], normType: VectorNormType): Tensor0[V] = + normType match + case VectorNormType.L1 => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = 1)) + case VectorNormType.L2 => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = 2)) + case VectorNormType.Ord(p) => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = p)) + case VectorNormType.Inf => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = Jax.jnp.inf)) + + /** Computes the matrix norm of the tensor `t` based on the specified `normType`. + * + * @param t The input tensor for which to compute the norm. + * @param normType The type of norm to compute (Frobenius, Nuclear, Spectral, One, or Inf). + * @return A new 0-D tensor containing the computed norm of the input tensor. + */ + def norm[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], normType: MatrixNormType): Tensor0[V] = + normType match + case MatrixNormType.Frobenius => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = "fro")) + case MatrixNormType.Nuclear => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = "nuc")) + case MatrixNormType.Spectral => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = 2)) + case MatrixNormType.One => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = 1)) + case MatrixNormType.Inf => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = Jax.jnp.inf)) + + /** Computes the element-wise norm of the tensor `t` l2 norm along the last axis. + * @param t The input tensor for which to compute the norm. + * + * @return A new 0-D tensor containing the computed norm of the input tensor. + */ + def norm[T <: Tuple, V: IsFloating](t: Tensor[T, V]): Tensor0[V] = + Tensor0(Jax.jnp.linalg.norm(t.jaxValue)) + + /** Cholesky factorization. + * + * @param t The input tensor to be factorized. It must be a symmetric positive-definite matrix. + * @param upper If true, the upper-triangular Cholesky factor is returned + * @param symmetrizeInput If true, the input matrix is symmetrized before factorization to ensure numerical stability. + * @return a triangular matrix representing the cholesky factor + * + * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.cholesky.html#jax.numpy.linalg.cholesky JAX documentation]] for more details on the underlying implementation. + */ + def cholesky[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], upper: Boolean = false, symmetrizeInput: Boolean = true): Tensor2[L1, L2, V] = + Tensor(Jax.jnp.linalg.cholesky(t.jaxValue, upper = upper, symmetrize_input = symmetrizeInput)) + + /** Computes the QR factorization of the tensor `t`. + * + * @param t The input tensor to be factorized. It must be a 2D matrix. + * @param mode The mode of the QR factorization (Reduced or Complete). + * @return A tuple containing two tensors: the orthogonal matrix Q and the upper-triangular matrix R. + * + * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.qr.html#jax.numpy.linalg.qr JAX documentation]] for more details on the underlying implementation. + */ + def qr[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], mode: QRMode = QRMode.Reduced): (q: Tensor2[L1, L2, V], r: Tensor2[L1, L2, V]) = + val qr = Jax.jnp.linalg.qr( + t.jaxValue, + mode = mode match + case QRMode.Reduced => "reduced" + case QRMode.Complete => "complete" + ) + (q = Tensor(qr.bracketAccess(0)), r = Tensor(qr.bracketAccess(1))) + + /** Computes the eigenvalues and eigenvectors of a symmetric matrix `t`. + * @param t The input tensor representing a symmetric matrix. + * @param upper If true, the upper-triangular part of the matrix is used. + * @param symmetrizeInput If true, the input matrix is symmetrized before computation to ensure numerical stability. + * @return A tuple containing two tensors: the eigenvalues and the corresponding eigenvectors of the input matrix. + * + * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.eigh.html#jax.numpy.linalg.eigh JAX documentation]] for more details on the underlying implementation. + */ + def eigh[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], upper: Boolean = false, symmetrizeInput: Boolean = true) + : (eigenvalues: Tensor1[L1, V], eigenvectors: Tensor2[L1, L2, V]) = + + val ret = Jax.jnp.linalg.eigh(t.jaxValue, UPLO = if upper then "U" else "L", symmetrize_input = symmetrizeInput) + val eigenvalues: Tensor1[L1, V] = Tensor(ret.bracketAccess(0)) + val eigenvectors: Tensor2[L1, L2, V] = Tensor(ret.bracketAccess(1)) + (eigenvalues = eigenvalues, eigenvectors = eigenvectors) + + /** Computes the singular value decomposition (SVD) of the tensor `t`. + * + * @param t The input tensor to be decomposed. + * @param fullMatrices If true, compute the full-sized U and Vh matrices; if false, compute the reduced-sized matrices. + * @param hermitian If true, the input is a Hermitian matrix. + * @return A tuple containing three tensors: the left singular vectors U, the singular values S, and the right singular vectors Vh. + * + * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.svd.html#jax.numpy.linalg.svd JAX documentation]] for more details on the underlying implementation. + */ + def svd[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], fullMatrices: Boolean = false, hermitian: Boolean = false) + : (U: Tensor2[L1, L2, V], S: Tensor1[L1, V], Vh: Tensor2[L1, L2, V]) = + + val ret = Jax.jnp.linalg.svd(t.jaxValue, full_matrices = fullMatrices, hermitian = hermitian) + val u: Tensor2[L1, L2, V] = Tensor(ret.bracketAccess(0)) + val s: Tensor1[L1, V] = Tensor(ret.bracketAccess(1)) + val vh: Tensor2[L1, L2, V] = Tensor(ret.bracketAccess(2)) + (U = u, S = s, Vh = vh) + + /** Solves the linear equation Ax = b for x, where A is a square matrix and b is a vector. + * + * @param a The input tensor representing the square matrix A. + * @param b The input tensor representing the vector b. + * @return A new tensor containing the solution vector x. + * + * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.solve.html#jax.numpy.linalg.solve JAX documentation]] for more details on the underlying implementation. + */ + def solve[L1: Label, L2: Label, V: IsFloating](a: Tensor2[L1, L2, V], b: Tensor1[L1, V]): Tensor1[L2, V] = + Tensor(Jax.jnp.linalg.solve(a.jaxValue, b.jaxValue)) diff --git a/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala b/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala new file mode 100644 index 0000000..dbaac38 --- /dev/null +++ b/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala @@ -0,0 +1,169 @@ +package dimwit.linalg + +import dimwit.* + +class LinearAlgebraTests extends DimwitTest: + + describe("Vector norms"): + val v = Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f)) + + it("L1 norm"): + LinearAlgebra.norm(v, LinearAlgebra.VectorNormType.L1).item shouldBe 7.0f +- 1e-5f + + it("L2 norm"): + LinearAlgebra.norm(v, LinearAlgebra.VectorNormType.L2).item shouldBe 5.0f +- 1e-5f + + it("Ord(1) norm equals L1"): + LinearAlgebra.norm(v, LinearAlgebra.VectorNormType.Ord(1)).item shouldBe + LinearAlgebra.norm(v, LinearAlgebra.VectorNormType.L1).item +- 1e-5f + + it("Ord(3) norm"): + // (3^3 + 4^3)^(1/3) = (27 + 64)^(1/3) = 91^(1/3) ≈ 4.4979 + LinearAlgebra.norm(v, LinearAlgebra.VectorNormType.Ord(3)).item shouldBe + Math.pow(91.0, 1.0 / 3.0).toFloat +- 1e-4f + + it("Inf norm (max abs value)"): + LinearAlgebra.norm(v, LinearAlgebra.VectorNormType.Inf).item shouldBe 4.0f +- 1e-5f + + describe("Matrix norms"): + // [[3, 0], [4, 0]]: easy to reason about column/row sums + val m = Tensor2(Axis[A], Axis[B]).fromArray( + Array(Array(3.0f, 0.0f), Array(4.0f, 0.0f)) + ) + + it("Frobenius norm"): + // sqrt(3^2 + 4^2) = 5 + LinearAlgebra.norm(m, LinearAlgebra.MatrixNormType.Frobenius).item shouldBe 5.0f +- 1e-5f + + it("Nuclear norm"): + // singular values of [[3,0],[4,0]] are 5 and 0; nuclear = sum = 5 + LinearAlgebra.norm(m, LinearAlgebra.MatrixNormType.Nuclear).item shouldBe 5.0f +- 1e-5f + + it("Spectral norm (ord=2)"): + // largest singular value = 5 + LinearAlgebra.norm(m, LinearAlgebra.MatrixNormType.Spectral).item shouldBe 5.0f +- 1e-5f + + it("One norm (max absolute column sum)"): + // col 0 sum = 3+4=7, col 1 sum = 0 → 7 + LinearAlgebra.norm(m, LinearAlgebra.MatrixNormType.One).item shouldBe 7.0f +- 1e-5f + + it("Inf norm (max absolute row sum)"): + // row 0 sum = 3, row 1 sum = 4 → 4 + LinearAlgebra.norm(m, LinearAlgebra.MatrixNormType.Inf).item shouldBe 4.0f +- 1e-5f + + describe("Cholesky factorization"): + + val lower = LinearAlgebra.cholesky( + Tensor2(Axis[A], Axis[Prime[A]]).fromArray(Array(Array(4.0f, 0f), Array(2.0f, 3.0f))), + upper = false + ) + val spd = lower.dot(Axis[Prime[A]])(lower) // make it symmetric positive-definite + + it("Lower-triangular factor has exact values"): + val L = LinearAlgebra.cholesky(spd, upper = false) + L should approxEqual(lower, tolerance = 1e-5f) + + it("correctly reconstructs the original matrix"): + val L = LinearAlgebra.cholesky(spd, upper = false, symmetrizeInput = false) + val reconstructed = L.dot(Axis[Prime[A]])(L) + reconstructed should approxEqual(spd, tolerance = 1e-5f) + + // Shared diagonal test matrix for eigh/svd: [[3, 0], [0, 5]] + // Eigenvalues (ascending): [3, 5]; singular values (descending): [5, 3] + val diagMat = Tensor2(Axis[A], Axis[Prime[A]]).fromArray( + Array(Array(3.0f, 0.0f), Array(0.0f, 5.0f)) + ) + // identity as Tensor2[A, Prime[A]] — the type produced by contracting Prime[A] from a Tensor2[A, Prime[A]] with itself + val identityAP = Tensor2(Axis[A], Axis[Prime[A]]).fromArray( + Array(Array(1.0f, 0.0f), Array(0.0f, 1.0f)) + ) + + describe("Eigendecomposition (eigh)"): + + it("eigenvalues of a diagonal matrix are its diagonal entries (ascending)"): + val (eigenvalues, _) = LinearAlgebra.eigh(diagMat) + eigenvalues should approxEqual( + Tensor1(Axis[A]).fromArray(Array(3.0f, 5.0f)), + tolerance = 1e-5f + ) + + it("eigenvalues sum equals trace"): + val (eigenvalues, _) = LinearAlgebra.eigh(diagMat) + eigenvalues.sum.item shouldBe diagMat.sum.item +- 1e-4f + + it("eigenvectors of a diagonal matrix are the standard basis (up to sign)"): + val (_, eigenvectors) = LinearAlgebra.eigh(diagMat) + // |V| should be identity (sign-agnostic) + val absEigvecs = eigenvectors.abs + val expected = Tensor2(Axis[A], Axis[Prime[A]]).fromArray( + Array(Array(1.0f, 0.0f), Array(0.0f, 1.0f)) + ) + absEigvecs should approxEqual(expected, tolerance = 1e-5f) + + it("eigenvectors are orthonormal: V @ V^T = I"): + val (_, eigenvectors) = LinearAlgebra.eigh(diagMat) + val vvt = eigenvectors.dot(Axis[Prime[A]])(eigenvectors) + vvt should approxEqual(identityAP, tolerance = 1e-5f) + + describe("QR factorization"): + + // Non-trivial 2×2 matrix; expected properties are sign-agnostic + val qrMat = Tensor2(Axis[A], Axis[Prime[A]]).fromArray( + Array(Array(3.0f, 2.0f), Array(4.0f, 1.0f)) + ) + + it("Q is (column-)orthonormal: Q @ Q^T = I"): + val (q, _) = LinearAlgebra.qr(qrMat) + val qqt = q.dot(Axis[Prime[A]])(q) + qqt should approxEqual(identityAP, tolerance = 1e-5f) + + it("R is upper triangular: lower-left element is zero"): + val (_, r) = LinearAlgebra.qr(qrMat) + r.slice(Axis[A].at(1)).slice(Axis[Prime[A]].at(0)).item shouldBe 0.0f +- 1e-5f + + it("Frobenius norm is preserved: ||A||_F = ||R||_F (since Q is orthogonal)"): + val (_, r) = LinearAlgebra.qr(qrMat) + LinearAlgebra.norm(r, LinearAlgebra.MatrixNormType.Frobenius).item shouldBe + LinearAlgebra.norm(qrMat, LinearAlgebra.MatrixNormType.Frobenius).item +- 1e-4f + + describe("Singular value decomposition (SVD)"): + + it("singular values of a diagonal matrix are its diagonal entries (descending)"): + val (_, s, _) = LinearAlgebra.svd(diagMat) + s should approxEqual( + Tensor1(Axis[A]).fromArray(Array(5.0f, 3.0f)), + tolerance = 1e-5f + ) + + it("singular values sum equals nuclear norm"): + val (_, s, _) = LinearAlgebra.svd(diagMat) + s.sum.item shouldBe + LinearAlgebra.norm(diagMat, LinearAlgebra.MatrixNormType.Nuclear).item +- 1e-4f + + it("largest singular value equals spectral norm"): + val (_, s, _) = LinearAlgebra.svd(diagMat) + s.max.item shouldBe + LinearAlgebra.norm(diagMat, LinearAlgebra.MatrixNormType.Spectral).item +- 1e-4f + + it("U is orthonormal: U @ U^T = I"): + val (u, _, _) = LinearAlgebra.svd(diagMat) + // u: Tensor2[A, Prime[A]], contracting Prime[A] gives Tensor2[A, Prime[A]] + val uut = u.dot(Axis[Prime[A]])(u) + uut should approxEqual(identityAP, tolerance = 1e-5f) + + it("Vh is orthonormal: Vh @ Vh^T = I"): + val (_, _, vh) = LinearAlgebra.svd(diagMat) + // vh: Tensor2[A, Prime[A]], contracting Prime[A] gives Tensor2[A, Prime[A]] + val vhvht = vh.dot(Axis[Prime[A]])(vh) + vhvht should approxEqual(identityAP, tolerance = 1e-5f) + + describe("Linear solve (Ax = b)"): + // A = [[2, 1], [1, 3]], b = [5, 10] → exact solution x = [1, 3] + val solveA = Tensor2(Axis[A], Axis[Prime[A]]).fromArray( + Array(Array(2.0f, 1.0f), Array(1.0f, 3.0f)) + ) + val solveB = Tensor1(Axis[A]).fromArray(Array(5.0f, 10.0f)) + + it("solution satisfies A x = b"): + val x = LinearAlgebra.solve(solveA, solveB) + x should approxEqual(Tensor1(Axis[Prime[A]]).fromArray(Array(1.0f, 3.0f)), tolerance = 1e-5f) From c8d018edbfb4bb9362101d67406923be630ebc90 Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Thu, 2 Jul 2026 17:26:33 +0200 Subject: [PATCH 2/6] implement LinearAlgebraOps extension methods in terms using the LinearAlgebra module --- .../tensor/tensorops/LinearAlgebraOps.scala | 83 ++++++------------- 1 file changed, 27 insertions(+), 56 deletions(-) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/LinearAlgebraOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/LinearAlgebraOps.scala index db52ae1..6b67a42 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/LinearAlgebraOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/LinearAlgebraOps.scala @@ -1,6 +1,6 @@ package dimwit.tensor.tensorops -import dimwit.jax.Jax +import dimwit.linalg.LinearAlgebra import dimwit.tensor.Axis import dimwit.tensor.Label import dimwit.tensor.Labels @@ -11,9 +11,6 @@ import dimwit.tensor.Tensor1 import dimwit.tensor.Tensor2 import dimwit.tensor.TensorOps.IsFloating import dimwit.tensor.TensorOps.IsNumber -import me.shadaj.scalapy.py -import me.shadaj.scalapy.py.SeqConverters -import me.shadaj.scalapy.readwrite.Writer object LinearAlgebraOps: @@ -22,33 +19,25 @@ object LinearAlgebraOps: /** Extracts the diagonal along the given two axes (with optional offset), * replacing them by a new 1D axis labeled L1. * - * @param axis1 The first axis along which to extract the diagonal. - * @param axis2 The second axis along which to extract the diagonal. - * @param offset The offset of the diagonal from the main diagonal. Positive values indicate diagonals above the main diagonal, while negative values indicate diagonals below it. - * @return A new tensor with the diagonal extracted, where the two specified axes are replaced by a new 1D axis labeled L1. + * @see [[LinearAlgebra.diagonal]] for details */ def diagonal[L1: Label, L2: Label](axis1: Axis[L1], axis2: Axis[L2], offset: Int = 0)(using ev: AxesRemover[T, (L1, L2)], labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes *: L1 *: EmptyTuple, V] = - Tensor(Jax.jnp.diagonal(t.jaxValue, offset = offset, axis1 = ev.indices(0), axis2 = ev.indices(1))) + ): Tensor[ev.RemainingAxes *: L1 *: EmptyTuple, V] = LinearAlgebra.diagonal(t, axis1, axis2, offset) extension [L1: Label, L2: Label, V](t: Tensor2[L1, L2, V]) /** return the diagonal of the tensor `t` along the specified axes. - * The resulting 1D tensor has a single axis labeled L1, representing the diagonal index over the original (L1, L2) axes. - * - * @return A new tensor1 with representing the diagonal. It uses the Label of the first axis (L1) as the label for the resulting 1D tensor. + * @see [[LinearAlgebra.diagonal]] for details */ - def diagonal: Tensor1[L1, V] = t.diagonal(0) + def diagonal: Tensor1[L1, V] = LinearAlgebra.diagonal(t, Axis[L1], Axis[L2]).asInstanceOf[Tensor[Tuple1[L1], V]] /** return the diagonal of the tensor `t` along the specified axes. - * The resulting 1D tensor has a single axis labeled L1, representing the diagonal index over the original (L1, L2) axes. - * - * @param offset The offset of the diagonal from the main diagonal. Positive values indicate diagonals above the main diagonal, while negative values indicate diagonals below it. - * @return A new tensor1 with representing the diagonal. It uses the Label of the first axis (L1) as the label for the resulting 1D tensor. + * @see [[LinearAlgebra.diagonal]] for details */ - def diagonal(offset: Int): Tensor1[L1, V] = Tensor(Jax.jnp.diagonal(t.jaxValue, offset = offset)) + def diagonal(offset: Int): Tensor1[L1, V] = + LinearAlgebra.diagonal(t, Axis[L1], Axis[L2], offset).asInstanceOf[Tensor[Tuple1[L1], V]] // --------------------------------------------------------- // IsNumber operations (IsFloat or IsInt) @@ -56,66 +45,48 @@ object LinearAlgebraOps: extension [T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V]) - /** Computes the trace of the tensor `t` along the specified axes (L1, L2) with an optional offset. - * The resulting tensor has the two specified axes removed, and the remaining axes are preserved. - * - * @param axis1 The first axis along which to compute the trace. - * @param axis2 The second axis along which to compute the trace. - * @param offset The offset of the diagonal from the main diagonal. Positive values indicate diagonals above the main diagonal, while negative values indicate diagonals below it. + /** Computes the trace of the tensor. * - * @return A new tensor with the trace computed, where the two specified axes are removed, and the remaining axes are preserved. + * @see [[LinearAlgebra.trace]] for details */ def trace[L1: Label, L2: Label](axis1: Axis[L1], axis2: Axis[L2], offset: Int = 0)(using ev: AxesRemover[T, (L1, L2)], labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes, V] = Tensor(Jax.jnp.trace(t.jaxValue, offset = offset, axis1 = ev.indices(0), axis2 = ev.indices(1))) + ): Tensor[ev.RemainingAxes, V] = LinearAlgebra.trace(t, axis1, axis2, offset) extension [L1: Label, L2: Label, V: IsNumber](t: Tensor2[L1, L2, V]) /** Computes the trace of the tensor + * + * @see [[LinearAlgebra.trace]] for details */ def trace: Tensor0[V] = t.trace(0) - /** Computes the trace of the tensor with an optional offset. - * - * @param offset The offset of the diagonal from the main diagonal. Positive values indicate diagonals above the main diagonal, - * while negative values indicate diagonals below it. - */ - def trace(offset: Int): Tensor0[V] = Tensor0(Jax.jnp.trace(t.jaxValue, offset = offset)) + /** Computes the trace. @see [[LinearAlgebra.trace]] for details */ + def trace(offset: Int): Tensor0[V] = LinearAlgebra.trace(t, Axis[L1], Axis[L2], offset) extension [T <: Tuple: Labels, V: IsFloating](t: Tensor[T, V]) - /** Computes the L2 norm of the tensor t. + /** Computes the element wise L2 norm of the tensor t. + * + * @see [[LinearAlgebra.norm]] for details */ - def norm: Tensor0[V] = Tensor0(Jax.jnp.linalg.norm(t.jaxValue)) + def norm: Tensor0[V] = LinearAlgebra.norm(t) - /** Computes the inverse of the tensor t along the last two axes. - * The first axes of the tensors are preserved, while the last - * two axes are replaced by their inverses. - * The tensor must be square along the last two axes. - * - * @return a new tensor with the same shape as t, but with the last two axes replaced by their inverses. + /** @see [[LinearAlgebra.inv]] for details */ - def inv: Tensor[T, V] = Tensor(Jax.jnp.linalg.inv(t.jaxValue)) + def inv: Tensor[T, V] = LinearAlgebra.inv(t) /** Computes the determinant of the tensor `t` along the specified axes (L1, L2) - * - * @param axis1 The first axis along which to compute the determinant. - * @param axis2 The second axis along which to compute the determinant. - * @return A new tensor with the determinant computed, where the two specified axes are removed + * @see [[LinearAlgebra.det]] for details */ def det[L1: Label, L2: Label](axis1: Axis[L1], axis2: Axis[L2])(using ev: AxesRemover[T, (L1, L2)], labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes, V] = - // JAX det only works on the last two axes (-2, -1). We must move the user's selected axes to the end. - val moved = Jax.jnp.moveaxis( - t.jaxValue, - source = ev.indices.toPythonProxy, - destination = Seq(-2, -1).toPythonProxy - ) - Tensor(Jax.jnp.linalg.det(moved)) + ): Tensor[ev.RemainingAxes, V] = LinearAlgebra.det(t, axis1, axis2) extension [L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V]) - /** computes the determinant of the 2-D tensor t */ - def det: Tensor0[V] = Tensor0(Jax.jnp.linalg.det(t.jaxValue)) + /** computes the determinant of the 2-D tensor t + * @see [[LinearAlgebra.det]] for details + */ + def det: Tensor0[V] = LinearAlgebra.det(t, Axis[L1], Axis[L2]) From 770534ef29a8660211eb3a20a2492732c4d89f8b Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Sat, 4 Jul 2026 17:19:14 +0200 Subject: [PATCH 3/6] introduce additional axes for qr, eigh and svd --- .../scala/dimwit/linalg/LinearAlgebra.scala | 33 +++++----- .../dimwit/linalg/LinearAlgebraTests.scala | 64 ++++++++++--------- 2 files changed, 51 insertions(+), 46 deletions(-) diff --git a/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala b/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala index 5360dc9..d6b75ac 100644 --- a/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala +++ b/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala @@ -69,12 +69,8 @@ object LinearAlgebra: ): Tensor[ev.RemainingAxes *: L1 *: EmptyTuple, V] = Tensor(Jax.jnp.diagonal(t.jaxValue, offset = offset, axis1 = ev.indices(0), axis2 = ev.indices(1))) - /** Computes the inverse of the tensor t along the last two axes. - * The first axes of the tensors are preserved, while the last - * two axes are replaced by their inverses. - * The tensor must be square along the last two axes. - * - * @return a new tensor with the same shape as t, but with the last two axes replaced by their inverses. + /** Computes the inverse of the tensor t + * @return a new tensor with the same shape as t */ def inv[T <: Tuple: Labels, V: IsFloating](t: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.linalg.inv(t.jaxValue)) @@ -143,34 +139,37 @@ object LinearAlgebra: /** Computes the QR factorization of the tensor `t`. * * @param t The input tensor to be factorized. It must be a 2D matrix. + * @param basisAxis An axis Label denoting the basis axis of the output Q and R matrices. * @param mode The mode of the QR factorization (Reduced or Complete). * @return A tuple containing two tensors: the orthogonal matrix Q and the upper-triangular matrix R. * * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.qr.html#jax.numpy.linalg.qr JAX documentation]] for more details on the underlying implementation. */ - def qr[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], mode: QRMode = QRMode.Reduced): (q: Tensor2[L1, L2, V], r: Tensor2[L1, L2, V]) = + def qr[L1: Label, L2: Label, LBasis: Label, V: IsFloating](t: Tensor2[L1, L2, V], basisAxis: Axis[LBasis], mode: QRMode = QRMode.Reduced): (q: Tensor2[L1, LBasis, V], r: Tensor2[LBasis, L2, V]) = val qr = Jax.jnp.linalg.qr( t.jaxValue, mode = mode match case QRMode.Reduced => "reduced" case QRMode.Complete => "complete" ) - (q = Tensor(qr.bracketAccess(0)), r = Tensor(qr.bracketAccess(1))) + (q = Tensor[(L1, LBasis), V](qr.bracketAccess(0)), r = Tensor[(LBasis, L2), V](qr.bracketAccess(1))) /** Computes the eigenvalues and eigenvectors of a symmetric matrix `t`. * @param t The input tensor representing a symmetric matrix. + * @param eigAxis An axis Label denoting the axis of the output eigenvalues tensor. + * @param spaceAxis An axis Label denoting the axis of the output eigenvectors tensor. * @param upper If true, the upper-triangular part of the matrix is used. * @param symmetrizeInput If true, the input matrix is symmetrized before computation to ensure numerical stability. * @return A tuple containing two tensors: the eigenvalues and the corresponding eigenvectors of the input matrix. * * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.eigh.html#jax.numpy.linalg.eigh JAX documentation]] for more details on the underlying implementation. */ - def eigh[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], upper: Boolean = false, symmetrizeInput: Boolean = true) - : (eigenvalues: Tensor1[L1, V], eigenvectors: Tensor2[L1, L2, V]) = + def eigh[L1: Label, L2: Label, LEig: Label, LSpace: Label, V: IsFloating](t: Tensor2[L1, L2, V], eigAxis: Axis[LEig], spaceAxis: Axis[LSpace], upper: Boolean = false, symmetrizeInput: Boolean = true) + : (eigenvalues: Tensor1[LEig, V], eigenvectors: Tensor2[LSpace, LEig, V]) = val ret = Jax.jnp.linalg.eigh(t.jaxValue, UPLO = if upper then "U" else "L", symmetrize_input = symmetrizeInput) - val eigenvalues: Tensor1[L1, V] = Tensor(ret.bracketAccess(0)) - val eigenvectors: Tensor2[L1, L2, V] = Tensor(ret.bracketAccess(1)) + val eigenvalues: Tensor1[LEig, V] = Tensor(ret.bracketAccess(0)) + val eigenvectors: Tensor2[LSpace, LEig, V] = Tensor(ret.bracketAccess(1)) (eigenvalues = eigenvalues, eigenvectors = eigenvectors) /** Computes the singular value decomposition (SVD) of the tensor `t`. @@ -182,13 +181,13 @@ object LinearAlgebra: * * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.svd.html#jax.numpy.linalg.svd JAX documentation]] for more details on the underlying implementation. */ - def svd[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], fullMatrices: Boolean = false, hermitian: Boolean = false) - : (U: Tensor2[L1, L2, V], S: Tensor1[L1, V], Vh: Tensor2[L1, L2, V]) = + def svd[L1: Label, L2: Label, LBasis: Label, V: IsFloating](t: Tensor2[L1, L2, V], basisAxis: Axis[LBasis], fullMatrices: Boolean = false, hermitian: Boolean = false) + : (U: Tensor2[L1, LBasis, V], S: Tensor1[LBasis, V], Vh: Tensor2[LBasis, L2, V]) = val ret = Jax.jnp.linalg.svd(t.jaxValue, full_matrices = fullMatrices, hermitian = hermitian) - val u: Tensor2[L1, L2, V] = Tensor(ret.bracketAccess(0)) - val s: Tensor1[L1, V] = Tensor(ret.bracketAccess(1)) - val vh: Tensor2[L1, L2, V] = Tensor(ret.bracketAccess(2)) + val u: Tensor2[L1, LBasis, V] = Tensor(ret.bracketAccess(0)) + val s: Tensor1[LBasis, V] = Tensor(ret.bracketAccess(1)) + val vh: Tensor2[LBasis, L2, V] = Tensor(ret.bracketAccess(2)) (U = u, S = s, Vh = vh) /** Solves the linear equation Ax = b for x, where A is a square matrix and b is a vector. diff --git a/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala b/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala index dbaac38..96e9f23 100644 --- a/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala +++ b/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala @@ -73,89 +73,95 @@ class LinearAlgebraTests extends DimwitTest: val diagMat = Tensor2(Axis[A], Axis[Prime[A]]).fromArray( Array(Array(3.0f, 0.0f), Array(0.0f, 5.0f)) ) - // identity as Tensor2[A, Prime[A]] — the type produced by contracting Prime[A] from a Tensor2[A, Prime[A]] with itself - val identityAP = Tensor2(Axis[A], Axis[Prime[A]]).fromArray( + + def identity[LRow: Label, LCol: Label] = Tensor2(Axis[LRow], Axis[LCol]).fromArray( Array(Array(1.0f, 0.0f), Array(0.0f, 1.0f)) ) describe("Eigendecomposition (eigh)"): - + trait LEigen derives Label + trait LSpace derives Label it("eigenvalues of a diagonal matrix are its diagonal entries (ascending)"): - val (eigenvalues, _) = LinearAlgebra.eigh(diagMat) + val (eigenvalues, _) = LinearAlgebra.eigh(diagMat, Axis[LEigen], Axis[LSpace]) eigenvalues should approxEqual( - Tensor1(Axis[A]).fromArray(Array(3.0f, 5.0f)), + Tensor1(Axis[LEigen]).fromArray(Array(3.0f, 5.0f)), tolerance = 1e-5f ) it("eigenvalues sum equals trace"): - val (eigenvalues, _) = LinearAlgebra.eigh(diagMat) + val (eigenvalues, _) = LinearAlgebra.eigh(diagMat, Axis[LEigen], Axis[LSpace]) eigenvalues.sum.item shouldBe diagMat.sum.item +- 1e-4f it("eigenvectors of a diagonal matrix are the standard basis (up to sign)"): - val (_, eigenvectors) = LinearAlgebra.eigh(diagMat) + val (_, eigenvectors) = LinearAlgebra.eigh(diagMat, Axis[LEigen], Axis[LSpace]) // |V| should be identity (sign-agnostic) val absEigvecs = eigenvectors.abs - val expected = Tensor2(Axis[A], Axis[Prime[A]]).fromArray( + val expected = Tensor2(Axis[LSpace], Axis[LEigen]).fromArray( Array(Array(1.0f, 0.0f), Array(0.0f, 1.0f)) ) absEigvecs should approxEqual(expected, tolerance = 1e-5f) it("eigenvectors are orthonormal: V @ V^T = I"): - val (_, eigenvectors) = LinearAlgebra.eigh(diagMat) - val vvt = eigenvectors.dot(Axis[Prime[A]])(eigenvectors) - vvt should approxEqual(identityAP, tolerance = 1e-5f) + val (_, eigenvectors) = LinearAlgebra.eigh(diagMat, Axis[LEigen], Axis[LSpace]) + val vvt = eigenvectors.dot(Axis[LSpace])(eigenvectors) + val expected = identity[LEigen, Prime[LEigen]] + vvt should approxEqual(expected, tolerance = 1e-5f) describe("QR factorization"): + trait LBasis derives Label + // Non-trivial 2×2 matrix; expected properties are sign-agnostic val qrMat = Tensor2(Axis[A], Axis[Prime[A]]).fromArray( Array(Array(3.0f, 2.0f), Array(4.0f, 1.0f)) ) it("Q is (column-)orthonormal: Q @ Q^T = I"): - val (q, _) = LinearAlgebra.qr(qrMat) - val qqt = q.dot(Axis[Prime[A]])(q) - qqt should approxEqual(identityAP, tolerance = 1e-5f) + val (q, _) = LinearAlgebra.qr(qrMat, Axis[LBasis]) + val qqt = q.dot(Axis[LBasis])(q) + val expected = identity[A, Prime[A]] + qqt should approxEqual(expected, tolerance = 1e-5f) it("R is upper triangular: lower-left element is zero"): - val (_, r) = LinearAlgebra.qr(qrMat) - r.slice(Axis[A].at(1)).slice(Axis[Prime[A]].at(0)).item shouldBe 0.0f +- 1e-5f + val (_, r) = LinearAlgebra.qr(qrMat, Axis[LBasis]) + r.slice(Axis[LBasis].at(1)).slice(Axis[Prime[A]].at(0)).item shouldBe 0.0f +- 1e-5f it("Frobenius norm is preserved: ||A||_F = ||R||_F (since Q is orthogonal)"): - val (_, r) = LinearAlgebra.qr(qrMat) + val (_, r) = LinearAlgebra.qr(qrMat, Axis[LBasis]) LinearAlgebra.norm(r, LinearAlgebra.MatrixNormType.Frobenius).item shouldBe LinearAlgebra.norm(qrMat, LinearAlgebra.MatrixNormType.Frobenius).item +- 1e-4f describe("Singular value decomposition (SVD)"): - + trait LBasis derives Label it("singular values of a diagonal matrix are its diagonal entries (descending)"): - val (_, s, _) = LinearAlgebra.svd(diagMat) + val (_, s, _) = LinearAlgebra.svd(diagMat, Axis[LBasis]) s should approxEqual( - Tensor1(Axis[A]).fromArray(Array(5.0f, 3.0f)), + Tensor1(Axis[LBasis]).fromArray(Array(5.0f, 3.0f)), tolerance = 1e-5f ) it("singular values sum equals nuclear norm"): - val (_, s, _) = LinearAlgebra.svd(diagMat) + trait LBasis derives Label + val (_, s, _) = LinearAlgebra.svd(diagMat, Axis[LBasis]) s.sum.item shouldBe LinearAlgebra.norm(diagMat, LinearAlgebra.MatrixNormType.Nuclear).item +- 1e-4f it("largest singular value equals spectral norm"): - val (_, s, _) = LinearAlgebra.svd(diagMat) + val (_, s, _) = LinearAlgebra.svd(diagMat, Axis[LBasis]) s.max.item shouldBe LinearAlgebra.norm(diagMat, LinearAlgebra.MatrixNormType.Spectral).item +- 1e-4f it("U is orthonormal: U @ U^T = I"): - val (u, _, _) = LinearAlgebra.svd(diagMat) - // u: Tensor2[A, Prime[A]], contracting Prime[A] gives Tensor2[A, Prime[A]] - val uut = u.dot(Axis[Prime[A]])(u) - uut should approxEqual(identityAP, tolerance = 1e-5f) + val (u, _, _) = LinearAlgebra.svd(diagMat, Axis[LBasis]) + val uut = u.dot(Axis[A])(u) + val expected = identity[LBasis, Prime[LBasis]] + uut should approxEqual(expected, tolerance = 1e-5f) it("Vh is orthonormal: Vh @ Vh^T = I"): - val (_, _, vh) = LinearAlgebra.svd(diagMat) - // vh: Tensor2[A, Prime[A]], contracting Prime[A] gives Tensor2[A, Prime[A]] + val (_, _, vh) = LinearAlgebra.svd(diagMat, Axis[LBasis]) val vhvht = vh.dot(Axis[Prime[A]])(vh) - vhvht should approxEqual(identityAP, tolerance = 1e-5f) + val expected = identity[LBasis, Prime[LBasis]] + vhvht should approxEqual(expected, tolerance = 1e-5f) describe("Linear solve (Ax = b)"): // A = [[2, 1], [1, 3]], b = [5, 10] → exact solution x = [1, 3] From bad09569d6af02f3c42795fd9b9673b823380ba0 Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Tue, 7 Jul 2026 06:29:13 +0200 Subject: [PATCH 4/6] simplify linear algebra operation to work only on rank 2 tensors --- .../scala/dimwit/linalg/LinearAlgebra.scala | 79 +++++++------------ .../tensor/tensorops/LinearAlgebraOps.scala | 54 +++---------- 2 files changed, 39 insertions(+), 94 deletions(-) diff --git a/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala b/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala index d6b75ac..846d9ff 100644 --- a/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala +++ b/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala @@ -4,7 +4,6 @@ import dimwit.jax.Jax import dimwit.tensor.Axis import dimwit.tensor.Label import dimwit.tensor.Labels -import dimwit.tensor.ShapeTypeHelpers.AxesRemover import dimwit.tensor.Tensor import dimwit.tensor.Tensor0 import dimwit.tensor.Tensor1 @@ -12,7 +11,6 @@ import dimwit.tensor.Tensor2 import dimwit.tensor.TensorOps.IsFloating import dimwit.tensor.TensorOps.IsNumber import me.shadaj.scalapy.py -import me.shadaj.scalapy.py.SeqConverters /** Common linear algebra operations. */ @@ -35,59 +33,38 @@ object LinearAlgebra: case Reduced case Complete - /** Computes the determinant of the tensor `t` along the specified axes (L1, L2) + /** Computes the determinant of the tensor `t` * * @param t The input tensor from which to compute the determinant. - * @param axis1 The first axis along which to compute the determinant. - * @param axis2 The second axis along which to compute the determinant. - * @return A new tensor with the determinant computed, where the two specified axes are removed + * @return The determinant of the input tensor */ - def det[T <: Tuple: Labels, L1: Label, L2: Label, V: IsFloating](t: Tensor[T, V], axis1: Axis[L1], axis2: Axis[L2])(using - ev: AxesRemover[T, (L1, L2)], - labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes, V] = - // JAX det only works on the last two axes (-2, -1). We must move the user's selected axes to the end. - val moved = Jax.jnp.moveaxis( - t.jaxValue, - source = ev.indices.toPythonProxy, - destination = Seq(-2, -1).toPythonProxy - ) - Tensor(Jax.jnp.linalg.det(moved)) + def det[LRow: Label, LCol: Label, V: IsFloating](t: Tensor[(LRow, LCol), V]): Tensor0[V] = + Tensor(Jax.jnp.linalg.det(t.jaxValue)) - /** Extracts the diagonal along the given two axes (with optional offset), - * replacing them by a new 1D axis labeled L1. + /** Extracts the diagonal, with an optional offset, * * @param t The input tensor from which to extract the diagonal. - * @param axis1 The first axis along which to extract the diagonal. - * @param axis2 The second axis along which to extract the diagonal. + * @param diagAxis A new axis Label, representing the axis of the output diagonal tensor. * @param offset The offset of the diagonal from the main diagonal. Positive values indicate diagonals above the main diagonal, while negative values indicate diagonals below it. - * @return A new tensor with the diagonal extracted, where the two specified axes are replaced by a new 1D axis labeled L1. + * @return A new tensor containing the extracted diagonal elements of the input tensor. */ - def diagonal[T <: Tuple, L1: Label, L2: Label, V](t: Tensor[T, V], axis1: Axis[L1], axis2: Axis[L2], offset: Int = 0)(using - ev: AxesRemover[T, (L1, L2)], - labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes *: L1 *: EmptyTuple, V] = - Tensor(Jax.jnp.diagonal(t.jaxValue, offset = offset, axis1 = ev.indices(0), axis2 = ev.indices(1))) - - /** Computes the inverse of the tensor t - * @return a new tensor with the same shape as t + def diagonal[LRow: Label, LCol: Label, LDiag: Label, V](t: Tensor2[LRow, LCol, V], diagAxis: Axis[LDiag], offset: Int = 0): Tensor1[LDiag, V] = + Tensor(Jax.jnp.diagonal(t.jaxValue, offset = offset)) + + /** Computes the inverse of the rank 2 tensor t + * @return a new tensor, representing the inverse of t */ - def inv[T <: Tuple: Labels, V: IsFloating](t: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.linalg.inv(t.jaxValue)) + def inv[LRow: Label, LCol: Label, V: IsFloating](t: Tensor2[LRow, LCol, V]): Tensor2[LCol, LRow, V] = Tensor(Jax.jnp.linalg.inv(t.jaxValue)) - /** Computes the trace of the tensor `t` along the specified axes (L1, L2) with an optional offset. - * The resulting tensor has the two specified axes removed, and the remaining axes are preserved. + /** Computes the trace of the tensor `t` with an optional offset. * * @param t The input tensor from which to compute the trace. - * @param axis1 The first axis along which to compute the trace. - * @param axis2 The second axis along which to compute the trace. * @param offset The offset of the diagonal from the main diagonal. Positive values indicate diagonals above the main diagonal, while negative values indicate diagonals below it. * - * @return A new tensor with the trace computed, where the two specified axes are removed, and the remaining axes are preserved. + * @return The trace of the input tensor */ - def trace[T <: Tuple, L1: Label, L2: Label, V: IsNumber](t: Tensor[T, V], axis1: Axis[L1], axis2: Axis[L2], offset: Int = 0)(using - ev: AxesRemover[T, (L1, L2)], - labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes, V] = Tensor(Jax.jnp.trace(t.jaxValue, offset = offset, axis1 = ev.indices(0), axis2 = ev.indices(1))) + def trace[LRow: Label, LCol: Label, V: IsNumber](t: Tensor2[LRow, LCol, V], offset: Int = 0): Tensor0[V] = + Tensor0(Jax.jnp.trace(t.jaxValue, offset = offset)) /** Computes the vector norm of the tensor `t` based on the specified `normType`. * @@ -95,7 +72,7 @@ object LinearAlgebra: * @param normType The type of norm to compute (L1, L2, Ord(p), or Inf). * @return A new 0-D tensor containing the computed norm of the input tensor. */ - def norm[L1: Label, V: IsFloating](t: Tensor1[L1, V], normType: VectorNormType): Tensor0[V] = + def norm[L: Label, V: IsFloating](t: Tensor1[L, V], normType: VectorNormType): Tensor0[V] = normType match case VectorNormType.L1 => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = 1)) case VectorNormType.L2 => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = 2)) @@ -108,7 +85,7 @@ object LinearAlgebra: * @param normType The type of norm to compute (Frobenius, Nuclear, Spectral, One, or Inf). * @return A new 0-D tensor containing the computed norm of the input tensor. */ - def norm[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], normType: MatrixNormType): Tensor0[V] = + def norm[LRow: Label, LCol: Label, V: IsFloating](t: Tensor2[LRow, LCol, V], normType: MatrixNormType): Tensor0[V] = normType match case MatrixNormType.Frobenius => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = "fro")) case MatrixNormType.Nuclear => Tensor0(Jax.jnp.linalg.norm(t.jaxValue, ord = "nuc")) @@ -133,7 +110,7 @@ object LinearAlgebra: * * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.cholesky.html#jax.numpy.linalg.cholesky JAX documentation]] for more details on the underlying implementation. */ - def cholesky[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], upper: Boolean = false, symmetrizeInput: Boolean = true): Tensor2[L1, L2, V] = + def cholesky[LRow: Label, LCol: Label, V: IsFloating](t: Tensor2[LRow, LCol, V], upper: Boolean = false, symmetrizeInput: Boolean = true): Tensor2[LRow, LCol, V] = Tensor(Jax.jnp.linalg.cholesky(t.jaxValue, upper = upper, symmetrize_input = symmetrizeInput)) /** Computes the QR factorization of the tensor `t`. @@ -145,14 +122,14 @@ object LinearAlgebra: * * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.qr.html#jax.numpy.linalg.qr JAX documentation]] for more details on the underlying implementation. */ - def qr[L1: Label, L2: Label, LBasis: Label, V: IsFloating](t: Tensor2[L1, L2, V], basisAxis: Axis[LBasis], mode: QRMode = QRMode.Reduced): (q: Tensor2[L1, LBasis, V], r: Tensor2[LBasis, L2, V]) = + def qr[LRow: Label, LCol: Label, LBasis: Label, V: IsFloating](t: Tensor2[LRow, LCol, V], basisAxis: Axis[LBasis], mode: QRMode = QRMode.Reduced): (q: Tensor2[LRow, LBasis, V], r: Tensor2[LBasis, LCol, V]) = val qr = Jax.jnp.linalg.qr( t.jaxValue, mode = mode match case QRMode.Reduced => "reduced" case QRMode.Complete => "complete" ) - (q = Tensor[(L1, LBasis), V](qr.bracketAccess(0)), r = Tensor[(LBasis, L2), V](qr.bracketAccess(1))) + (q = Tensor[(LRow, LBasis), V](qr.bracketAccess(0)), r = Tensor[(LBasis, LCol), V](qr.bracketAccess(1))) /** Computes the eigenvalues and eigenvectors of a symmetric matrix `t`. * @param t The input tensor representing a symmetric matrix. @@ -164,7 +141,7 @@ object LinearAlgebra: * * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.eigh.html#jax.numpy.linalg.eigh JAX documentation]] for more details on the underlying implementation. */ - def eigh[L1: Label, L2: Label, LEig: Label, LSpace: Label, V: IsFloating](t: Tensor2[L1, L2, V], eigAxis: Axis[LEig], spaceAxis: Axis[LSpace], upper: Boolean = false, symmetrizeInput: Boolean = true) + def eigh[LRow: Label, LCol: Label, LEig: Label, LSpace: Label, V: IsFloating](t: Tensor2[LRow, LCol, V], eigAxis: Axis[LEig], spaceAxis: Axis[LSpace], upper: Boolean = false, symmetrizeInput: Boolean = true) : (eigenvalues: Tensor1[LEig, V], eigenvectors: Tensor2[LSpace, LEig, V]) = val ret = Jax.jnp.linalg.eigh(t.jaxValue, UPLO = if upper then "U" else "L", symmetrize_input = symmetrizeInput) @@ -181,13 +158,13 @@ object LinearAlgebra: * * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.svd.html#jax.numpy.linalg.svd JAX documentation]] for more details on the underlying implementation. */ - def svd[L1: Label, L2: Label, LBasis: Label, V: IsFloating](t: Tensor2[L1, L2, V], basisAxis: Axis[LBasis], fullMatrices: Boolean = false, hermitian: Boolean = false) - : (U: Tensor2[L1, LBasis, V], S: Tensor1[LBasis, V], Vh: Tensor2[LBasis, L2, V]) = + def svd[LRow: Label, LCol: Label, LBasis: Label, V: IsFloating](t: Tensor2[LRow, LCol, V], basisAxis: Axis[LBasis], fullMatrices: Boolean = false, hermitian: Boolean = false) + : (U: Tensor2[LRow, LBasis, V], S: Tensor1[LBasis, V], Vh: Tensor2[LBasis, LCol, V]) = val ret = Jax.jnp.linalg.svd(t.jaxValue, full_matrices = fullMatrices, hermitian = hermitian) - val u: Tensor2[L1, LBasis, V] = Tensor(ret.bracketAccess(0)) + val u: Tensor2[LRow, LBasis, V] = Tensor(ret.bracketAccess(0)) val s: Tensor1[LBasis, V] = Tensor(ret.bracketAccess(1)) - val vh: Tensor2[LBasis, L2, V] = Tensor(ret.bracketAccess(2)) + val vh: Tensor2[LBasis, LCol, V] = Tensor(ret.bracketAccess(2)) (U = u, S = s, Vh = vh) /** Solves the linear equation Ax = b for x, where A is a square matrix and b is a vector. @@ -198,5 +175,5 @@ object LinearAlgebra: * * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.solve.html#jax.numpy.linalg.solve JAX documentation]] for more details on the underlying implementation. */ - def solve[L1: Label, L2: Label, V: IsFloating](a: Tensor2[L1, L2, V], b: Tensor1[L1, V]): Tensor1[L2, V] = + def solve[LRow: Label, LCol: Label, V: IsFloating](a: Tensor2[LRow, LCol, V], b: Tensor1[LRow, V]): Tensor1[LCol, V] = Tensor(Jax.jnp.linalg.solve(a.jaxValue, b.jaxValue)) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/LinearAlgebraOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/LinearAlgebraOps.scala index 6b67a42..c61ac31 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/LinearAlgebraOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/LinearAlgebraOps.scala @@ -4,7 +4,6 @@ import dimwit.linalg.LinearAlgebra import dimwit.tensor.Axis import dimwit.tensor.Label import dimwit.tensor.Labels -import dimwit.tensor.ShapeTypeHelpers.AxesRemover import dimwit.tensor.Tensor import dimwit.tensor.Tensor0 import dimwit.tensor.Tensor1 @@ -14,46 +13,23 @@ import dimwit.tensor.TensorOps.IsNumber object LinearAlgebraOps: - extension [T <: Tuple: Labels, V](t: Tensor[T, V]) - - /** Extracts the diagonal along the given two axes (with optional offset), - * replacing them by a new 1D axis labeled L1. - * - * @see [[LinearAlgebra.diagonal]] for details - */ - def diagonal[L1: Label, L2: Label](axis1: Axis[L1], axis2: Axis[L2], offset: Int = 0)(using - ev: AxesRemover[T, (L1, L2)], - labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes *: L1 *: EmptyTuple, V] = LinearAlgebra.diagonal(t, axis1, axis2, offset) - extension [L1: Label, L2: Label, V](t: Tensor2[L1, L2, V]) - /** return the diagonal of the tensor `t` along the specified axes. + /** return the diagonal of the tensor `t` * @see [[LinearAlgebra.diagonal]] for details */ - def diagonal: Tensor1[L1, V] = LinearAlgebra.diagonal(t, Axis[L1], Axis[L2]).asInstanceOf[Tensor[Tuple1[L1], V]] + def diagonal[LDiag: Label](diagAxis: Axis[LDiag]): Tensor1[LDiag, V] = LinearAlgebra.diagonal(t, diagAxis, 0) - /** return the diagonal of the tensor `t` along the specified axes. + /** return the diagonal of the tensor `t` * @see [[LinearAlgebra.diagonal]] for details */ - def diagonal(offset: Int): Tensor1[L1, V] = - LinearAlgebra.diagonal(t, Axis[L1], Axis[L2], offset).asInstanceOf[Tensor[Tuple1[L1], V]] + def diagonal[LDiag: Label](diagAxis: Axis[LDiag], offset: Int): Tensor1[LDiag, V] = + LinearAlgebra.diagonal(t, diagAxis, offset) // --------------------------------------------------------- // IsNumber operations (IsFloat or IsInt) // --------------------------------------------------------- - extension [T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V]) - - /** Computes the trace of the tensor. - * - * @see [[LinearAlgebra.trace]] for details - */ - def trace[L1: Label, L2: Label](axis1: Axis[L1], axis2: Axis[L2], offset: Int = 0)(using - ev: AxesRemover[T, (L1, L2)], - labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes, V] = LinearAlgebra.trace(t, axis1, axis2, offset) - extension [L1: Label, L2: Label, V: IsNumber](t: Tensor2[L1, L2, V]) /** Computes the trace of the tensor @@ -63,7 +39,7 @@ object LinearAlgebraOps: def trace: Tensor0[V] = t.trace(0) /** Computes the trace. @see [[LinearAlgebra.trace]] for details */ - def trace(offset: Int): Tensor0[V] = LinearAlgebra.trace(t, Axis[L1], Axis[L2], offset) + def trace(offset: Int): Tensor0[V] = LinearAlgebra.trace(t, offset) extension [T <: Tuple: Labels, V: IsFloating](t: Tensor[T, V]) @@ -73,20 +49,12 @@ object LinearAlgebraOps: */ def norm: Tensor0[V] = LinearAlgebra.norm(t) - /** @see [[LinearAlgebra.inv]] for details - */ - def inv: Tensor[T, V] = LinearAlgebra.inv(t) - - /** Computes the determinant of the tensor `t` along the specified axes (L1, L2) + extension [LRow: Label, LCol: Label, V: IsFloating](t: Tensor2[LRow, LCol, V]) + /** computes the determinant of the 2-D tensor t * @see [[LinearAlgebra.det]] for details */ - def det[L1: Label, L2: Label](axis1: Axis[L1], axis2: Axis[L2])(using - ev: AxesRemover[T, (L1, L2)], - labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes, V] = LinearAlgebra.det(t, axis1, axis2) + def det: Tensor0[V] = LinearAlgebra.det(t) - extension [L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V]) - /** computes the determinant of the 2-D tensor t - * @see [[LinearAlgebra.det]] for details + /** @see [[LinearAlgebra.inv]] for details */ - def det: Tensor0[V] = LinearAlgebra.det(t, Axis[L1], Axis[L2]) + def inv: Tensor2[LCol, LRow, V] = LinearAlgebra.inv(t) From 98e143b95c2fc45589011ddaa18abc432b96da0d Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Tue, 7 Jul 2026 06:48:04 +0200 Subject: [PATCH 5/6] add an separate axis label LSing for the singular axis in svd --- .../scala/dimwit/linalg/LinearAlgebra.scala | 8 +++++--- .../dimwit/linalg/LinearAlgebraTests.scala | 18 ++++++++++-------- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala b/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala index 846d9ff..824836a 100644 --- a/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala +++ b/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala @@ -152,18 +152,20 @@ object LinearAlgebra: /** Computes the singular value decomposition (SVD) of the tensor `t`. * * @param t The input tensor to be decomposed. + * @param basisAxis An new axis Label denoting the basis axis of the output U and Vh matrices. + * @param singularValuesAxis An new axis Label denoting the axis of the output singular values tensor. * @param fullMatrices If true, compute the full-sized U and Vh matrices; if false, compute the reduced-sized matrices. * @param hermitian If true, the input is a Hermitian matrix. * @return A tuple containing three tensors: the left singular vectors U, the singular values S, and the right singular vectors Vh. * * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.svd.html#jax.numpy.linalg.svd JAX documentation]] for more details on the underlying implementation. */ - def svd[LRow: Label, LCol: Label, LBasis: Label, V: IsFloating](t: Tensor2[LRow, LCol, V], basisAxis: Axis[LBasis], fullMatrices: Boolean = false, hermitian: Boolean = false) - : (U: Tensor2[LRow, LBasis, V], S: Tensor1[LBasis, V], Vh: Tensor2[LBasis, LCol, V]) = + def svd[LRow: Label, LCol: Label, LBasis: Label, LSing: Label, V: IsFloating](t: Tensor2[LRow, LCol, V], basisAxis: Axis[LBasis], singularValuesAxis: Axis[LSing], fullMatrices: Boolean = false, hermitian: Boolean = false) + : (U: Tensor2[LRow, LBasis, V], S: Tensor1[LSing, V], Vh: Tensor2[LBasis, LCol, V]) = val ret = Jax.jnp.linalg.svd(t.jaxValue, full_matrices = fullMatrices, hermitian = hermitian) val u: Tensor2[LRow, LBasis, V] = Tensor(ret.bracketAccess(0)) - val s: Tensor1[LBasis, V] = Tensor(ret.bracketAccess(1)) + val s: Tensor1[LSing, V] = Tensor(ret.bracketAccess(1)) val vh: Tensor2[LBasis, LCol, V] = Tensor(ret.bracketAccess(2)) (U = u, S = s, Vh = vh) diff --git a/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala b/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala index 96e9f23..5ac45a3 100644 --- a/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala +++ b/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala @@ -133,32 +133,34 @@ class LinearAlgebraTests extends DimwitTest: describe("Singular value decomposition (SVD)"): trait LBasis derives Label + trait LSing derives Label + it("singular values of a diagonal matrix are its diagonal entries (descending)"): - val (_, s, _) = LinearAlgebra.svd(diagMat, Axis[LBasis]) + val (_, s, _) = LinearAlgebra.svd(diagMat, Axis[LBasis], Axis[LSing]) s should approxEqual( - Tensor1(Axis[LBasis]).fromArray(Array(5.0f, 3.0f)), + Tensor1(Axis[LSing]).fromArray(Array(5.0f, 3.0f)), tolerance = 1e-5f ) it("singular values sum equals nuclear norm"): trait LBasis derives Label - val (_, s, _) = LinearAlgebra.svd(diagMat, Axis[LBasis]) + val (_, s, _) = LinearAlgebra.svd(diagMat, Axis[LBasis], Axis[LSing]) s.sum.item shouldBe LinearAlgebra.norm(diagMat, LinearAlgebra.MatrixNormType.Nuclear).item +- 1e-4f it("largest singular value equals spectral norm"): - val (_, s, _) = LinearAlgebra.svd(diagMat, Axis[LBasis]) + val (_, s, _) = LinearAlgebra.svd(diagMat, Axis[LBasis], Axis[LSing]) s.max.item shouldBe LinearAlgebra.norm(diagMat, LinearAlgebra.MatrixNormType.Spectral).item +- 1e-4f it("U is orthonormal: U @ U^T = I"): - val (u, _, _) = LinearAlgebra.svd(diagMat, Axis[LBasis]) - val uut = u.dot(Axis[A])(u) - val expected = identity[LBasis, Prime[LBasis]] + val (u, _, _) = LinearAlgebra.svd(diagMat, Axis[LBasis], Axis[LSing]) + val uut = u.dot(Axis[LBasis])(u) + val expected = identity[A, Prime[A]] uut should approxEqual(expected, tolerance = 1e-5f) it("Vh is orthonormal: Vh @ Vh^T = I"): - val (_, _, vh) = LinearAlgebra.svd(diagMat, Axis[LBasis]) + val (_, _, vh) = LinearAlgebra.svd(diagMat, Axis[LBasis], Axis[LSing]) val vhvht = vh.dot(Axis[Prime[A]])(vh) val expected = identity[LBasis, Prime[LBasis]] vhvht should approxEqual(expected, tolerance = 1e-5f) From 4b463e0f6e22310d60de1ea6cd0233a177c8be18 Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Tue, 7 Jul 2026 10:51:22 +0200 Subject: [PATCH 6/6] add tests for diagonal methods --- .../dimwit/linalg/LinearAlgebraTests.scala | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala b/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala index 5ac45a3..1117822 100644 --- a/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala +++ b/core/src/test/scala/dimwit/linalg/LinearAlgebraTests.scala @@ -165,6 +165,27 @@ class LinearAlgebraTests extends DimwitTest: val expected = identity[LBasis, Prime[LBasis]] vhvht should approxEqual(expected, tolerance = 1e-5f) + describe("Diagonal extraction"): + trait LDiag derives Label + + it("extracts the main diagonal of a diagonal matrix"): + diagMat.diagonal(Axis[LDiag]) should approxEqual( + Tensor1(Axis[LDiag]).fromArray(Array(3.0f, 5.0f)), + tolerance = 1e-5f + ) + + it("extracts the main diagonal of a non-diagonal matrix"): + val m = Tensor2(Axis[A], Axis[Prime[A]]).fromArray( + Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f)) + ) + m.diagonal(Axis[LDiag]) should approxEqual( + Tensor1(Axis[LDiag]).fromArray(Array(1.0f, 4.0f)), + tolerance = 1e-5f + ) + + it("diagonal sums equal the trace"): + diagMat.diagonal(Axis[LDiag]).sum.item shouldBe diagMat.trace.item +- 1e-5f + describe("Linear solve (Ax = b)"): // A = [[2, 1], [1, 3]], b = [5, 10] → exact solution x = [1, 3] val solveA = Tensor2(Axis[A], Axis[Prime[A]]).fromArray(