From c521e95ce919429543f34e63275af722811694de Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Tue, 7 Jul 2026 21:01:31 +0200 Subject: [PATCH 1/2] add ravel method to TensorTree It is sometimes useful to be able to flatten a Tensortree into a single Vector and unflatten it again. This is, for example useful for working with Jacobian or hessian, which are only defined for tensors and tuple of tensors, but not arbitrary case classes. --- .../scala/dimwit/autodiff/TensorTree.scala | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/core/src/main/scala/dimwit/autodiff/TensorTree.scala b/core/src/main/scala/dimwit/autodiff/TensorTree.scala index 99eca73..ec7eb0c 100644 --- a/core/src/main/scala/dimwit/autodiff/TensorTree.scala +++ b/core/src/main/scala/dimwit/autodiff/TensorTree.scala @@ -2,6 +2,7 @@ package dimwit.autodiff import dimwit.jax.Jax import dimwit.tensor.* +import dimwit.tensor.DType.Float32 import me.shadaj.scalapy.py import me.shadaj.scalapy.py.SeqConverters @@ -61,6 +62,34 @@ trait TensorTree[P]: object TensorTree: // extends TensorTreeLowPriority: def apply[P](using pt: TensorTree[P]): TensorTree[P] = pt + /** Return a flatten function and unflatten function for a parameter structure. + * + * Takes a `reference` instance to capture the pytree structure (shapes of all + * leaves) for the unflatten function. The returned flatten function works on + * any `P` of the same structure. + * + * Delegates to JAX's `jax.flatten_util.ravel_pytree`. All parameters are cast + * to Float32 during flattening. + * + * Example: + * {{{ + * val (flatten, unflatten) = TensorTree.ravel(initParams, Axis[L]) + * val flat: Tensor1[L, Float32] = flatten(params) + * val reconstructed: Params = unflatten(flat) + * }}} + */ + def ravel[P, L: Label](reference: P, axis: Axis[L])(using + tt: TensorTree[P], + flatTree: TensorTree[Tensor1[L, Float32]] + ): (flatten: P => Tensor1[L, Float32], unflatten: Tensor1[L, Float32] => P) = + val flattenUtil = py.module("jax.flatten_util") + val result = flattenUtil.ravel_pytree(tt.toPyTree(reference)).as[py.Dynamic] + val unflattenPy = result.bracketAccess(1).as[py.Dynamic] + val flatten = (p: P) => + flatTree.fromPyTree(flattenUtil.ravel_pytree(tt.toPyTree(p)).as[py.Dynamic].bracketAccess(0)) + val unflatten = (v: Tensor1[L, Float32]) => tt.fromPyTree(unflattenPy(flatTree.toPyTree(v))) + (flatten = flatten, unflatten = unflatten) + /** Generic instance for any Tensor[Q, V] with labels Q and value V */ given genericTensorInstance[Q <: Tuple, V](using n: Labels[Q]): TensorTree[Tensor[Q, V]] with From 8837e9e9d98e4903596fa4b429d7c4bd325e6bdf Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Tue, 7 Jul 2026 21:02:00 +0200 Subject: [PATCH 2/2] add hessian method to Autodiff --- AGENTS.md | 118 +++++++++++------- core/src/main/resources/python/jax_helper.py | 3 + .../main/scala/dimwit/autodiff/Autodiff.scala | 24 ++++ .../scala/dimwit/autodiff/AutodiffSuite.scala | 30 +++++ mdocs/AGENTS.md | 32 +++++ 5 files changed, 164 insertions(+), 43 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 312be6a..dd3dcd4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -829,6 +829,38 @@ val jacRev = Autodiff.jacRev(linearMap) val jacFwd = Autodiff.jacFwd(linearMap) ``` +### Hessian Matrices + +```scala +import dimwit.* +import dimwit.autodiff.* + +trait A derives Label +trait B derives Label + +// Hessian of f: R -> R, f(x) = x² +val hScalar = Autodiff.hessian((x: Tensor0[Float32]) => x * x) +val xScalar = Tensor0(3.0f) +println(s"Hessian of x² at x=3: ${hScalar(xScalar)}") // 2.0 + +// Hessian of f: R² -> R, f(x) = sum(x²) +def sumSquares(x: Tensor1[A, Float32]): Tensor0[Float32] = (x * x).sum +val hVec = Autodiff.hessian(sumSquares) +val xVec = Tensor1(Axis[A]).fromArray(Array(1.0f, 5.0f)) +println(s"Hessian of sum(x²): ${hVec(xVec)}") // 2 * identity matrix + +// Block Hessian of f: R² x R² -> R, f(x1, x2) = sum(x1 * x2) +def mixed(x1: Tensor1[A, Float32], x2: Tensor1[A, Float32]): Tensor0[Float32] = + (x1 * x2).sum +val hBlock = Autodiff.hessian(mixed.tupled) +val x1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) +val x2 = Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f)) +val ((h_x1x1, h_x1x2), (h_x2x1, h_x2x2)) = hBlock(x1, x2) +println(s"Block Hessian shapes: ${h_x1x1.shape}, ${h_x1x2.shape}, ${h_x2x1.shape}, ${h_x2x2.shape}") +``` + +**Note**: `Autodiff.hessian` is only defined for scalar-output functions (`f: In => Tensor0[V]`). For vector-output functions, use `Autodiff.jacobian` instead. + --- ## Training Workflows @@ -1059,12 +1091,12 @@ trait D derives Label val intTensor = Tensor1(Axis[A]).fromArray(Array(1, 2, 3)) val wrong = intTensor.exp // exp requires IsFloating constraint // error: -// value exp is not a member of dimwit.tensor.Tensor1[MdocApp11.this.A, dimwit.tensor.DType.Int32]. +// value exp is not a member of dimwit.tensor.Tensor1[MdocApp12.this.A, dimwit.tensor.DType.Int32]. // An extension method was tried, but could not be fully constructed: // -// dimwit.exp[Tuple1[MdocApp11.this.A], +// dimwit.exp[Tuple1[MdocApp12.this.A], // (dimwit.tensor.DType.Int32 : dimwit.tensor.DType)](this.intTensor)( -// dimwit.tensor.Labels.concat[MdocApp11.this.A, EmptyTuple.type]( +// dimwit.tensor.Labels.concat[MdocApp12.this.A, EmptyTuple.type]( // this.A.derived$Label, dimwit.tensor.Labels.namesOfEmpty), // /* missing */ // summon[ @@ -1083,12 +1115,12 @@ val wrong = intTensor.exp // exp requires IsFloating constraint val boolTensor = Tensor1(Axis[A]).fromArray(Array(true, false, true)) val wrong = boolTensor.mean // error: -// value mean is not a member of dimwit.tensor.Tensor1[MdocApp11.this.A, dimwit.tensor.DType.Bool]. +// value mean is not a member of dimwit.tensor.Tensor1[MdocApp12.this.A, dimwit.tensor.DType.Bool]. // An extension method was tried, but could not be fully constructed: // -// dimwit.mean[Tuple1[MdocApp11.this.A], +// dimwit.mean[Tuple1[MdocApp12.this.A], // (dimwit.tensor.DType.Bool : dimwit.tensor.DType)](this.boolTensor)( -// dimwit.tensor.Labels.concat[MdocApp11.this.A, EmptyTuple.type]( +// dimwit.tensor.Labels.concat[MdocApp12.this.A, EmptyTuple.type]( // this.A.derived$Label, dimwit.tensor.Labels.namesOfEmpty), // /* missing */ // summon[ @@ -1112,9 +1144,9 @@ val t1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) val t2 = Tensor1(Axis[B]).fromArray(Array(3.0f, 4.0f, 5.0f)) val wrong = t1 + t2 // Different labels AND different sizes // error: -// Found: (MdocApp11.this.t2 : -// dimwit.tensor.Tensor1[MdocApp11.this.B, dimwit.tensor.DType.Float32]) -// Required: dimwit.tensor.Tensor[Tuple1[MdocApp11.this.A], +// Found: (MdocApp12.this.t2 : +// dimwit.tensor.Tensor1[MdocApp12.this.B, dimwit.tensor.DType.Float32]) +// Required: dimwit.tensor.Tensor[Tuple1[MdocApp12.this.A], // (dimwit.tensor.DType.Float32 : dimwit.tensor.DType)] ``` @@ -1124,22 +1156,22 @@ val m1 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f))) // Shape val m2 = Tensor2(Axis[C], Axis[D]).fromArray(Array(Array(3.0f), Array(4.0f))) // Shape: (2, 1) val wrong = m1.dot(Axis[B])(m2) // Axis[B] not in m2 // error: -// Axis[MdocApp11.this.B] not found in Tensor[(MdocApp11.this.C, MdocApp11.this.D)]. +// Axis[MdocApp12.this.B] not found in Tensor[(MdocApp12.this.C, MdocApp12.this.D)]. // I found: // // dimwit.tensor.ShapeTypeHelpers.AxisRemover.bridge[ -// (MdocApp11.this.C, MdocApp11.this.D), MdocApp11.this.B, R]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp11.this.C, -// MdocApp11.this.D *: EmptyTuple.type, MdocApp11.this.B]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp11.this.D, -// EmptyTuple.type, MdocApp11.this.B]( +// (MdocApp12.this.C, MdocApp12.this.D), MdocApp12.this.B, R]( +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp12.this.C, +// MdocApp12.this.D *: EmptyTuple.type, MdocApp12.this.B]( +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp12.this.D, +// EmptyTuple.type, MdocApp12.this.B]( // dimwit.tensor.ShapeTypeHelpers.AxisIndex.concatRight[A, B², L]) // ), // ???) // -// But given instance concatRight in object AxisIndex does not match type dimwit.tensor.ShapeTypeHelpers.AxisIndex[EmptyTuple.type, MdocApp11.this.B] +// But given instance concatRight in object AxisIndex does not match type dimwit.tensor.ShapeTypeHelpers.AxisIndex[EmptyTuple.type, MdocApp12.this.B] // -// where: B is a trait in class MdocApp11 +// where: B is a trait in class MdocApp12 // B² is a type variable with constraint <: Tuple // . ``` @@ -1152,7 +1184,7 @@ val t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f))) val wrong = t + 10.0f // Should use +! for scalar broadcast // error: // Found: (10.0f : Float) -// Required: dimwit.tensor.Tensor[(MdocApp11.this.A, MdocApp11.this.B), +// Required: dimwit.tensor.Tensor[(MdocApp12.this.A, MdocApp12.this.B), // (dimwit.tensor.DType.Float32 : dimwit.tensor.DType)] ``` @@ -1163,32 +1195,32 @@ val t2 = Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f)) // This works but is semantically wrong (use + instead) val wrong = t1 +! t2 // error: -// Cannot broadcast tensors of shapes Tuple1[MdocApp11.this.A] and Tuple1[MdocApp11.this.A]. If same shape no broadcasting allowed!. +// Cannot broadcast tensors of shapes Tuple1[MdocApp12.this.A] and Tuple1[MdocApp12.this.A]. If same shape no broadcasting allowed!. // I found: // // dimwit.tensor.tensorops.TensorOpsUtil.Broadcast.broadcastLeft[ -// Tuple1[MdocApp11.this.A], Tuple1[MdocApp11.this.A], +// Tuple1[MdocApp12.this.A], Tuple1[MdocApp12.this.A], // (dimwit.tensor.DType.Float32 : dimwit.tensor.DType)]( -// dimwit.tensor.Labels.concat[MdocApp11.this.A, EmptyTuple.type]( +// dimwit.tensor.Labels.concat[MdocApp12.this.A, EmptyTuple.type]( // this.A.derived$Label, dimwit.tensor.Labels.namesOfEmpty), -// dimwit.tensor.Labels.concat[MdocApp11.this.A, EmptyTuple.type]( +// dimwit.tensor.Labels.concat[MdocApp12.this.A, EmptyTuple.type]( // this.A.derived$Label, dimwit.tensor.Labels.namesOfEmpty), -// dimwit.tensor.TupleHelpers.StrictSubset.derive[Tuple1[MdocApp11.this.A], -// Tuple1[MdocApp11.this.A]]( -// dimwit.tensor.TupleHelpers.Subset.head[MdocApp11.this.A, EmptyTuple.type, -// Tuple1[MdocApp11.this.A]]( -// dimwit.tensor.TupleHelpers.SetMember.found[MdocApp11.this.A, +// dimwit.tensor.TupleHelpers.StrictSubset.derive[Tuple1[MdocApp12.this.A], +// Tuple1[MdocApp12.this.A]]( +// dimwit.tensor.TupleHelpers.Subset.head[MdocApp12.this.A, EmptyTuple.type, +// Tuple1[MdocApp12.this.A]]( +// dimwit.tensor.TupleHelpers.SetMember.found[MdocApp12.this.A, // EmptyTuple.type], -// dimwit.tensor.TupleHelpers.Subset.empty[Tuple1[MdocApp11.this.A]]), +// dimwit.tensor.TupleHelpers.Subset.empty[Tuple1[MdocApp12.this.A]]), // /* missing */ // summon[ -// scala.util.NotGiven[Tuple1[MdocApp11.this.A] =:= -// Tuple1[MdocApp11.this.A]] +// scala.util.NotGiven[Tuple1[MdocApp12.this.A] =:= +// Tuple1[MdocApp12.this.A]] // ] // ) // ) // -// But no implicit values were found that match type scala.util.NotGiven[Tuple1[MdocApp11.this.A] =:= Tuple1[MdocApp11.this.A]]. +// But no implicit values were found that match type scala.util.NotGiven[Tuple1[MdocApp12.this.A] =:= Tuple1[MdocApp12.this.A]]. // val wrong = t1 +! t2 // ^^ ``` @@ -1200,24 +1232,24 @@ val wrong = t1 +! t2 val t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f))) val wrong = t.sum(Axis[C]) // Axis[C] not in tensor // error: -// Axis[MdocApp11.this.C] not found in Tensor[(MdocApp11.this.A, MdocApp11.this.B)]. +// Axis[MdocApp12.this.C] not found in Tensor[(MdocApp12.this.A, MdocApp12.this.B)]. // I found: // // dimwit.tensor.ShapeTypeHelpers.AxisRemover.bridge[ -// (MdocApp11.this.A, MdocApp11.this.B), MdocApp11.this.C, R]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp11.this.A, -// MdocApp11.this.B *: EmptyTuple.type, MdocApp11.this.C]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp11.this.B, -// EmptyTuple.type, MdocApp11.this.C]( +// (MdocApp12.this.A, MdocApp12.this.B), MdocApp12.this.C, R]( +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp12.this.A, +// MdocApp12.this.B *: EmptyTuple.type, MdocApp12.this.C]( +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp12.this.B, +// EmptyTuple.type, MdocApp12.this.C]( // dimwit.tensor.ShapeTypeHelpers.AxisIndex.concatRight[A², B², L]) // ), // ???) // -// But given instance concatRight in object AxisIndex does not match type dimwit.tensor.ShapeTypeHelpers.AxisIndex[EmptyTuple.type, MdocApp11.this.C] +// But given instance concatRight in object AxisIndex does not match type dimwit.tensor.ShapeTypeHelpers.AxisIndex[EmptyTuple.type, MdocApp12.this.C] // -// where: A is a trait in class MdocApp11 +// where: A is a trait in class MdocApp12 // A² is a type variable with constraint <: Tuple -// B is a trait in class MdocApp11 +// B is a trait in class MdocApp12 // B² is a type variable with constraint <: Tuple // . ``` @@ -1227,7 +1259,7 @@ val wrong = t.sum(Axis[C]) // Axis[C] not in tensor val t = Tensor2(Axis[A], Axis[B]).fill(1.0f) val wrong = t.vmap(Axis[C])(_.sum) // Axis[C] doesn't exist // error: -// value fill is not a member of dimwit.tensor.Tensor2.DefaultsFactory[MdocApp11.this.A, MdocApp11.this.B] +// value fill is not a member of dimwit.tensor.Tensor2.DefaultsFactory[MdocApp12.this.A, MdocApp12.this.B] ``` ### Gradient Errors @@ -1269,8 +1301,8 @@ val wrong = Autodiff.grad(nonScalar) // Use jacobian instead // t2Tree²: dimwit.autodiff.TensorTree[T2²], outTree³: // dimwit.autodiff.TensorTree[dimwit.tensor.Tensor0[V³]]): (T1², T2²) => // dimwit.autodiff.Grad[(T1², T2²)] -// match arguments (dimwit.tensor.Tensor1[MdocApp11.this.A, dimwit.Float32] => -// dimwit.tensor.Tensor1[MdocApp11.this.A, dimwit.Float32]) +// match arguments (dimwit.tensor.Tensor1[MdocApp12.this.A, dimwit.Float32] => +// dimwit.tensor.Tensor1[MdocApp12.this.A, dimwit.Float32]) // // where: T1 is a type variable // T1² is a type variable diff --git a/core/src/main/resources/python/jax_helper.py b/core/src/main/resources/python/jax_helper.py index 3218119..fe1b761 100644 --- a/core/src/main/resources/python/jax_helper.py +++ b/core/src/main/resources/python/jax_helper.py @@ -77,6 +77,9 @@ def jacrev(f): def jacobian(f): return wrap(jax.jacobian, f) +def hessian(f): + return wrap(jax.hessian, f) + def jit(f): return wrap(jax.jit, f) diff --git a/core/src/main/scala/dimwit/autodiff/Autodiff.scala b/core/src/main/scala/dimwit/autodiff/Autodiff.scala index 871733c..03f8ee1 100644 --- a/core/src/main/scala/dimwit/autodiff/Autodiff.scala +++ b/core/src/main/scala/dimwit/autodiff/Autodiff.scala @@ -21,6 +21,13 @@ object Autodiff: case h *: t => GradientTensorVsInput[h, OutShape, V] *: GradientTensorVsInput[t, OutShape, V] case Tensor[inS, v2] => Tensor[PrimeConcatType[OutShape, inS], V] + type Hessian[In] = HessianProduct[In, In] + + type HessianProduct[In, Out] = Out match + case EmptyTuple => EmptyTuple + case h *: t => HessianProduct[In, h] *: HessianProduct[In, t] + case Tensor[outS, v] => GradientTensorVsInput[In, outS, v] + // TODO replace with TupledFunction when available (no longer experimental) def grad[T1, T2, V: IsFloating](f: (T1, T2) => Tensor0[V])(using t1Tree: TensorTree[T1], t2Tree: TensorTree[T2], outTree: TensorTree[Tensor0[V]]): (T1, T2) => Grad[(T1, T2)] = (t1, t2) => grad(f.tupled)((t1, t2)) def grad[T1, T2, T3, V: IsFloating](f: (T1, T2, T3) => Tensor0[V])(using t1Tree: TensorTree[T1], t2Tree: TensorTree[T2], t3Tree: TensorTree[T3], outTree: TensorTree[Tensor0[V]]): (T1, T2, T3) => Grad[(T1, T2, T3)] = (t1, t2, t3) => grad(f.tupled)((t1, t2, t3)) @@ -103,3 +110,20 @@ object Autodiff: outTree.toPyTree(f(inTree.fromPyTree(jxpr))) val jpy = Jax.jax_helper.jacfwd(fpy) (params: In) => gradTree.fromPyTree(jpy(inTree.toPyTree(params))) + + def hessian[In, V: IsFloating](f: In => Tensor0[V])(using + inTree: TensorTree[In], + outTree: TensorTree[Tensor0[V]], + hessTree: TensorTree[Hessian[In]] + ): In => Hessian[In] = + val fpy = (jxpr: py.Dynamic) => + OnError.traceStack: + val x = inTree.fromPyTree(jxpr) + outTree.toPyTree(f(x)) + + val hpy = Jax.jax_helper.hessian(fpy) + + (params: In) => + val xpy = inTree.toPyTree(params) + val res = hpy(xpy) + hessTree.fromPyTree(res) diff --git a/core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala b/core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala index af60d77..f9cdd4b 100644 --- a/core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala +++ b/core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala @@ -107,6 +107,36 @@ class AutodiffSuite extends DimwitTest: x2_dx1 should approxEqual(Tensor2.eye(x2.extent(Axis[A]), x2.vtype) *! Tensor0(1.0f)) x2_dx2 should approxEqual(Tensor.like(x2_dx2).fill(0f)) + describe("hessian"): + describe("single parameter function"): + it("Hessian of f(x) = x^2"): + def f(x: Tensor0[Float32]) = x * x + val hf = Autodiff.hessian(f) + + val x = Tensor0(3.0f) + hf(x) shouldEqual Tensor0(2.0f) + + it("Hessian of f(x) = sum(x^2)"): + def f(x: Tensor1[A, Float32]) = (x * x).sum + val hf = Autodiff.hessian(f) + + val x = Tensor1(Axis[A]).fromArray(Array(1.0f, 5.0f)) + hf(x) should approxEqual(Tensor2.eye(x.extent(Axis[A]), x.vtype) *! 2.0f) + + it("Hessian of f(x1, x2) = sum(x1 * x2)"): + def f(x1: Tensor1[A, Float32], x2: Tensor1[A, Float32]): Tensor0[Float32] = (x1 * x2).sum + val hf = Autodiff.hessian(f.tupled) + + val x1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) + val x2 = Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f)) + val (x1Grad, x2Grad) = hf(x1, x2) + val (x1_dx1, x1_dx2) = x1Grad + val (x2_dx1, x2_dx2) = x2Grad + x1_dx1 should approxEqual(Tensor.like(x1_dx1).fill(0f)) + x1_dx2 should approxEqual(Tensor2.eye(x1.extent(Axis[A]), x1.vtype) *! Tensor0(1.0f)) + x2_dx1 should approxEqual(Tensor2.eye(x2.extent(Axis[A]), x2.vtype) *! Tensor0(1.0f)) + x2_dx2 should approxEqual(Tensor.like(x2_dx2).fill(0f)) + describe("Complex application"): it("case class support"): case class Params(w: Tensor1[A, Float32], b: Tensor0[Float32]) diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md index 3904081..0cff8cf 100644 --- a/mdocs/AGENTS.md +++ b/mdocs/AGENTS.md @@ -636,6 +636,38 @@ val jacRev = Autodiff.jacRev(linearMap) val jacFwd = Autodiff.jacFwd(linearMap) ``` +### Hessian Matrices + +```scala mdoc:reset:silent +import dimwit.* +import dimwit.autodiff.* + +trait A derives Label +trait B derives Label + +// Hessian of f: R -> R, f(x) = x² +val hScalar = Autodiff.hessian((x: Tensor0[Float32]) => x * x) +val xScalar = Tensor0(3.0f) +println(s"Hessian of x² at x=3: ${hScalar(xScalar)}") // 2.0 + +// Hessian of f: R² -> R, f(x) = sum(x²) +def sumSquares(x: Tensor1[A, Float32]): Tensor0[Float32] = (x * x).sum +val hVec = Autodiff.hessian(sumSquares) +val xVec = Tensor1(Axis[A]).fromArray(Array(1.0f, 5.0f)) +println(s"Hessian of sum(x²): ${hVec(xVec)}") // 2 * identity matrix + +// Block Hessian of f: R² x R² -> R, f(x1, x2) = sum(x1 * x2) +def mixed(x1: Tensor1[A, Float32], x2: Tensor1[A, Float32]): Tensor0[Float32] = + (x1 * x2).sum +val hBlock = Autodiff.hessian(mixed.tupled) +val x1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) +val x2 = Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f)) +val ((h_x1x1, h_x1x2), (h_x2x1, h_x2x2)) = hBlock(x1, x2) +println(s"Block Hessian shapes: ${h_x1x1.shape}, ${h_x1x2.shape}, ${h_x2x1.shape}, ${h_x2x2.shape}") +``` + +**Note**: `Autodiff.hessian` is only defined for scalar-output functions (`f: In => Tensor0[V]`). For vector-output functions, use `Autodiff.jacobian` instead. + --- ## Training Workflows