Skip to content

Add method for computing the hessian of a scalar-valued function #126

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
marcelluethi wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add method for computing the hessian of a scalar-valued function #126

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c521e95
add ravel method to TensorTree
marcelluethi Jul 7, 2026
8837e9e
add hessian method to Autodiff
marcelluethi Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
118 changes: 75 additions & 43 deletions AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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[
Expand All @@ -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[
Expand All @@ -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)]
```

Expand All @@ -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
// .
```
Expand All @@ -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)]
```

Expand All @@ -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
// ^^
```
Expand All @@ -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
// .
```
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions core/src/main/resources/python/jax_helper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
24 changes: 24 additions & 0 deletions core/src/main/scala/dimwit/autodiff/Autodiff.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down Expand Up @@ -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)
29 changes: 29 additions & 0 deletions core/src/main/scala/dimwit/autodiff/TensorTree.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
30 changes: 30 additions & 0 deletions core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
32 changes: 32 additions & 0 deletions mdocs/AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading