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/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 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