From 96c575c93e38056dc476dfe9bfc494916f855bc3 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Sun, 16 Aug 2026 10:04:29 +0200 Subject: [PATCH] Fix issues found in the type logic. Some type logic uses match types which often can't be reduced. Switching to more powerful type class inference. * Add test cases. Most failed on old type logic (issue detection) * Fixes of issues. Tests now run * Overall cleanup: Remove unused or duplicated type logic --- AGENTS.md | 58 ++-- .../main/scala/dimwit/autodiff/Autodiff.scala | 124 ++++++-- .../scala/dimwit/linalg/LinearAlgebra.scala | 2 +- core/src/main/scala/dimwit/package.scala | 23 +- core/src/main/scala/dimwit/prime/Prime.scala | 92 ++++++ core/src/main/scala/dimwit/tensor/DType.scala | 2 +- .../src/main/scala/dimwit/tensor/Labels.scala | 29 +- .../dimwit/tensor/ShapeTypeHelpers.scala | 22 +- .../src/main/scala/dimwit/tensor/Tensor.scala | 2 +- .../scala/dimwit/tensor/TensorEvidence.scala | 119 ++++++++ .../scala/dimwit/tensor/TupleHelpers.scala | 274 ++++-------------- .../tensor/tensorops/ContractionOps.scala | 9 +- .../tensor/tensorops/ConvolutionOps.scala | 3 +- .../tensor/tensorops/FunctionalOps.scala | 65 +++-- .../tensor/tensorops/StructuralOps.scala | 140 ++++----- .../scala/dimwit/tensortree/TensorTree.scala | 48 +-- .../dimwit/tensortree/TensorTreeFormat.scala | 1 - .../main/scala/dimwit/tensortree/TreeOf.scala | 8 +- .../scala/dimwit/autodiff/AutodiffSuite.scala | 212 +++++++++++--- .../tensor/TensorOpsFunctionalSuite.scala | 31 ++ .../tensor/TensorOpsStructureSuite.scala | 57 +++- .../scala/dimwit/dataset/MNISTLoader.scala | 6 +- 22 files changed, 812 insertions(+), 515 deletions(-) create mode 100644 core/src/main/scala/dimwit/prime/Prime.scala create mode 100644 core/src/main/scala/dimwit/tensor/TensorEvidence.scala diff --git a/AGENTS.md b/AGENTS.md index af0d497..d24a4de 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -104,9 +104,9 @@ val wrong = t.sum(Axis[C]) // dimwit.tensor.ShapeTypeHelpers.AxisRemover.bridge[ // (repl.MdocSession.MdocApp.A, repl.MdocSession.MdocApp.B), // repl.MdocSession.MdocApp.C, R]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[repl.MdocSession.MdocApp.A, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[repl.MdocSession.MdocApp.A, // repl.MdocSession.MdocApp.B *: EmptyTuple.type, repl.MdocSession.MdocApp.C]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[repl.MdocSession.MdocApp.B, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[repl.MdocSession.MdocApp.B, // EmptyTuple.type, repl.MdocSession.MdocApp.C]( // dimwit.tensor.ShapeTypeHelpers.AxisIndex.concatRight[A², B², L]) // ), @@ -209,7 +209,7 @@ val summed = wrongAxis.sum(Axis[B]) // B not in shape! // // dimwit.tensor.ShapeTypeHelpers.AxisRemover.bridge[ // Tuple1[repl.MdocSession.MdocApp.A], repl.MdocSession.MdocApp.B, R]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[repl.MdocSession.MdocApp.A, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[repl.MdocSession.MdocApp.A, // EmptyTuple.type, repl.MdocSession.MdocApp.B]( // dimwit.tensor.ShapeTypeHelpers.AxisIndex.concatRight[A², B², L]), // ???) @@ -316,9 +316,9 @@ val wrong = t.sum(Axis[C]) // // dimwit.tensor.ShapeTypeHelpers.AxisRemover.bridge[ // (MdocApp0.this.A, MdocApp0.this.B), MdocApp0.this.C, R]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp0.this.A, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp0.this.A, // MdocApp0.this.B *: EmptyTuple.type, MdocApp0.this.C]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp0.this.B, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp0.this.B, // EmptyTuple.type, MdocApp0.this.C]( // dimwit.tensor.ShapeTypeHelpers.AxisIndex.concatRight[A², B², L]) // ), @@ -437,9 +437,9 @@ val wrong = m1.dot(Axis[B])(m2) // // dimwit.tensor.ShapeTypeHelpers.AxisRemover.bridge[ // (MdocApp1.this.C, MdocApp1.this.D), MdocApp1.this.B, R]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp1.this.C, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp1.this.C, // MdocApp1.this.D *: EmptyTuple.type, MdocApp1.this.B]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp1.this.D, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp1.this.D, // EmptyTuple.type, MdocApp1.this.B]( // dimwit.tensor.ShapeTypeHelpers.AxisIndex.concatRight[A, B², L]) // ), @@ -456,9 +456,9 @@ val wrong = m1.dot(Axis[B])(m2) // // dimwit.tensor.ShapeTypeHelpers.AxisRemover.bridge[ // (MdocApp1.this.C, MdocApp1.this.D), MdocApp1.this.B, R]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp1.this.C, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp1.this.C, // MdocApp1.this.D *: EmptyTuple.type, MdocApp1.this.B]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp1.this.D, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp1.this.D, // EmptyTuple.type, MdocApp1.this.B]( // dimwit.tensor.ShapeTypeHelpers.AxisIndex.concatRight[A, B², L]) // ), @@ -1100,8 +1100,8 @@ val wrong = intTensor.exp // exp requires IsFloating constraint // An extension method was tried, but could not be fully constructed: // // dimwit.exp[Tuple1[MdocApp12.this.A], dimwit.tensor.DType.Int32](this.intTensor)( -// dimwit.tensor.Labels.concat[MdocApp12.this.A, EmptyTuple.type]( -// this.A.derived$Label, dimwit.tensor.Labels.namesOfEmpty), +// dimwit.tensor.Labels.consTuple[MdocApp12.this.A, EmptyTuple.type]( +// this.A.derived$Label, dimwit.tensor.Labels.emptyTuple), // /* missing */ // summon[dimwit.tensor.TensorOps.IsFloating[dimwit.tensor.DType.Int32]] // ) @@ -1121,8 +1121,8 @@ val wrong = boolTensor.mean // // dimwit.mean[Tuple1[MdocApp12.this.A], dimwit.tensor.DType.Bool](this.boolTensor) // ( -// dimwit.tensor.Labels.concat[MdocApp12.this.A, EmptyTuple.type]( -// this.A.derived$Label, dimwit.tensor.Labels.namesOfEmpty), +// dimwit.tensor.Labels.consTuple[MdocApp12.this.A, EmptyTuple.type]( +// this.A.derived$Label, dimwit.tensor.Labels.emptyTuple), // /* missing */ // summon[dimwit.tensor.TensorOps.IsFloating[dimwit.tensor.DType.Bool]] // ) @@ -1158,9 +1158,9 @@ val wrong = m1.dot(Axis[B])(m2) // Axis[B] not in m2 // // dimwit.tensor.ShapeTypeHelpers.AxisRemover.bridge[ // (MdocApp12.this.C, MdocApp12.this.D), MdocApp12.this.B, R]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp12.this.C, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp12.this.C, // MdocApp12.this.D *: EmptyTuple.type, MdocApp12.this.B]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp12.this.D, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp12.this.D, // EmptyTuple.type, MdocApp12.this.B]( // dimwit.tensor.ShapeTypeHelpers.AxisIndex.concatRight[A, B², L]) // ), @@ -1198,17 +1198,17 @@ val wrong = t1 +! t2 // dimwit.tensor.tensorops.TensorOpsUtil.Broadcast.broadcastLeft[ // Tuple1[MdocApp12.this.A], Tuple1[MdocApp12.this.A], // dimwit.tensor.DType.Float32]( -// dimwit.tensor.Labels.concat[MdocApp12.this.A, EmptyTuple.type]( -// this.A.derived$Label, dimwit.tensor.Labels.namesOfEmpty), -// dimwit.tensor.Labels.concat[MdocApp12.this.A, EmptyTuple.type]( -// this.A.derived$Label, dimwit.tensor.Labels.namesOfEmpty), -// dimwit.tensor.TupleHelpers.StrictSubset.derive[Tuple1[MdocApp12.this.A], +// dimwit.tensor.Labels.consTuple[MdocApp12.this.A, EmptyTuple.type]( +// this.A.derived$Label, dimwit.tensor.Labels.emptyTuple), +// dimwit.tensor.Labels.consTuple[MdocApp12.this.A, EmptyTuple.type]( +// this.A.derived$Label, dimwit.tensor.Labels.emptyTuple), +// dimwit.tensor.TupleHelpers.StrictSubset.bridge[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.Subset.consTuple²[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[MdocApp12.this.A]]), +// dimwit.tensor.TupleHelpers.Subset.emptyTuple²[Tuple1[MdocApp12.this.A]]), // /* missing */ // summon[ // scala.util.NotGiven[Tuple1[MdocApp12.this.A] =:= @@ -1217,7 +1217,13 @@ val wrong = t1 +! t2 // ) // ) // -// But no implicit values were found that match type scala.util.NotGiven[Tuple1[MdocApp12.this.A] =:= Tuple1[MdocApp12.this.A]]. +// But no implicit values were found that match type scala.util.NotGiven[Tuple1[MdocApp12.this.A] =:= Tuple1[MdocApp12.this.A]] +// +// where: consTuple is a given instance in object Labels +// consTuple² is a given instance in object Subset +// emptyTuple is a given instance in object Labels +// emptyTuple² is a given instance in object Subset +// . // val wrong = t1 +! t2 // ^^ ``` @@ -1234,9 +1240,9 @@ val wrong = t.sum(Axis[C]) // Axis[C] not in tensor // // dimwit.tensor.ShapeTypeHelpers.AxisRemover.bridge[ // (MdocApp12.this.A, MdocApp12.this.B), MdocApp12.this.C, R]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp12.this.A, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp12.this.A, // MdocApp12.this.B *: EmptyTuple.type, MdocApp12.this.C]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.tail[MdocApp12.this.B, +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp12.this.B, // EmptyTuple.type, MdocApp12.this.C]( // dimwit.tensor.ShapeTypeHelpers.AxisIndex.concatRight[A², B², L]) // ), diff --git a/core/src/main/scala/dimwit/autodiff/Autodiff.scala b/core/src/main/scala/dimwit/autodiff/Autodiff.scala index 43bd7a4..28d7946 100644 --- a/core/src/main/scala/dimwit/autodiff/Autodiff.scala +++ b/core/src/main/scala/dimwit/autodiff/Autodiff.scala @@ -3,32 +3,100 @@ package dimwit.autodiff import dimwit.python.PyIndex.itemAt import dimwit.OnError import dimwit.jax.Jax +import dimwit.prime.PrimeConcat import dimwit.tensor.Tensor import dimwit.tensor.Tensor0 import dimwit.tensor.TensorOps.IsFloating -import dimwit.tensor.TupleHelpers.PrimeConcatType import dimwit.tensortree.TensorTree import me.shadaj.scalapy.py -object Autodiff: - - type Gradient[In, Out] = Out match - case EmptyTuple => EmptyTuple - case h *: t => Gradient[In, h] *: Gradient[In, t] - case Tensor[outS, v] => GradientTensorVsInput[In, outS, v] - case _ => EmptyTuple +import scala.NamedTuple.NamedTuple +import scala.annotation.implicitNotFound +import scala.deriving.Mirror - type GradientTensorVsInput[In, OutShape <: Tuple, V] = In match - case EmptyTuple => EmptyTuple - 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] +object Autodiff: - 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] + /** The derivative of a function `In => Out`: the structure of `Out`, with every + * tensor in it replaced by its derivative with respect to the whole of `In`. + * + * Instances are open, so this stays in step with [[TensorTree]] - a structure + * a user has given a `TensorTree` can be given a `Gradient` too. + */ + @implicitNotFound( + "Cannot express the derivative of ${Out} with respect to ${In}. Both must be built from tensors, tuples, named tuples or case classes with a TensorTree instance" + ) + trait Gradient[In, Out]: + type Result + + object Gradient extends GradientLowPriority: + + type Aux[In, Out, R] = Gradient[In, Out] { type Result = R } + + private[autodiff] def instance[In, Out, R]: Aux[In, Out, R] = + new Gradient[In, Out]: + type Result = R + + /** An output tensor is differentiated against every tensor in the input. */ + given tensor[In, OutShape <: Tuple, V, R](using + vsInput: GradientTensorVsInput.Aux[In, OutShape, V, R] + ): Aux[In, Tensor[OutShape, V], R] = instance + + given emptyTuple[In]: Aux[In, EmptyTuple, EmptyTuple] = instance + + given consTuple[In, H, HR, T <: Tuple, TR <: Tuple](using + head: Aux[In, H, HR], + tail: Aux[In, T, TR] + ): Aux[In, H *: T, HR *: TR] = instance + + given namedTuple[In, N <: Tuple, Vs <: Tuple, R <: Tuple](using + values: Aux[In, Vs, R] + ): Aux[In, NamedTuple[N, Vs], NamedTuple[N, R]] = instance + + trait GradientLowPriority: + /** A case class output becomes a named tuple of its field derivatives, + * keeping the field names. Lower priority than the tuple cases, since + * tuples are Products too. + */ + given product[In, P <: Product, Names <: Tuple, Elems <: Tuple, R <: Tuple](using + m: Mirror.ProductOf[P] { type MirroredElemLabels = Names; type MirroredElemTypes = Elems }, + elems: Gradient.Aux[In, Elems, R] + ): Gradient.Aux[In, P, NamedTuple[Names, R]] = Gradient.instance + + /** The derivative of one output tensor of shape `OutShape` with respect to the + * whole input structure `In`. Mirrors [[Gradient]], recursing on the input. + */ + trait GradientTensorVsInput[In, OutShape <: Tuple, V]: + type Result + + object GradientTensorVsInput extends GradientTensorVsInputLowPriority: + + type Aux[In, OutShape <: Tuple, V, R] = GradientTensorVsInput[In, OutShape, V] { type Result = R } + + private[autodiff] def instance[In, OutShape <: Tuple, V, R]: Aux[In, OutShape, V, R] = + new GradientTensorVsInput[In, OutShape, V]: + type Result = R + + /** Output axes first, then the input axes, primed where they collide. */ + given tensor[InShape <: Tuple, InV, OutShape <: Tuple, V, O <: Tuple](using + concat: PrimeConcat.Aux[OutShape, InShape, O] + ): Aux[Tensor[InShape, InV], OutShape, V, Tensor[O, V]] = instance + + given emptyTuple[OutShape <: Tuple, V]: Aux[EmptyTuple, OutShape, V, EmptyTuple] = instance + + given consTuple[H, HR, T <: Tuple, TR <: Tuple, OutShape <: Tuple, V](using + head: Aux[H, OutShape, V, HR], + tail: Aux[T, OutShape, V, TR] + ): Aux[H *: T, OutShape, V, HR *: TR] = instance + + given namedTuple[N <: Tuple, Vs <: Tuple, R <: Tuple, OutShape <: Tuple, V](using + values: Aux[Vs, OutShape, V, R] + ): Aux[NamedTuple[N, Vs], OutShape, V, NamedTuple[N, R]] = instance + + trait GradientTensorVsInputLowPriority: + given product[P <: Product, Names <: Tuple, Elems <: Tuple, R <: Tuple, OutShape <: Tuple, V](using + m: Mirror.ProductOf[P] { type MirroredElemLabels = Names; type MirroredElemTypes = Elems }, + elems: GradientTensorVsInput.Aux[Elems, OutShape, V, R] + ): GradientTensorVsInput.Aux[P, OutShape, V, NamedTuple[Names, R]] = GradientTensorVsInput.instance // 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)) @@ -76,8 +144,9 @@ object Autodiff: def jacobian[In, Out](f: In => Out)(using inTree: TensorTree[In], outTree: TensorTree[Out], - gradTree: TensorTree[Gradient[In, Out]] - ): In => Gradient[In, Out] = + gradient: Gradient[In, Out], + gradTree: TensorTree[gradient.Result] + ): In => gradient.Result = val fpy = (jxpr: py.Dynamic) => OnError.traceStack: @@ -94,8 +163,9 @@ object Autodiff: def jacRev[In, Out](f: In => Out)(using inTree: TensorTree[In], outTree: TensorTree[Out], - gradTree: TensorTree[Gradient[In, Out]] - ): In => Gradient[In, Out] = + gradient: Gradient[In, Out], + gradTree: TensorTree[gradient.Result] + ): In => gradient.Result = val fpy = (jxpr: py.Dynamic) => OnError.traceStack: outTree.toPyTree(f(inTree.fromPyTree(jxpr))) @@ -105,8 +175,9 @@ object Autodiff: def jacFwd[In, Out](f: In => Out)(using inTree: TensorTree[In], outTree: TensorTree[Out], - gradTree: TensorTree[Gradient[In, Out]] - ): In => Gradient[In, Out] = + gradient: Gradient[In, Out], + gradTree: TensorTree[gradient.Result] + ): In => gradient.Result = val fpy = (jxpr: py.Dynamic) => OnError.traceStack: outTree.toPyTree(f(inTree.fromPyTree(jxpr))) @@ -116,8 +187,9 @@ object Autodiff: def hessian[In, V: IsFloating](f: In => Tensor0[V])(using inTree: TensorTree[In], outTree: TensorTree[Tensor0[V]], - hessTree: TensorTree[Hessian[In]] - ): In => Hessian[In] = + hess: Gradient[In, In], + hessTree: TensorTree[hess.Result] + ): In => hess.Result = val fpy = (jxpr: py.Dynamic) => OnError.traceStack: val x = inTree.fromPyTree(jxpr) diff --git a/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala b/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala index f7334cb..1afc629 100644 --- a/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala +++ b/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala @@ -39,7 +39,7 @@ object LinearAlgebra: * @param t The input tensor from which to compute the determinant. * @return The determinant of the input tensor */ - def det[LRow: Label, LCol: Label, V: IsFloating](t: Tensor[(LRow, LCol), V]): Tensor0[V] = + def det[LRow: Label, LCol: Label, V: IsFloating](t: Tensor2[LRow, LCol, V]): Tensor0[V] = Tensor(Jax.jnp.linalg.det(t.jaxValue)) /** Extracts the diagonal, with an optional offset, diff --git a/core/src/main/scala/dimwit/package.scala b/core/src/main/scala/dimwit/package.scala index 67191f7..5999c35 100644 --- a/core/src/main/scala/dimwit/package.scala +++ b/core/src/main/scala/dimwit/package.scala @@ -15,24 +15,6 @@ package object dimwit: object StringLabelMath: infix type *[A <: String, B <: String] = A + "*" + B - trait Prime[T] - object Prime: - given [L](using label: Label[L]): Label[Prime[L]] with - val name: String = s"${label.name}'" - - type RemovePrimes[T <: Tuple] <: Tuple = T match - case EmptyTuple => EmptyTuple - case Prime[l] *: tail => l *: RemovePrimes[tail] - case h *: tail => h *: RemovePrimes[tail] - - extension [T <: Tuple: Labels, V](tensor: Tensor[T, V]) - def dropPrimes: Tensor[RemovePrimes[T], V] = - given newLabels: Labels[RemovePrimes[T]] with - val names: List[String] = - val oldLabels = summon[Labels[T]] - oldLabels.names.toList.map(_.replace("'", "")) - Tensor[RemovePrimes[T], V](tensor.jaxValue) - def gc(): Unit = System.gc() Jax.gc() @@ -58,7 +40,7 @@ package object dimwit: val name: String = s"${labelA.name}+${labelB.name}" // Export tensor and related types - export dimwit.tensor.{Tensor, Tensor0, Tensor1, Tensor2, Tensor3, TypedIndex} + export dimwit.tensor.{Tensor, Tensor0, Tensor1, Tensor2, Tensor3, Tensor4, TypedIndex} export dimwit.tensor.{Shape, Shape0, Shape1, Shape2, Shape3} export dimwit.tensor.DType export dimwit.tensor.DType.{BFloat16, Float16, Float32, Float64, Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, Bool} @@ -77,6 +59,9 @@ package object dimwit: } export dimwit.tensor.ShapeTypeHelpers.{AxisInTensor, AxisIndex, AxisRemover, AxisReplacer, AxisIndices, AxesRemover, AxesConditionalRemover, SharedAxisRemover} + // Export the Prime axis marker and the type classes that manipulate it + export dimwit.prime.{Prime, PrimeRemover, PrimeRest, PrimeConcat} + // Export operations export dimwit.tensor.TensorOps.* export dimwit.tensor.ValueOps.* diff --git a/core/src/main/scala/dimwit/prime/Prime.scala b/core/src/main/scala/dimwit/prime/Prime.scala new file mode 100644 index 0000000..d2bd7af --- /dev/null +++ b/core/src/main/scala/dimwit/prime/Prime.scala @@ -0,0 +1,92 @@ +package dimwit.prime + +import dimwit.tensor.Label +import dimwit.tensor.Labels +import dimwit.tensor.Tensor +import dimwit.tensor.TupleHelpers.SetMember + +/** Marks an axis label as a distinct copy of `T`. + * + * Operations that would otherwise produce the same axis twice - the two sides of + * an outer product, or an output axis meeting the input axis it is differentiated + * against - prime the second occurrence so that the two stay distinguishable. + * `Label[Prime[T]]` renders as the label of `T` followed by a `'`. + */ +trait Prime[T] + +object Prime: + + given [L](using label: Label[L]): Label[Prime[L]] with + val name: String = s"${label.name}'" + + extension [T <: Tuple: Labels, V](tensor: Tensor[T, V]) + /** Drops one `Prime` wrapper from every primed axis of the shape. */ + def dropPrimes(using remover: PrimeRemover[T]): Tensor[remover.Out, V] = + given droppedLabels: Labels[remover.Out] with + val names: List[String] = summon[Labels[T]].names.map(_.stripSuffix("'")) + Tensor[remover.Out, V](tensor.jaxValue) + +/** Removes one `Prime` wrapper from every primed axis of a shape. */ +trait PrimeRemover[T <: Tuple]: + type Out <: Tuple + +object PrimeRemover extends PrimeRemoverLowPriority: + type Aux[T <: Tuple, O <: Tuple] = PrimeRemover[T] { type Out = O } + + private[prime] def instance[T <: Tuple, O <: Tuple]: Aux[T, O] = + new PrimeRemover[T]: + type Out = O + + given emptyTuple: Aux[EmptyTuple, EmptyTuple] = instance + + /** A primed head loses its wrapper. */ + given primedTuple[L, T <: Tuple, O <: Tuple](using tail: Aux[T, O]): Aux[Prime[L] *: T, L *: O] = instance + +trait PrimeRemoverLowPriority: + /** Any other head is carried over unchanged. */ + given plainTuple[H, T <: Tuple, O <: Tuple](using tail: PrimeRemover.Aux[T, O]): PrimeRemover.Aux[H *: T, H *: O] = + PrimeRemover.instance + +/** Primes every axis of `Incoming` that already occurs in `Fixed`, so that the two + * shapes can be put side by side without an axis appearing twice. + */ +trait PrimeRest[Fixed <: Tuple, Incoming <: Tuple]: + type Out <: Tuple + +object PrimeRest extends PrimeRestLowPriority: + type Aux[Fixed <: Tuple, Incoming <: Tuple, O <: Tuple] = + PrimeRest[Fixed, Incoming] { type Out = O } + + private[prime] def instance[Fixed <: Tuple, Incoming <: Tuple, O <: Tuple]: Aux[Fixed, Incoming, O] = + new PrimeRest[Fixed, Incoming]: + type Out = O + + given emptyTuple[Fixed <: Tuple]: Aux[Fixed, EmptyTuple, EmptyTuple] = instance + + /** The head collides with an axis of `Fixed`, so it is primed. */ + given collidingTuple[Fixed <: Tuple, H, T <: Tuple, TailOut <: Tuple](using + member: SetMember[H, Fixed], + tail: Aux[Fixed, T, TailOut] + ): Aux[Fixed, H *: T, Prime[H] *: TailOut] = instance + +trait PrimeRestLowPriority: + /** Membership could not be proven - the head is an abstract type parameter, say - + * so assume it does not collide and leave it unprimed. + */ + given distinctTuple[Fixed <: Tuple, H, T <: Tuple, TailOut <: Tuple](using + tail: PrimeRest.Aux[Fixed, T, TailOut] + ): PrimeRest.Aux[Fixed, H *: T, H *: TailOut] = PrimeRest.instance + +/** Puts `R2` after `R1`, priming the axes of `R2` that already occur in `R1`. */ +trait PrimeConcat[R1 <: Tuple, R2 <: Tuple]: + type Out <: Tuple + +object PrimeConcat: + type Aux[R1 <: Tuple, R2 <: Tuple, O <: Tuple] = + PrimeConcat[R1, R2] { type Out = O } + + given bridge[R1 <: Tuple, R2 <: Tuple, Suffix <: Tuple](using + rest: PrimeRest.Aux[R1, R2, Suffix] + ): PrimeConcat.Aux[R1, R2, Tuple.Concat[R1, Suffix]] = + new PrimeConcat[R1, R2]: + type Out = Tuple.Concat[R1, Suffix] diff --git a/core/src/main/scala/dimwit/tensor/DType.scala b/core/src/main/scala/dimwit/tensor/DType.scala index d7a65f7..a59c39f 100644 --- a/core/src/main/scala/dimwit/tensor/DType.scala +++ b/core/src/main/scala/dimwit/tensor/DType.scala @@ -12,7 +12,7 @@ import java.nio.ByteOrder object DType: sealed trait UInt8 - given uint8IsFloating: IsInteger[UInt8] with + given uint8IsInteger: IsInteger[UInt8] with def dtype: DType = DType.UInt8 sealed trait UInt16 diff --git a/core/src/main/scala/dimwit/tensor/Labels.scala b/core/src/main/scala/dimwit/tensor/Labels.scala index d565c3c..c934859 100644 --- a/core/src/main/scala/dimwit/tensor/Labels.scala +++ b/core/src/main/scala/dimwit/tensor/Labels.scala @@ -2,8 +2,6 @@ package dimwit.tensor import scala.quoted.* -import Tuple.:* - @scala.annotation.implicitNotFound(""" An axis label ${T} was given or inferred, which does not have a Label instance. Ensure that all axis types ${T} are defined with 'derives Label' (e.g. 'trait T derives Label') @@ -35,30 +33,15 @@ private class LabelsImpl[T](val names: List[String]) extends Labels[T] object Labels extends LabelsLowPriority: - given namesOfEmpty: Labels[EmptyTuple] = new LabelsImpl[EmptyTuple](Nil) + given emptyTuple: Labels[EmptyTuple] = new LabelsImpl[EmptyTuple](Nil) given lift[A](using v: Label[A]): Labels[A] = new LabelsImpl[A](List(v.name)) - given [A, B](using a: Labels[A], b: Labels[B]): Labels[(A, B)] = new LabelsImpl[(A, B)](a.names ++ b.names) - given [A, B, C](using a: Labels[A], b: Labels[B], c: Labels[C]): Labels[(A, B, C)] = new LabelsImpl[(A, B, C)](a.names ++ b.names ++ c.names) - given [A, B, C, D](using a: Labels[A], b: Labels[B], c: Labels[C], d: Labels[D]): Labels[(A, B, C, D)] = new LabelsImpl[(A, B, C, D)](a.names ++ b.names ++ c.names ++ d.names) - given [A, B, C, D, E](using a: Labels[A], b: Labels[B], c: Labels[C], d: Labels[D], e: Labels[E]): Labels[(A, B, C, D, E)] = new LabelsImpl[(A, B, C, D, E)](a.names ++ b.names ++ c.names ++ d.names ++ e.names) - given [A, B, C, D, E, F](using a: Labels[A], b: Labels[B], c: Labels[C], d: Labels[D], e: Labels[E], f: Labels[F]): Labels[(A, B, C, D, E, F)] = new LabelsImpl[(A, B, C, D, E, F)](a.names ++ b.names ++ c.names ++ d.names ++ e.names ++ f.names) - - given concat[head, tail <: Tuple](using - v: Label[head], - t: Labels[tail] - ): Labels[head *: tail] = new LabelsImpl[head *: tail]( - v.name :: t.names - ) - - given append[head, tail <: Tuple](using - v: Label[head], - t: Labels[tail] - ): Labels[tail :* head] = new LabelsImpl[tail :* head]( - t.names :+ v.name - ) + given consTuple[H, T <: Tuple](using + head: Label[H], + tail: Labels[T] + ): Labels[H *: T] = new LabelsImpl[H *: T](head.name :: tail.names) private trait LabelsLowPriority: - given [T1 <: Tuple, T2 <: Tuple](using n1: Labels[T1], n2: Labels[T2]): Labels[Tuple.Concat[T1, T2]] = + given concatTuple[T1 <: Tuple, T2 <: Tuple](using n1: Labels[T1], n2: Labels[T2]): Labels[Tuple.Concat[T1, T2]] = new LabelsImpl(n1.names ++ n2.names) diff --git a/core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala b/core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala index 4192921..37f54eb 100644 --- a/core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala +++ b/core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala @@ -38,10 +38,10 @@ object ShapeTypeHelpers: def apply[T <: Tuple, L](using idx: AxisIndex[T, L]): Int = idx.index - given head[L, Tail <: Tuple]: AxisIndex[L *: Tail, L] with + given found[L, Tail <: Tuple]: AxisIndex[L *: Tail, L] with val index = 0 - given tail[H, T <: Tuple, L](using + given search[H, T <: Tuple, L](using next: AxisIndex[T, L] ): AxisIndex[H *: T, L] with val index = 1 + next.index @@ -108,7 +108,7 @@ object ShapeTypeHelpers: type Out = Tuple.Concat[R, T] def index = 0 - given recurse[H, T <: Tuple, A, R <: Tuple, TailOut <: Tuple](using + given search[H, T <: Tuple, A, R <: Tuple, TailOut <: Tuple](using ne: NotGiven[H =:= A], tailSplice: Splice.Aux[T, A, R, TailOut] ): Splice.Aux[H *: T, A, R, H *: TailOut] = @@ -140,7 +140,7 @@ object ShapeTypeHelpers: case _: (head *: tail) => summonInline[AxisIndex[InTuple, head]].index :: indicesOfList[InTuple, tail] - inline given [T <: Tuple, ToFind <: Tuple]: AxisIndices[T, ToFind] = AxisIndicesImpl[T, ToFind](indicesOfList[T, ToFind]) + inline given indices[T <: Tuple, ToFind <: Tuple]: AxisIndices[T, ToFind] = AxisIndicesImpl[T, ToFind](indicesOfList[T, ToFind]) end AxisIndices @@ -186,12 +186,12 @@ object ShapeTypeHelpers: object SharedAxisRemover: type Aux[S <: Tuple, A, O <: Tuple] = SharedAxisRemover[S, A] { type RemainingAxes = O } - given empty[Axis]: SharedAxisRemover.Aux[EmptyTuple, Axis, EmptyTuple] = new SharedAxisRemover[EmptyTuple, Axis]: + given emptyTuple[Axis]: SharedAxisRemover.Aux[EmptyTuple, Axis, EmptyTuple] = new SharedAxisRemover[EmptyTuple, Axis]: type RemainingAxes = EmptyTuple def indices = Nil def shapesLabels = Nil - given cons[H <: Tuple, T <: Tuple, Axis, R <: Tuple, TailOut <: Tuple](using + given consTuple[H <: Tuple, T <: Tuple, Axis, R <: Tuple, TailOut <: Tuple](using evH: AxisRemover.Aux[H, Axis, R], evT: SharedAxisRemover.Aux[T, Axis, TailOut], rLabels: Labels[R] @@ -206,10 +206,10 @@ object ShapeTypeHelpers: def extract(t: T): Map[String, Int] object DimExtractor: - given DimExtractor[EmptyTuple] with + given emptyTuple: DimExtractor[EmptyTuple] with def extract(t: EmptyTuple) = Map.empty - given [L, Tail <: Tuple](using + given consTuple[L, Tail <: Tuple](using label: Label[L], tailExtractor: DimExtractor[Tail] ): DimExtractor[AxisExtent[L] *: Tail] with @@ -217,10 +217,6 @@ object ShapeTypeHelpers: val size = t.head.size Map(label.name -> size) ++ tailExtractor.extract(t.tail) - given single[L](using label: Label[L]): DimExtractor[AxisExtent[L]] with - def extract(t: AxisExtent[L]) = - Map(label.name -> t.size) - /** Merges multiple axes in a tensor shape into a single axis. * * NewShape is the resulting shape after merging the axes. @@ -241,7 +237,7 @@ object ShapeTypeHelpers: case head *: tail => MergeLabelsRec[tail, Acc |*| head] object MergeLabels: - given [T <: Tuple: Labels]: Label[MergeLabels[T]] with + given mergedLabel[T <: Tuple: Labels]: Label[MergeLabels[T]] with def name = summon[Labels[T]].names.mkString("*") object AxesMerger: diff --git a/core/src/main/scala/dimwit/tensor/Tensor.scala b/core/src/main/scala/dimwit/tensor/Tensor.scala index 567e2d4..027dba5 100644 --- a/core/src/main/scala/dimwit/tensor/Tensor.scala +++ b/core/src/main/scala/dimwit/tensor/Tensor.scala @@ -1,6 +1,6 @@ package dimwit.tensor -import dimwit.Prime +import dimwit.prime.Prime import dimwit.hardware.Device import dimwit.jax.Jax import dimwit.jax.Jax.PyDynamic diff --git a/core/src/main/scala/dimwit/tensor/TensorEvidence.scala b/core/src/main/scala/dimwit/tensor/TensorEvidence.scala new file mode 100644 index 0000000..59f652d --- /dev/null +++ b/core/src/main/scala/dimwit/tensor/TensorEvidence.scala @@ -0,0 +1,119 @@ +package dimwit.tensor + +import dimwit.tensor.TupleHelpers.Remover +import dimwit.tensor.TupleHelpers.SetMember +import dimwit.|*| + +import scala.annotation.implicitNotFound +import scala.quoted.Expr +import scala.quoted.Quotes +import scala.quoted.Type +import scala.util.NotGiven + +/** Compile time checks relating a shape the user asked for to the shape a tensor + * actually has: is the new order a permutation of the old one, and can every axis + * of the new order be formed from the source shape? + * + * A missing axis is reported by [[ComputeMissing]] as a type rather than as a + * failed search, so that [[CheckValid]] can turn it into a readable error. + */ +object TensorEvidence: + + /** Can axis `A` be formed from source shape `S`? + * + * `Ignore` holds the axes whose size the caller supplied explicitly, which + * therefore need no counterpart in `S`. + */ + trait CanForm[A, S <: Tuple, Ignore <: Tuple] + + object CanForm: + given inSource[A, S <: Tuple, I <: Tuple](using + SetMember[A, S] + ): CanForm[A, S, I] = new CanForm[A, S, I] {} + + given inIgnore[A, S <: Tuple, I <: Tuple](using + NotGiven[SetMember[A, S]], + SetMember[A, I] + ): CanForm[A, S, I] = new CanForm[A, S, I] {} + + @implicitNotFound("The shape ${A} is not a valid permutation of ${B}.") + trait IsPermutation[A <: Tuple, B <: Tuple] + + object IsPermutation: + given emptyTuple: IsPermutation[EmptyTuple, EmptyTuple] with {} + + given consTuple[H, T <: Tuple, B <: Tuple, RemainingB <: Tuple](using + remover: Remover.Aux[B, H, RemainingB], + tail: IsPermutation[T, RemainingB] + ): IsPermutation[H *: T, B] with {} + + /** The outcome of [[ComputeMissing]]: either every axis can be formed, or the + * first one that cannot, kept in the type so it can be named in the error. + */ + sealed trait ValidationResult + final class AllOk extends ValidationResult + final class MissingAxis[A, InT <: Tuple] extends ValidationResult + + /** Walks the `Target` axes and reports the first one that cannot be formed from + * `Source`, in the type parameter `Res`. + */ + trait ComputeMissing[Target <: Tuple, Source <: Tuple, Ignore <: Tuple, Res <: ValidationResult] + + object ComputeMissing extends ComputeMissingLowPriority: + + given emptyTuple[S <: Tuple, I <: Tuple]: ComputeMissing[EmptyTuple, S, I, AllOk] = + new ComputeMissing[EmptyTuple, S, I, AllOk] {} + + /** The head can be formed, so the result is whatever the tail reports. */ + given headFound[H, T <: Tuple, S <: Tuple, I <: Tuple, Res <: ValidationResult](using + found: CanForm[H, S, I], + tail: ComputeMissing[T, S, I, Res] + ): ComputeMissing[H *: T, S, I, Res] = + new ComputeMissing[H *: T, S, I, Res] {} + + /** A composite head that could not be formed as a unit is split into its two + * components, which are then checked on their own. More specific than + * [[ComputeMissingLowPriority.headMissing]], so it is tried first. + */ + given headDecomposed[L, R, T <: Tuple, S <: Tuple, I <: Tuple, Res <: ValidationResult](using + notAUnit: NotGiven[CanForm[L |*| R, S, I]], + tail: ComputeMissing[L *: R *: T, S, I, Res] + ): ComputeMissing[(L |*| R) *: T, S, I, Res] = + new ComputeMissing[(L |*| R) *: T, S, I, Res] {} + + trait ComputeMissingLowPriority: + /** The head can be formed neither from the source nor from the explicit sizes: + * stop here and report it. + */ + given headMissing[H, T <: Tuple, S <: Tuple, I <: Tuple](using + missing: NotGiven[CanForm[H, S, I]], + notIgnored: NotGiven[SetMember[H, I]] + ): ComputeMissing[H *: T, S, I, MissingAxis[H, S]] = + new ComputeMissing[H *: T, S, I, MissingAxis[H, S]] {} + + /** Turns the [[ValidationResult]] into either a summonable instance or a + * compile error naming the missing axis. + */ + sealed trait CheckValid[R <: ValidationResult] + + object CheckValid: + given ok: CheckValid[AllOk] = new CheckValid[AllOk] {} + + inline given fail[A, SourceShape <: Tuple]: CheckValid[MissingAxis[A, SourceShape]] = + ${ failImpl[A, SourceShape] } + + def failImpl[A: Type, SourceShape <: Tuple: Type](using Quotes): Expr[CheckValid[MissingAxis[A, SourceShape]]] = + import scala.quoted.quotes.reflect.* + // Type.show gives the readable name (e.g. "A" instead of "package.A") + val name = Type.show[A] + val sourceShape = Type.show[SourceShape] + + report.errorAndAbort( + s"""❌ Missing Axis: '$name' in the source shape $sourceShape. There are a few possible reasons: + | 1. Missing axis $name is not present in the source shape $sourceShape. + | 👉 New structure must be based on source shape + | 2. Missing axis $name is present only in flattened form (e.g., $name|*|OtherAxis) in the source shape $sourceShape. This requires additional information to be unflattened. + | If you are unflattening (e.g. $name|*|OtherAxis -> $name, OtherAxis), you must provide the size of '$name' explicitly. + | 👉 Try: .rearrange(newOrder, (Axis[$name] -> size, ...)), where size is the length of $name after the unflattening. + |""".stripMargin + ) diff --git a/core/src/main/scala/dimwit/tensor/TupleHelpers.scala b/core/src/main/scala/dimwit/tensor/TupleHelpers.scala index f6379d6..ff5b8a9 100644 --- a/core/src/main/scala/dimwit/tensor/TupleHelpers.scala +++ b/core/src/main/scala/dimwit/tensor/TupleHelpers.scala @@ -1,80 +1,85 @@ package dimwit.tensor import scala.compiletime.ops -import scala.quoted.Expr -import scala.quoted.Quotes -import scala.quoted.Type import scala.util.NotGiven -/* Helpers for manipulating Tuple types */ +/* Helpers for manipulating Tuple types. */ object TupleHelpers: + /** Evidence that `S` is a subset of `T` but not all of it. */ trait StrictSubset[S <: Tuple, T <: Tuple] object StrictSubset: - given derive[S <: Tuple, T <: Tuple](using - ev: Subset[S, T], + given bridge[S <: Tuple, T <: Tuple](using + subset: Subset[S, T], notEq: NotGiven[S =:= T] ): StrictSubset[S, T] with {} + /** Evidence that every element of `S` also occurs in `T`. */ trait Subset[S <: Tuple, T <: Tuple] object Subset: - given empty[T <: Tuple]: Subset[EmptyTuple, T] with {} + given emptyTuple[T <: Tuple]: Subset[EmptyTuple, T] with {} - given head[H, STail <: Tuple, T <: Tuple](using - evH: SetMember[H, T], - evT: Subset[STail, T] + given consTuple[H, STail <: Tuple, T <: Tuple](using + head: SetMember[H, T], + tail: Subset[STail, T] ): Subset[H *: STail, T] with {} + /** Evidence that `K` occurs in `T`. */ trait SetMember[K, T <: Tuple] + object SetMember: given found[K, T <: Tuple]: SetMember[K, K *: T] with {} - given search[K, H, T <: Tuple](using ev: SetMember[K, T]): SetMember[K, H *: T] with {} + given search[K, H, T <: Tuple](using tail: SetMember[K, T]): SetMember[K, H *: T] with {} + /** Removes the first occurrence of `ToRemoveElement` from `T`. */ type Remover[T <: Tuple, ToRemoveElement] = RemoverAll[T, ToRemoveElement *: EmptyTuple] object Remover: type Aux[T <: Tuple, ToRemoveElement, O <: Tuple] = RemoverAll.Aux[T, ToRemoveElement *: EmptyTuple, O] + /** Removes the first occurrence of every element of `ToRemove` from `T`. + * + * The instances recurse on `ToRemove` first (`noKeys`, `multipleKeys`) and then + * search `T` for the one remaining key (`singleKeyFound`, `singleKeySearch`). + */ trait RemoverAll[T <: Tuple, ToRemove <: Tuple]: type Out <: Tuple - object RemoverAll extends LowPriorityRemoverAll: + object RemoverAll extends RemoverAllLowPriority: - // 0. The Aux type alias forces the compiler to resolve 'O' explicitly + /** The `Aux` alias forces the compiler to resolve `O` explicitly. */ type Aux[T <: Tuple, ToRemove <: Tuple, O <: Tuple] = RemoverAll[T, ToRemove] { type Out = O } - // 1. Base Case: Empty keys -> Return input as is - given emptyKeys[T <: Tuple]: Aux[T, EmptyTuple, T] = - new RemoverAll[T, EmptyTuple]: - type Out = T - - // 2. Chain Case: Process K1, then K2... - // We use Aux to capture 'Inter' and 'O' explicitly - given chain[T <: Tuple, K1, K2, Rest <: Tuple, Inter <: Tuple, O <: Tuple](using - r1: Aux[T, K1 *: EmptyTuple, Inter], - r2: Aux[Inter, K2 *: Rest, O] - ): Aux[T, K1 *: K2 *: Rest, O] = - new RemoverAll[T, K1 *: K2 *: Rest]: + private[tensor] def instance[T <: Tuple, ToRemove <: Tuple, O <: Tuple]: Aux[T, ToRemove, O] = + new RemoverAll[T, ToRemove]: type Out = O - // 3. Found Case: H is a subtype of K - // We explicitly return 'Tail' as the output - given singleFound[K, Tail <: Tuple]: Aux[K *: Tail, K *: EmptyTuple, Tail] = - new RemoverAll[K *: Tail, K *: EmptyTuple]: - type Out = Tail - - trait LowPriorityRemoverAll: - // 4. Search Case: Recurse - // We capture 'TailOut' as a type parameter to ensure it is fully resolved - given singleSearch[H, Tail <: Tuple, K, TailOut <: Tuple](using - next: RemoverAll.Aux[Tail, K *: EmptyTuple, TailOut] - ): RemoverAll.Aux[H *: Tail, K *: EmptyTuple, H *: TailOut] = - new RemoverAll[H *: Tail, K *: EmptyTuple]: - type Out = H *: TailOut - + /** Nothing left to remove. */ + given noKeys[T <: Tuple]: Aux[T, EmptyTuple, T] = instance + + /** Remove `K1`, then carry on with the rest. `Inter` names the intermediate + * tuple so that it, and `O`, are resolved explicitly. + */ + given multipleKeys[T <: Tuple, K1, K2, Rest <: Tuple, Inter <: Tuple, O <: Tuple](using + first: Aux[T, K1 *: EmptyTuple, Inter], + rest: Aux[Inter, K2 *: Rest, O] + ): Aux[T, K1 *: K2 *: Rest, O] = instance + + /** The single remaining key sits at the head, so drop it. */ + given singleKeyFound[K, Tail <: Tuple]: Aux[K *: Tail, K *: EmptyTuple, Tail] = instance + + trait RemoverAllLowPriority: + /** Keep the head and look for the key in the tail. `TailOut` is a type + * parameter so that it is fully resolved. + */ + given singleKeySearch[H, Tail <: Tuple, K, TailOut <: Tuple](using + tail: RemoverAll.Aux[Tail, K *: EmptyTuple, TailOut] + ): RemoverAll.Aux[H *: Tail, K *: EmptyTuple, H *: TailOut] = RemoverAll.instance + + /** Replaces the first occurrence of `Target` in `T` with `Replacement`. */ trait Replacer[T <: Tuple, Target, Replacement]: type Out <: Tuple @@ -82,189 +87,18 @@ object TupleHelpers: type Aux[T <: Tuple, Target, Replacement, O <: Tuple] = Replacer[T, Target, Replacement] { type Out = O } - given found[Target, Tail <: Tuple, Replacement]: Replacer[Target *: Tail, Target, Replacement] with - type Out = Replacement *: Tail + private[tensor] def instance[T <: Tuple, Target, Replacement, O <: Tuple]: Aux[T, Target, Replacement, O] = + new Replacer[T, Target, Replacement]: + type Out = O + + given found[Target, Tail <: Tuple, Replacement]: Aux[Target *: Tail, Target, Replacement, Replacement *: Tail] = instance trait ReplacerLowPriority: - given recurse[Head, Tail <: Tuple, Target, Replacement, TailOut <: Tuple](using - next: Replacer.Aux[Tail, Target, Replacement, TailOut] - ): Replacer[Head *: Tail, Target, Replacement] with - type Out = Head *: TailOut + given search[Head, Tail <: Tuple, Target, Replacement, TailOut <: Tuple](using + tail: Replacer.Aux[Tail, Target, Replacement, TailOut] + ): Replacer.Aux[Head *: Tail, Target, Replacement, Head *: TailOut] = Replacer.instance + /** A tuple of `N` elements, all of type `T`. */ type TupleNOf[N <: Int, T] <: Tuple = N match case 0 => EmptyTuple case _ => T *: TupleNOf[ops.int.-[N, 1], T] - - import dimwit.Prime - - type Member[X, T <: Tuple] <: Boolean = T match - case EmptyTuple => false - case X *: t => true - case _ *: t => Member[X, t] - - import dimwit.|*| - import scala.util.NotGiven - import scala.annotation.implicitNotFound - - object TensorEvidence: - - // --- Core Checks (Same as before) --- - - // 1. Does Source T contain Axis X? - trait Has[X, T <: Tuple] - object Has: - given head[X, T <: Tuple]: Has[X, X *: T] = new Has[X, X *: T] {} - given tail[X, H, T <: Tuple](using Has[X, T]): Has[X, H *: T] = new Has[X, H *: T] {} - - // 2. Can we form Axis A from Source S? (Handles A vs A|*|B) - trait CanForm[A, S <: Tuple, Ignore <: Tuple] - object CanForm: - // Case 1: Found directly in Source (Highest Priority) - given inSource[A, S <: Tuple, I <: Tuple](using - Has[A, S] - ): CanForm[A, S, I] = new CanForm[A, S, I] {} - - // Case 2: Found directly in Ignore List (Explicit Dims) - given inIgnore[A, S <: Tuple, I <: Tuple](using - NotGiven[Has[A, S]], - Has[A, I] - ): CanForm[A, S, I] = new CanForm[A, S, I] {} - - @implicitNotFound("The shape ${A} is not a valid permutation of ${B}.") - trait IsPermutation[A <: Tuple, B <: Tuple] - object IsPermutation: - given base: IsPermutation[EmptyTuple, EmptyTuple] with {} - - given recurse[H, T <: Tuple, B <: Tuple, RemainingB <: Tuple](using - evRem: Remover.Aux[B, H, RemainingB], - evNext: IsPermutation[T, RemainingB] - ): IsPermutation[H *: T, B] with {} - - // Result Types - sealed trait ValidationResult - final class AllOk extends ValidationResult - final class MissingAxis[A, InT <: Tuple] extends ValidationResult - - // 3. ComputeMissing: Walks through Target axes and finds the first missing one. - // It returns the result in the type parameter 'Res'. - trait ComputeMissing[Target <: Tuple, Source <: Tuple, Ignore <: Tuple, Res <: ValidationResult] - - object ComputeMissing extends ComputeMissingLowPriority: - - // Case 1: Target is empty -> All Good! - given empty[S <: Tuple, I <: Tuple]: ComputeMissing[EmptyTuple, S, I, AllOk] = - new ComputeMissing[EmptyTuple, S, I, AllOk] {} - - // Case 2: Head is Valid -> Continue checking Tail - given headFound[H, T <: Tuple, S <: Tuple, I <: Tuple, Res <: ValidationResult](using - found: CanForm[H, S, I], // Proof that Head exists - tailCheck: ComputeMissing[T, S, I, Res] // Recurse - ): ComputeMissing[H *: T, S, I, Res] = - new ComputeMissing[H *: T, S, I, Res] {} - - // Case 3: Composite Decompose (Target Head is L |*| R, and it was NOT found above) - // Strategy: Replace (L |*| R) with L, then R, in the search queue. - // Specificity: This matches (L |*| R) *: T, which is more specific than H *: T. - given decompose[L, R, T <: Tuple, S <: Tuple, I <: Tuple, Res <: ValidationResult](using - missingAsUnit: NotGiven[CanForm[L |*| R, S, I]], // Ensure we didn't miss the unit above - recurse: ComputeMissing[L *: R *: T, S, I, Res] // Flatten L and R into the stream - ): ComputeMissing[(L |*| R) *: T, S, I, Res] = - new ComputeMissing[(L |*| R) *: T, S, I, Res] {} - - trait ComputeMissingLowPriority: - // Case 4: Head is MISSING -> Stop and Report Error - // We use NotGiven to prove it's missing. This handles the 'else' branch safely. - given headMissing[H, T <: Tuple, S <: Tuple, I <: Tuple](using - missing: NotGiven[CanForm[H, S, I]], // Proof that Head is missing - notIgnored: NotGiven[Has[H, I]] - ): ComputeMissing[H *: T, S, I, MissingAxis[H, S]] = - new ComputeMissing[H *: T, S, I, MissingAxis[H, S]] {} - - // --- The Guard (Error Trigger) --- - - // 4. CheckValid: This checks the RESULT of the computation. - // If the result is AllOk, it compiles. - // If the result is MissingAxis[A], it fails with your message. - sealed trait CheckValid[R <: ValidationResult] - - object CheckValid: - // Case 1: Success. We provide an instance, so compilation proceeds. - given ok: CheckValid[AllOk] = new CheckValid[AllOk] {} - - def failImpl[A: Type, SourceShape <: Tuple: Type](using Quotes): Expr[CheckValid[MissingAxis[A, SourceShape]]] = - import scala.quoted.quotes.reflect.* - // Type.show[A] gives you the nice, readable name (e.g., "A" instead of "package.A") - val name = Type.show[A] - val sourceShape = Type.show[SourceShape] - - report.errorAndAbort( - s"""❌ Missing Axis: '$name' in the source shape $sourceShape. There are a few possible reasons: - | 1. Missing axis $name is not present in the source shape $sourceShape. - | 👉 New structure must be based on source shape - | 2. Missing axis $name is present only in flattened form (e.g., $name|*|OtherAxis) in the source shape $sourceShape. This requires additional information to be unflattened. - | If you are unflattening (e.g. $name|*|OtherAxis -> $name, OtherAxis), you must provide the size of '$name' explicitly. - | 👉 Try: .rearrange(newOrder, (Axis[$name] -> size, ...)), where size is the length of $name after the unflattening. - |""".stripMargin - ) - - // Case 2: Failure. We provide an instance that triggers a compile-time ERROR with a user-friendly message. - inline given fail[A, SourceShape <: Tuple]: CheckValid[MissingAxis[A, SourceShape]] = - ${ failImpl[A, SourceShape] } - export TensorEvidence.* - - trait PrimeRest[Fixed <: Tuple, Incoming <: Tuple]: - type Out <: Tuple - - trait PrimeRestLowPriority: - /** If nothing found, we can't proof Member (e.g. for generics), just assume they are different. */ - given assumeAbsent[Fixed <: Tuple, H, T <: Tuple, TailOut <: Tuple](using - tail: PrimeRest.Aux[Fixed, T, TailOut] - ): PrimeRest.Aux[Fixed, H *: T, H *: TailOut] = - new PrimeRest[Fixed, H *: T]: - type Out = H *: TailOut - - object PrimeRest extends PrimeRestLowPriority: - type Aux[Fixed <: Tuple, Incoming <: Tuple, O <: Tuple] = - PrimeRest[Fixed, Incoming] { type Out = O } - - given empty[Fixed <: Tuple]: PrimeRest.Aux[Fixed, EmptyTuple, EmptyTuple] = - new PrimeRest[Fixed, EmptyTuple]: - type Out = EmptyTuple - - given present[Fixed <: Tuple, H, T <: Tuple, TailOut <: Tuple](using - ev: Member[H, Fixed] =:= true, - tail: PrimeRest.Aux[Fixed, T, TailOut] - ): PrimeRest.Aux[Fixed, H *: T, Prime[H] *: TailOut] = - new PrimeRest[Fixed, H *: T]: - type Out = Prime[H] *: TailOut - - given absent[Fixed <: Tuple, H, T <: Tuple, TailOut <: Tuple](using - ev: Member[H, Fixed] =:= false, - tail: PrimeRest.Aux[Fixed, T, TailOut] - ): PrimeRest.Aux[Fixed, H *: T, H *: TailOut] = - new PrimeRest[Fixed, H *: T]: - type Out = H *: TailOut - - trait PrimeConcat[R1 <: Tuple, R2 <: Tuple]: - type Out <: Tuple - - object PrimeConcat: - type Aux[R1 <: Tuple, R2 <: Tuple, O <: Tuple] = - PrimeConcat[R1, R2] { type Out = O } - - given [R1 <: Tuple, R2 <: Tuple, Suffix <: Tuple](using - rest: PrimeRest.Aux[R1, R2, Suffix] - ): PrimeConcat.Aux[R1, R2, Tuple.Concat[R1, Suffix]] = - new PrimeConcat[R1, R2]: - type Out = Tuple.Concat[R1, Suffix] - - // Match type versions for compile-time reduction - // PrimeRest: Transform tuple elements, adding Prime[H] if H is in Fixed - type PrimeRestType[Fixed <: Tuple, Incoming <: Tuple] <: Tuple = Incoming match - case EmptyTuple => EmptyTuple - case h *: t => Member[h, Fixed] match - case true => Prime[h] *: PrimeRestType[Fixed, t] - case false => h *: PrimeRestType[Fixed, t] - - // PrimeConcat: Concatenate R1 with primed version of R2 - type PrimeConcatType[R1 <: Tuple, R2 <: Tuple] = Tuple.Concat[R1, PrimeRestType[R1, R2]] diff --git a/core/src/main/scala/dimwit/tensor/tensorops/ContractionOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/ContractionOps.scala index f29ff4d..33bce6e 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/ContractionOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/ContractionOps.scala @@ -1,11 +1,11 @@ package dimwit.tensor.tensorops import dimwit.jax.Jax +import dimwit.prime.PrimeConcat import dimwit.tensor.Axis import dimwit.tensor.Labels import dimwit.tensor.ShapeTypeHelpers.AxisRemover import dimwit.tensor.Tensor -import dimwit.tensor.TupleHelpers.PrimeConcat import me.shadaj.scalapy.py import me.shadaj.scalapy.py.SeqConverters import me.shadaj.scalapy.readwrite.Writer @@ -56,13 +56,6 @@ object ContractionOps: * * @param axis The pair of axes along which to contract. Each axis must be present in its respective tensor. * @param other The other tensor to contract with. - * - * Example usage: - * {{{ - * val t1: Tensor[("A", "B", "C"), Float] = ??? - * val t2: Tensor[("D", "E, "F), Float] = ??? - * val result = t1.dot(Axis[A]->Axis[D])(t2) - * }}} */ @targetName("dotOn") def dot[ diff --git a/core/src/main/scala/dimwit/tensor/tensorops/ConvolutionOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/ConvolutionOps.scala index 1656d25..d98b754 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/ConvolutionOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/ConvolutionOps.scala @@ -7,6 +7,7 @@ import dimwit.tensor.Label import dimwit.tensor.Labels import dimwit.tensor.ShapeTypeHelpers.AxisIndex import dimwit.tensor.Tensor +import dimwit.tensor.Tensor3 import dimwit.tensor.TensorOps.IsFloating import dimwit.tensor.TensorOps.swap import me.shadaj.scalapy.py @@ -41,7 +42,7 @@ object ConvolutionOps: * @return A new tensor representing the result of the convolution operation. */ def conv1d[OutChannel: Label]( - kernel: Tensor[(S1, InChannel, OutChannel), V], + kernel: Tensor3[S1, InChannel, OutChannel, V], stride: Stride1[S1] | Int = 1, padding: Padding = Padding.SAME ): Tensor[S1 *: OutChannel *: EmptyTuple, V] = diff --git a/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala index d882011..925d9ca 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala @@ -12,24 +12,48 @@ import dimwit.tensor.ShapeTypeHelpers.SharedAxisRemover import dimwit.tensor.Tensor import dimwit.tensor.Tensor0 import dimwit.tensor.tensorops.FunctionalOps.ZipVmap.TensorsOf +import dimwit.tensortree.TensorTree import me.shadaj.scalapy.py import me.shadaj.scalapy.py.SeqConverters import me.shadaj.scalapy.readwrite.Reader import me.shadaj.scalapy.readwrite.Writer -import dimwit.tensortree.TensorTree -import dimwit.tensor.ShapeTypeHelpers.UnwrapAxes -import dimwit.tensor.ShapeTypeHelpers.AxesRemover + +import scala.NamedTuple.NamedTuple +import scala.annotation.implicitNotFound object FunctionalOps: - type PrependAxes[Axes <: Tuple, FOut] = Axes match - case EmptyTuple => FOut - case h *: t => PrependAxis[h, PrependAxes[t, FOut]] + /** Prepends the axis `L` to every tensor of the tensor tree `FOut`: the result + * type of a `vmap`/`zipvmap` whose body returned `FOut`. + * + * Instances cover tensors, tuples, named tuples and any nesting of those. + */ + @implicitNotFound( + "Cannot prepend Axis[${L}] to ${FOut}. A vmap/zipvmap body may return a Tensor, a tuple, a named tuple, or any nesting of those" + ) + trait PrependAxis[L, FOut]: + type Out + + object PrependAxis: + + type Aux[L, FOut, Out0] = PrependAxis[L, FOut] { type Out = Out0 } + + private def instance[L, FOut, Out0]: Aux[L, FOut, Out0] = + new PrependAxis[L, FOut]: + type Out = Out0 + + given tensor[L, Shape <: Tuple, V]: Aux[L, Tensor[Shape, V], Tensor[L *: Shape, V]] = instance + + given emptyTuple[L]: Aux[L, EmptyTuple, EmptyTuple] = instance + + given consTuple[L, H, HOut, T <: Tuple, TOut <: Tuple](using + head: Aux[L, H, HOut], + tail: Aux[L, T, TOut] + ): Aux[L, H *: T, HOut *: TOut] = instance - type PrependAxis[L, FOut] = FOut match - case Tensor[shape, v] => Tensor[L *: shape, v] - case EmptyTuple => EmptyTuple - case h *: t => PrependAxis[L, h] *: PrependAxis[L, t] + given namedTuple[L, N <: Tuple, V <: Tuple, VOut <: Tuple](using + values: Aux[L, V, VOut] + ): Aux[L, NamedTuple[N, V], NamedTuple[N, VOut]] = instance object ZipVmap: @@ -56,9 +80,9 @@ object FunctionalOps: * * Example usage: * {{{ - * val tensor1: Tensor[(A, B), Int] = ... - * val tensor2: Tensor[(A, B), Int] = ... - * val result: Tensor[(A, C), Int] = ZipVmap.zipvmap(Axis[A])(tensor1, tensor2) { case (t1, t2) => + * val tensor1: Tensor2[A, B, Int] = ... + * val tensor2: Tensor2[A, B, Int] = ... + * val result: Tensor2[A, C, Int] = ZipVmap.zipvmap(Axis[A])(tensor1, tensor2) { case (t1, t2) => * // Perform operations on t1 and t2, which are tensors with axis A removed, and return a new tensor * ... * } @@ -72,9 +96,10 @@ object FunctionalOps: )( f: TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]] => FOut )(using + prependAxis: PrependAxis[L, FOut], toPyTree: TensorTree[FOut], - fromPyTree: TensorTree[PrependAxis[L, FOut]] - ): PrependAxis[L, FOut] = + fromPyTree: TensorTree[prependAxis.Out] + ): prependAxis.Out = val fpy = (args: py.Dynamic) => OnError.traceStack: val tensorList = args.as[Seq[py.Dynamic]].zip(ev.shapesLabels).map: (jaxArr, labels) => @@ -112,9 +137,10 @@ object FunctionalOps: )( f: TensorsOf[ev.RemainingAxes, (V, V)] => FOut )(using + prependAxis: PrependAxis[L, FOut], toPyTree: TensorTree[FOut], - fromPyTree: TensorTree[PrependAxis[L, FOut]] - ): PrependAxis[L, FOut] = + fromPyTree: TensorTree[prependAxis.Out] + ): prependAxis.Out = ZipVmap.zipvmap(axis)(t, other)(f) /** Vectorized mapping over a specified axis of the tensor. @@ -130,10 +156,11 @@ object FunctionalOps: )( f: Tensor[ev.RemainingAxes, V] => FOut )(using + prependAxis: PrependAxis[VmapAxis, FOut], toPyTree: TensorTree[FOut], - fromPyTree: TensorTree[PrependAxis[VmapAxis, FOut]], + fromPyTree: TensorTree[prependAxis.Out], labels: Labels[ev.RemainingAxes] - ): PrependAxis[VmapAxis, FOut] = + ): prependAxis.Out = val fpy = (jxpr: Jax.PyDynamic) => OnError.traceStack: val innerTensor = Tensor[ev.RemainingAxes, V](jxpr) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala index 4ab8d4c..fd41342 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala @@ -27,14 +27,13 @@ import dimwit.tensor.ShapeTypeHelpers.MergeLabels import dimwit.tensor.ShapeTypeHelpers.UnwrapAxes import dimwit.tensor.ShapeTypeHelpers.UnwrapDims import dimwit.tensor.Tensor -import dimwit.tensor.Tensor0 import dimwit.tensor.Tensor1 import dimwit.tensor.TupleHelpers import dimwit.tensor.TupleHelpers.StrictSubset -import dimwit.tensor.TupleHelpers.TensorEvidence.CheckValid -import dimwit.tensor.TupleHelpers.TensorEvidence.ComputeMissing -import dimwit.tensor.TupleHelpers.TensorEvidence.IsPermutation -import dimwit.tensor.TupleHelpers.TensorEvidence.ValidationResult +import dimwit.tensor.TensorEvidence.CheckValid +import dimwit.tensor.TensorEvidence.ComputeMissing +import dimwit.tensor.TensorEvidence.IsPermutation +import dimwit.tensor.TensorEvidence.ValidationResult import dimwit.|+| import me.shadaj.scalapy.py import me.shadaj.scalapy.py.SeqConverters @@ -46,19 +45,55 @@ import scala.util.NotGiven object StructuralOps: - private object Util: + /** Inserts axis `New` directly after axis `Anchor`. */ + trait AxisInserter[T <: Tuple, Anchor, New]: + type Out <: Tuple + + object AxisInserter extends AxisInserterLowPriority: + type Aux[T <: Tuple, Anchor, New, O <: Tuple] = AxisInserter[T, Anchor, New] { type Out = O } + + private[tensorops] def instance[T <: Tuple, Anchor, New, O <: Tuple]: Aux[T, Anchor, New, O] = + new AxisInserter[T, Anchor, New]: + type Out = O + + given found[Anchor, New, T <: Tuple]: Aux[Anchor *: T, Anchor, New, Anchor *: New *: T] = instance + + trait AxisInserterLowPriority: + given search[H, Anchor, New, T <: Tuple, O <: Tuple](using + tail: AxisInserter.Aux[T, Anchor, New, O] + ): AxisInserter.Aux[H *: T, Anchor, New, H *: O] = AxisInserter.instance + + /** Exchanges two axes in a shape, leaving every other axis in place. + * + * The three priority tiers resolve `swap(a, a)` to [[firstTuple]] rather than + * leaving it ambiguous between [[firstTuple]] and [[AxisSwapperSecond.secondTuple]]. + */ + trait AxisSwapper[T <: Tuple, L1, L2]: + type Out <: Tuple + + object AxisSwapper extends AxisSwapperSecond: + type Aux[T <: Tuple, L1, L2, O <: Tuple] = AxisSwapper[T, L1, L2] { type Out = O } + + private[tensorops] def instance[T <: Tuple, L1, L2, O <: Tuple]: Aux[T, L1, L2, O] = + new AxisSwapper[T, L1, L2]: + type Out = O - type InsertBefore[T <: Tuple, A, B] <: Tuple = T match - case EmptyTuple => B *: EmptyTuple - case A *: tail => B *: A *: tail - case h *: tail => h *: InsertBefore[tail, A, B] + given emptyTuple[L1, L2]: Aux[EmptyTuple, L1, L2, EmptyTuple] = instance - type InsertAfter[T <: Tuple, A, B] <: Tuple = T match - case EmptyTuple => B *: EmptyTuple - case A *: tail => A *: B *: tail - case h *: tail => h *: InsertAfter[tail, A, B] + given firstTuple[L1, L2, T <: Tuple, O <: Tuple](using tail: Aux[T, L1, L2, O]): Aux[L1 *: T, L1, L2, L2 *: O] = instance + + trait AxisSwapperSecond extends AxisSwapperOther: + given secondTuple[L1, L2, T <: Tuple, O <: Tuple](using + tail: AxisSwapper.Aux[T, L1, L2, O] + ): AxisSwapper.Aux[L2 *: T, L1, L2, L1 *: O] = AxisSwapper.instance + + trait AxisSwapperOther: + given otherTuple[H, L1, L2, T <: Tuple, O <: Tuple](using + tail: AxisSwapper.Aux[T, L1, L2, O] + ): AxisSwapper.Aux[H *: T, L1, L2, H *: O] = AxisSwapper.instance + + private object Util: - type SliceIndex = Int | List[Int] | Range | Tensor0[Int32] type ExtractLabel[X] = X match case AxisAtIndex[l] => l case AxisAtRange[l] => l @@ -71,61 +106,38 @@ object StructuralOps: object SliceLabelExtractor: - given empty: SliceLabelExtractor[EmptyTuple, EmptyTuple] = + given emptyTuple: SliceLabelExtractor[EmptyTuple, EmptyTuple] = new SliceLabelExtractor[EmptyTuple, EmptyTuple] {} - // New givens for AxisSelector types - given consAxisAtIndex[L, Tail <: Tuple, TailOut <: Tuple](using + given atIndexTuple[L, Tail <: Tuple, TailOut <: Tuple](using tailExt: SliceLabelExtractor[Tail, TailOut] ): SliceLabelExtractor[AxisAtIndex[L] *: Tail, L *: TailOut] = new SliceLabelExtractor[AxisAtIndex[L] *: Tail, L *: TailOut] {} - given consAxisAtRange[L, Tail <: Tuple, TailOut <: Tuple](using + given atRangeTuple[L, Tail <: Tuple, TailOut <: Tuple](using tailExt: SliceLabelExtractor[Tail, TailOut] ): SliceLabelExtractor[AxisAtRange[L] *: Tail, TailOut] = new SliceLabelExtractor[AxisAtRange[L] *: Tail, TailOut] {} - given consAxisAtIndices[L, Tail <: Tuple, TailOut <: Tuple](using + given atIndicesTuple[L, Tail <: Tuple, TailOut <: Tuple](using tailExt: SliceLabelExtractor[Tail, TailOut] ): SliceLabelExtractor[AxisAtIndices[L] *: Tail, TailOut] = new SliceLabelExtractor[AxisAtIndices[L] *: Tail, TailOut] {} - given consAxisAtTupleIndices[L, I <: NonEmptyTuple, Tail <: Tuple, TailOut <: Tuple](using + given atTupleIndicesTuple[L, I <: NonEmptyTuple, Tail <: Tuple, TailOut <: Tuple](using tailExt: SliceLabelExtractor[Tail, TailOut] ): SliceLabelExtractor[AxisAtTupleIndices[L, I] *: Tail, TailOut] = new SliceLabelExtractor[AxisAtTupleIndices[L, I] *: Tail, TailOut] {} - given consAxisAtTensorIndex[L, Tail <: Tuple, TailOut <: Tuple](using + given atTensorIndexTuple[L, Tail <: Tuple, TailOut <: Tuple](using tailExt: SliceLabelExtractor[Tail, TailOut] ): SliceLabelExtractor[AxisAtTensorIndex[L] *: Tail, L *: TailOut] = new SliceLabelExtractor[AxisAtTensorIndex[L] *: Tail, L *: TailOut] {} - // Keep backward compatibility with tuple syntax - given consInt[L, Tail <: Tuple, TailOut <: Tuple](using - tailExt: SliceLabelExtractor[Tail, TailOut] - ): SliceLabelExtractor[(Axis[L], Int) *: Tail, L *: TailOut] = - new SliceLabelExtractor[(Axis[L], Int) *: Tail, L *: TailOut] {} - - given consTensor0Int[L, Tail <: Tuple, TailOut <: Tuple](using - tailExt: SliceLabelExtractor[Tail, TailOut] - ): SliceLabelExtractor[(Axis[L], Tensor0[Int32]) *: Tail, L *: TailOut] = - new SliceLabelExtractor[(Axis[L], Tensor0[Int32]) *: Tail, L *: TailOut] {} - - given consSeq[L, SeqT <: Seq[Int], Tail <: Tuple, TailOut <: Tuple](using - tailExt: SliceLabelExtractor[Tail, TailOut] - ): SliceLabelExtractor[(Axis[L], SeqT) *: Tail, TailOut] = - new SliceLabelExtractor[(Axis[L], SeqT) *: Tail, TailOut] {} - - type Swap[T <: Tuple, A, B] <: Tuple = T match - case EmptyTuple => EmptyTuple - case A *: tail => B *: Swap[tail, A, B] - case B *: tail => A *: Swap[tail, A, B] - case h *: tail => h *: Swap[tail, A, B] - @implicitNotFound("The axis ${L} is already present in the tensor shape ${T}.") trait AxisAbsent[T, L] object AxisAbsent: - given [T <: Tuple, L](using NotGiven[Tuple.Contains[T, L] =:= true]): AxisAbsent[T, L] = new AxisAbsent[T, L] {} + given notContained[T <: Tuple, L](using NotGiven[Tuple.Contains[T, L] =:= true]): AxisAbsent[T, L] = new AxisAbsent[T, L] {} import Util.* @@ -200,15 +212,16 @@ object StructuralOps: afterAxis: Axis[L] )(using newLabel: Label[NewL], - axisIndex: AxisIndex[T, L] - ): Tensor[InsertAfter[T, L, NewL], V] = + axisIndex: AxisIndex[T, L], + inserter: AxisInserter[T, L, NewL] + ): Tensor[inserter.Out, V] = require(tensors.nonEmpty, "Cannot stack an empty sequence of tensors") val axisIdx = axisIndex.index + 1 // we are inserting after the given axis, so shift by 1 val jaxValuesSeq = tensors.map(_.jaxValue).toPythonProxy val stackedJaxValue = Jax.jnp.stack(jaxValuesSeq, axis = axisIdx) val names = summon[Labels[T]].names val newNames = names.take(axisIdx) ++ Seq(newLabel.name) ++ names.drop(axisIdx) - given Labels[InsertAfter[T, L, NewL]] with + given Labels[inserter.Out] with val names = newNames.toSeq Tensor(stackedJaxValue) @@ -267,13 +280,13 @@ object StructuralOps: object ValidConcat: type Aux[T1 <: Tuple, T2 <: Tuple, O <: Tuple] = ValidConcat[T1, T2] { type Out = O } - given recursive[H, T1Tail <: Tuple, T2Tail <: Tuple, OutTail <: Tuple](using + given sameAxisTuple[H, T1Tail <: Tuple, T2Tail <: Tuple, OutTail <: Tuple](using next: ValidConcat.Aux[T1Tail, T2Tail, OutTail] ): ValidConcat[H *: T1Tail, H *: T2Tail] with type Out = H *: OutTail def index: Int = next.index + 1 - given concatAxis[H1, H2, Tail <: Tuple](using + given concatAxisTuple[H1, H2, Tail <: Tuple](using isDifferent: NotGiven[H1 =:= H2] ): ValidConcat[H1 *: Tail, H2 *: Tail] with type Out = (H1 |+| H2) *: Tail @@ -290,7 +303,7 @@ object StructuralOps: object Deconcatenator extends DeconcatenatorLowPriority: type Aux[L, C <: Tuple] = Deconcatenator[L] { type Components = C } - given recursive[A, B, CA <: Tuple, CB <: Tuple](using + given concatenated[A, B, CA <: Tuple, CB <: Tuple](using da: Aux[A, CA], db: Aux[B, CB] ): Aux[A |+| B, Tuple.Concat[CA, CB]] = @@ -299,7 +312,7 @@ object StructuralOps: def labels = da.labels ++ db.labels trait DeconcatenatorLowPriority: - given base[L](using l: Label[L]): Deconcatenator.Aux[L, L *: EmptyTuple] = + given single[L](using l: Label[L]): Deconcatenator.Aux[L, L *: EmptyTuple] = new Deconcatenator[L]: type Components = L *: EmptyTuple def labels = List(l) @@ -312,12 +325,12 @@ object StructuralOps: type Aux[C <: Tuple, F <: Tuple, S, V, O <: Tuple] = TensorTupleMaker[C, F, S, V] { type Out = O } - given empty[F <: Tuple, S, V]: Aux[EmptyTuple, F, S, V, EmptyTuple] = + given emptyTuple[F <: Tuple, S, V]: Aux[EmptyTuple, F, S, V, EmptyTuple] = new TensorTupleMaker[EmptyTuple, F, S, V]: type Out = EmptyTuple def apply(a: Seq[Jax.PyDynamic], c: List[Label[?]], o: Seq[String], i: Int) = EmptyTuple - given cons[Head, Tail <: Tuple, F <: Tuple, S, V, NewShape <: Tuple](using + given consTuple[Head, Tail <: Tuple, F <: Tuple, S, V, NewShape <: Tuple](using replacer: TupleHelpers.Replacer[F, S, Head] { type Out = NewShape }, tailMaker: TensorTupleMaker[Tail, F, S, V] ): Aux[Head *: Tail, F, S, V, Tensor[NewShape, V] *: tailMaker.Out] = @@ -518,7 +531,6 @@ object StructuralOps: targetDims.zip(inputList).foreach { case (dimIndex, input) => val dimSize = tensor.shape.dimensions(dimIndex) input match - // New AxisSelector types case AxisAtIndex(_, idx) => indicesBuffer(dimIndex) = py.Any.from(idx) case AxisAtRange(_, range) => @@ -529,17 +541,6 @@ object StructuralOps: indicesBuffer(dimIndex) = indices.toList.asInstanceOf[List[Int]].map(py.Any.from).toPythonCopy case AxisAtTensorIndex(_, tensorIdx) => indicesBuffer(dimIndex) = tensorIdx.jaxValue - // Backward compatibility with tuples - case (_, sliceIndex) => - sliceIndex match - case sliceSeq: List[Int] @unchecked => - indicesBuffer(dimIndex) = sliceSeq.map(py.Any.from).toPythonProxy - case range: Range @unchecked => - indicesBuffer(dimIndex) = PySlice(range.head, range.last + 1, range.step) - case idx: Int => - indicesBuffer(dimIndex) = py.Any.from(idx) - case tensorId: Tensor0[Int32] @unchecked => - indicesBuffer(dimIndex) = tensorId.jaxValue } Jax.Dynamic.global.tuple(indicesBuffer.toSeq.toPythonProxy) @@ -860,9 +861,10 @@ object StructuralOps: )(using labels: Labels[T], axisIndex1: AxisIndex[T, L1], - axisIndex2: AxisIndex[T, L2] - ): Tensor[Swap[T, L1, L2], V] = - given Labels[Swap[T, L1, L2]] with + axisIndex2: AxisIndex[T, L2], + swapper: AxisSwapper[T, L1, L2] + ): Tensor[swapper.Out, V] = + given Labels[swapper.Out] with def names = val originalNames = summon[Labels[T]].names val ax1Name = summon[Label[L1]].name diff --git a/core/src/main/scala/dimwit/tensortree/TensorTree.scala b/core/src/main/scala/dimwit/tensortree/TensorTree.scala index c38c3fd..c23fb67 100644 --- a/core/src/main/scala/dimwit/tensortree/TensorTree.scala +++ b/core/src/main/scala/dimwit/tensortree/TensorTree.scala @@ -104,7 +104,7 @@ object TensorTree: // extends TensorTreeLowPriority: /** 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 + given tensor[Q <: Tuple, V](using n: Labels[Q]): TensorTree[Tensor[Q, V]] with def map(t: Tensor[Q, V], f: [T <: Tuple, V2] => (Labels[T]) ?=> (Tensor[T, V2] => Tensor[T, V2])): Tensor[Q, V] = import TensorOps.retag f[Q, V](using n)(t.retag[Q](using n)) @@ -138,7 +138,7 @@ object TensorTree: // extends TensorTreeLowPriority: /** Tensor tree instance for an empty tree. This can be useful * for example for optimizers that don't have internal state */ - given TensorTree[Unit] with + given unit: TensorTree[Unit] with def map(p: Unit, f: [T <: Tuple, V] => (Labels[T]) ?=> (Tensor[T, V] => Tensor[T, V])): Unit = () def mapWithName(p: Unit, f: [T <: Tuple, V] => (Labels[T]) ?=> ((String, Tensor[T, V]) => Tensor[T, V]), path: String = ""): Unit = () def mapLeaves[A](p: Unit, f: [T <: Tuple, V] => (Labels[T]) ?=> (Tensor[T, V] => A)): Iterator[A] = Iterator.empty @@ -150,49 +150,9 @@ object TensorTree: // extends TensorTreeLowPriority: def toNumpyTree(p: Unit): Jax.PyAny = py.Dynamic.global.None def fromNumpyTree(pyVal: Jax.PyAny): Unit = () - /** Instance for a tuple of two tensors */ - given tupleInstance[P1, P2](using t1: TensorTree[P1], t2: TensorTree[P2]): TensorTree[(P1, P2)] with - def map(p: (P1, P2), f: [T <: Tuple, V] => (Labels[T]) ?=> (Tensor[T, V] => Tensor[T, V])): (P1, P2) = - (t1.map(p._1, f), t2.map(p._2, f)) - - def mapWithName(p: (P1, P2), f: [T <: Tuple, V] => (Labels[T]) ?=> ((String, Tensor[T, V]) => Tensor[T, V]), path: String = ""): (P1, P2) = - val p1Path = if path.isEmpty then "_1" else s"$path._1" - val p2Path = if path.isEmpty then "_2" else s"$path._2" - (t1.mapWithName(p._1, f, p1Path), t2.mapWithName(p._2, f, p2Path)) - - def mapLeaves[A](p: (P1, P2), f: [T <: Tuple, V] => (Labels[T]) ?=> (Tensor[T, V] => A)): Iterator[A] = - t1.mapLeaves(p._1, f) ++ t2.mapLeaves(p._2, f) - - def foreach(p: (P1, P2), f: [T <: Tuple, V] => (Labels[T]) ?=> (Tensor[T, V] => Unit)): Unit = - t1.foreach(p._1, f) - t2.foreach(p._2, f) - - def foreachWithName(p: (P1, P2), f: [T <: Tuple, V] => (Labels[T]) ?=> ((String, Tensor[T, V]) => Unit), path: String = ""): Unit = - val p1Path = if path.isEmpty then "_1" else s"$path._1" - val p2Path = if path.isEmpty then "_2" else s"$path._2" - t1.foreachWithName(p._1, f, p1Path) - t2.foreachWithName(p._2, f, p2Path) - - def zipMap(p1: (P1, P2), p2: (P1, P2), f: [T <: Tuple, V] => (Labels[T]) ?=> ((Tensor[T, V], Tensor[T, V]) => Tensor[T, V])): (P1, P2) = - (t1.zipMap(p1._1, p2._1, f), t2.zipMap(p1._2, p2._2, f)) - - def toPyTree(p: (P1, P2)): Jax.PyAny = - py.Dynamic.global.tuple(Seq(t1.toPyTree(p._1), t2.toPyTree(p._2)).toPythonProxy) - - def fromPyTree(pyVal: Jax.PyAny): (P1, P2) = - val pyTuple = pyVal.as[py.Dynamic] - (t1.fromPyTree(pyTuple.itemAt(0)), t2.fromPyTree(pyTuple.itemAt(1))) - - def toNumpyTree(p: (P1, P2)): Jax.PyAny = - py.Dynamic.global.tuple(Seq(t1.toNumpyTree(p._1), t2.toNumpyTree(p._2)).toPythonProxy) - - def fromNumpyTree(pyVal: Jax.PyAny): (P1, P2) = - val pyTuple = pyVal.as[py.Dynamic] - (t1.fromNumpyTree(pyTuple.itemAt(0)), t2.fromNumpyTree(pyTuple.itemAt(1))) - /** Instance for a list of tensor trees */ - given listInstance[P](using tp: TensorTree[P]): TensorTree[List[P]] with + given list[P](using tp: TensorTree[P]): TensorTree[List[P]] with def map(l: List[P], f: [T <: Tuple, V] => (Labels[T]) ?=> (Tensor[T, V] => Tensor[T, V])): List[P] = l.map(elem => tp.map(elem, f)) @@ -233,7 +193,7 @@ object TensorTree: // extends TensorTreeLowPriority: val len = py.Dynamic.global.len(pyList).as[Int] List.tabulate(len)(i => tp.fromNumpyTree(pyList.itemAt(i))) - given namedTupleInstance[N <: Tuple, V <: Tuple](using tt: TensorTree[V]): TensorTree[NamedTuple[N, V]] with + given namedTuple[N <: Tuple, V <: Tuple](using tt: TensorTree[V]): TensorTree[NamedTuple[N, V]] with def map(p: NamedTuple[N, V], f: [T <: Tuple, V2] => (Labels[T]) ?=> (Tensor[T, V2] => Tensor[T, V2])): NamedTuple[N, V] = tt.map(p.toTuple, f) diff --git a/core/src/main/scala/dimwit/tensortree/TensorTreeFormat.scala b/core/src/main/scala/dimwit/tensortree/TensorTreeFormat.scala index 028fbf7..490501d 100644 --- a/core/src/main/scala/dimwit/tensortree/TensorTreeFormat.scala +++ b/core/src/main/scala/dimwit/tensortree/TensorTreeFormat.scala @@ -1,6 +1,5 @@ package dimwit.tensortree -import dimwit.jax.Jax import me.shadaj.scalapy.py import java.nio.file.Path diff --git a/core/src/main/scala/dimwit/tensortree/TreeOf.scala b/core/src/main/scala/dimwit/tensortree/TreeOf.scala index d71b75a..7b13145 100644 --- a/core/src/main/scala/dimwit/tensortree/TreeOf.scala +++ b/core/src/main/scala/dimwit/tensortree/TreeOf.scala @@ -16,7 +16,7 @@ trait TreeOf[P, V] object TreeOf: // 1. Base case for Tensors - given [Q <: Tuple, V](using TensorTree[Tensor[Q, V]]): TreeOf[Tensor[Q, V], V] with {} + given tensor[Q <: Tuple, V](using TensorTree[Tensor[Q, V]]): TreeOf[Tensor[Q, V], V] with {} // 2. Inductive base cases for Tuples // This allows the compiler to step through the case class fields and lock in V. @@ -28,12 +28,12 @@ object TreeOf: )(using TensorTree[H *: T]): TreeOf[H *: T, V] with {} // 3. Standard collections - given listInstance[A: TensorTree, V](using TreeOf[A, V]): TreeOf[List[A], V] with {} + given list[A: TensorTree, V](using TreeOf[A, V]): TreeOf[List[A], V] with {} // given mapInstance[K, A, V](using TreeOf[A, V]): TreeOf[Map[K, A], V] with {} // 4. Named tuples, delegating to the TreeOf instance of the underlying value tuple - given namedTupleInstance[N <: Tuple, V <: Tuple: TensorTree, Fl](using TreeOf[V, Fl]): TreeOf[NamedTuple[N, V], Fl] with {} + given namedTuple[N <: Tuple, V <: Tuple: TensorTree, Fl](using TreeOf[V, Fl]): TreeOf[NamedTuple[N, V], Fl] with {} inline given derived[P <: Product: TensorTree, V](using evNotTuple: NotGiven[P <:< Tuple], @@ -100,7 +100,7 @@ object TreeOf: // helper typeclass trait IsFloatingTensor[P, V] object IsFloatingTensor: - given [T <: Tuple, V: IsFloating]: IsFloatingTensor[Tensor[T, V], V] with {} + given tensor[T <: Tuple, V: IsFloating]: IsFloatingTensor[Tensor[T, V], V] with {} // Scalar broadcast extensions (Tensor0 op Tree) extension [V: IsFloating](p2: Tensor0[V]) diff --git a/core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala b/core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala index f9cdd4b..3de3295 100644 --- a/core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala +++ b/core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala @@ -2,7 +2,9 @@ package dimwit.autodiff import dimwit.* import dimwit.Conversions.given -import dimwit.autodiff.Autodiff.Gradient + +/** A parameter tree, declared top level so its Mirror is available. */ +case class JacParams(w: Tensor1[A, Float32], b: Tensor1[B, Float32]) derives TensorTree class AutodiffSuite extends DimwitTest: @@ -70,42 +72,81 @@ class AutodiffSuite extends DimwitTest: val x = Tensor1(Axis[A]).fromArray(Array(1.0f, 1.0f)) jf(x) should approxEqual(Tensor2.eye(x.extent(Axis[A]), x.vtype) *! 2.0f) - describe("jacRev / jacFwd"): + describe("jacRev"): - // setup engines to test both modes in the same way - val engines = List( - ("jacRev", [In: TensorTree, Out: TensorTree] => (f: In => Out) => (gradTree: TensorTree[Gradient[In, Out]]) ?=> Autodiff.jacRev[In, Out](f)), - ("jacFwd", [In: TensorTree, Out: TensorTree] => (f: In => Out) => (gradTree: TensorTree[Gradient[In, Out]]) ?=> Autodiff.jacFwd[In, Out](f)) - ) + it("d¹ of f(x1, x2) = (x2, x1)"): + def f(x1: Tensor1[A, Float32], x2: Tensor1[A, Float32]): (Tensor1[A, Float32], Tensor1[A, Float32]) = (x2, x1) + val df = Autodiff.jacRev(f.tupled) + + val x1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 0.0f)) + val x2 = Tensor1(Axis[A]).fromArray(Array(0.0f, 1.0f)) + val ((x1_dx1, x1_dx2), (x2_dx1, x2_dx2)) = df(x1, x2) + + // the first output is x2, so it depends on x2 only, and the other way round + x1_dx1 should approxEqual(Tensor.like(x1_dx1).fill(0f)) + x1_dx2 should approxEqual(Tensor2.eye(x1.extent(Axis[A]), x1.vtype)) + x2_dx1 should approxEqual(Tensor2.eye(x2.extent(Axis[A]), x2.vtype)) + x2_dx2 should approxEqual(Tensor.like(x2_dx2).fill(0f)) + + it("d¹ of f: Tensor1[A] => Tensor1[B] keeps the output axis first"): + def f(x: Tensor1[A, Float32]): Tensor1[B, Float32] = x.relabel(Axis[A] -> Axis[B]) *! 2.0f + val df = Autodiff.jacRev(f) + + val x = Tensor1(Axis[A]).fromArray(Array(1.0f, 1.0f)) + df(x).axes shouldBe List("B", "A") + df(x) should approxEqual((Tensor2.eye(x.extent(Axis[A])) *! 2.0f).relabelAll((Axis[B], Axis[A]))) + + it("d² of f(x1, x2) = sum(x1 * x2)"): + def f(x1: Tensor1[A, Float32], x2: Tensor1[A, Float32]): Tensor0[Float32] = (x1 * x2).sum + val ddf = Autodiff.jacRev(Autodiff.jacRev(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 ((x1_dx1, x1_dx2), (x2_dx1, x2_dx2)) = ddf(x1, x2) + + // d²/dx1² and d²/dx2² vanish, the mixed partials are the identity + x1_dx1 should approxEqual(Tensor.like(x1_dx1).fill(0f)) + x1_dx2 should approxEqual(Tensor2.eye(x1.extent(Axis[A]), x1.vtype)) + x2_dx1 should approxEqual(Tensor2.eye(x2.extent(Axis[A]), x2.vtype)) + x2_dx2 should approxEqual(Tensor.like(x2_dx2).fill(0f)) + + describe("jacFwd"): - engines.foreach: - case (modeName, jacMode) => - it(s"$modeName d¹ on f: R² -> R², f(x) = swap(x)"): - def f(x1: Tensor1[A, Float32], x2: Tensor1[A, Float32]): (Tensor1[A, Float32], Tensor1[A, Float32]) = (x2, x1) - val df = jacMode(f.tupled) - val x1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 0.0f)) - val x2 = Tensor1(Axis[A]).fromArray(Array(0.0f, 1.0f)) - val (x1Grad, x2Grad) = df(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)) - x2_dx1 should approxEqual(Tensor2.eye(x2.extent(Axis[A]), x2.vtype)) - x2_dx2 should approxEqual(Tensor.like(x2_dx2).fill(0f)) - - it(s"$modeName d² on f: R² -> R, f(x1, x2) = sum(x1 * x2)"): - def f(x1: Tensor1[A, Float32], x2: Tensor1[A, Float32]): Tensor0[Float32] = (x1 * x2).sum - val df = jacMode(f.tupled) - val ddf = jacMode(df) - 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) = ddf(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)) + it("d¹ of f(x1, x2) = (x2, x1)"): + def f(x1: Tensor1[A, Float32], x2: Tensor1[A, Float32]): (Tensor1[A, Float32], Tensor1[A, Float32]) = (x2, x1) + val df = Autodiff.jacFwd(f.tupled) + + val x1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 0.0f)) + val x2 = Tensor1(Axis[A]).fromArray(Array(0.0f, 1.0f)) + val ((x1_dx1, x1_dx2), (x2_dx1, x2_dx2)) = df(x1, x2) + + // the first output is x2, so it depends on x2 only, and the other way round + x1_dx1 should approxEqual(Tensor.like(x1_dx1).fill(0f)) + x1_dx2 should approxEqual(Tensor2.eye(x1.extent(Axis[A]), x1.vtype)) + x2_dx1 should approxEqual(Tensor2.eye(x2.extent(Axis[A]), x2.vtype)) + x2_dx2 should approxEqual(Tensor.like(x2_dx2).fill(0f)) + + it("d¹ of f: Tensor1[A] => Tensor1[B] keeps the output axis first"): + def f(x: Tensor1[A, Float32]): Tensor1[B, Float32] = x.relabel(Axis[A] -> Axis[B]) *! 2.0f + val df = Autodiff.jacFwd(f) + + val x = Tensor1(Axis[A]).fromArray(Array(1.0f, 1.0f)) + df(x).axes shouldBe List("B", "A") + df(x) should approxEqual((Tensor2.eye(x.extent(Axis[A])) *! 2.0f).relabelAll((Axis[B], Axis[A]))) + + it("d² of f(x1, x2) = sum(x1 * x2)"): + def f(x1: Tensor1[A, Float32], x2: Tensor1[A, Float32]): Tensor0[Float32] = (x1 * x2).sum + val ddf = Autodiff.jacFwd(Autodiff.jacFwd(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 ((x1_dx1, x1_dx2), (x2_dx1, x2_dx2)) = ddf(x1, x2) + + // d²/dx1² and d²/dx2² vanish, the mixed partials are the identity + x1_dx1 should approxEqual(Tensor.like(x1_dx1).fill(0f)) + x1_dx2 should approxEqual(Tensor2.eye(x1.extent(Axis[A]), x1.vtype)) + x2_dx1 should approxEqual(Tensor2.eye(x2.extent(Axis[A]), x2.vtype)) + x2_dx2 should approxEqual(Tensor.like(x2_dx2).fill(0f)) describe("hessian"): describe("single parameter function"): @@ -137,6 +178,107 @@ 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("jacobian of a function whose input and output axes differ"): + + it("non-square jacobian: Tensor1[A] => Tensor1[B]"): + def f(x: Tensor1[A, Float32]): Tensor1[B, Float32] = x.relabel(Axis[A] -> Axis[B]) *! 2.0f + val jf = Autodiff.jacobian(f) + + val x = Tensor1(Axis[A]).fromArray(Array(1.0f, 1.0f)) + jf(x).axes shouldBe List("B", "A") + jf(x) should approxEqual((Tensor2.eye(x.extent(Axis[A])) *! 2.0f).relabelAll((Axis[B], Axis[A]))) + + it("primes an input axis that collides with an output axis"): + def f(x: Tensor2[A, B, Float32]): Tensor1[B, Float32] = x.sum(Axis[A]) + val jf = Autodiff.jacobian(f) + + val x = Tensor(Shape(Axis[A] -> 3, Axis[B] -> 2)).fill(1f) + val jac = jf(x) + jac.axes shouldBe List("B", "A", "B'") + jac.shape(Axis[A]) shouldBe 3 + // d(sum over A)_b / dx(a, b') is 1 exactly when b == b', for every a + jac.sum shouldEqual Tensor0(6.0f) + + it("primes a colliding input axis that is not the head of the output"): + def f(x: Tensor1[B, Float32]): Tensor2[A, B, Float32] = x.broadcastTo(Shape(Axis[A] -> 3, Axis[B] -> 2)) + val jf = Autodiff.jacobian(f) + + val x = Tensor1(Axis[B]).fromArray(Array(1.0f, 2.0f)) + val jac = jf(x) + jac.axes shouldBe List("A", "B", "B'") + // the broadcast copies x, so d out(a, b) / dx(b') is 1 exactly when b == b' + jac.sum shouldEqual Tensor0(6.0f) + + it("primes only the colliding axis of a multi-axis input"): + def f(x: Tensor2[A, B, Float32]): Tensor2[C, A, Float32] = + x.sum(Axis[B]).broadcastTo(Shape(Axis[C] -> 2, Axis[A] -> 3)) + val jf = Autodiff.jacobian(f) + + val x = Tensor(Shape(Axis[A] -> 3, Axis[B] -> 4)).fill(1f) + val jac = jf(x) + jac.axes shouldBe List("C", "A", "A'", "B") + jac.sum shouldEqual Tensor0(24.0f) + + it("jacobian over a tuple input with differing axes"): + def f(x: Tensor1[A, Float32], y: Tensor1[B, Float32]): Tensor0[Float32] = x.sum * y.sum + val jf = Autodiff.jacobian(f.tupled) + + val x = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) + val y = Tensor1(Axis[B]).fromArray(Array(3.0f, 4.0f)) + val (dx, dy) = jf(x, y) + dx should approxEqual(Tensor1(Axis[A]).fromArray(Array(7.0f, 7.0f))) + dy should approxEqual(Tensor1(Axis[B]).fromArray(Array(3.0f, 3.0f))) + + it("hessian of a scalar loss over two different axes"): + def f(x1: Tensor1[A, Float32], x2: Tensor1[B, Float32]): Tensor0[Float32] = x1.sum * x2.sum + val hf = Autodiff.hessian(f.tupled) + + val x1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) + val x2 = Tensor1(Axis[B]).fromArray(Array(3.0f, 4.0f)) + val (d1, d2) = hf(x1, x2) + val (d1_d1, d1_d2) = d1 + val (d2_d1, d2_d2) = d2 + d1_d1 should approxEqual(Tensor.like(d1_d1).fill(0f)) + d1_d2 should approxEqual(Tensor.like(d1_d2).fill(1f)) + d2_d1 should approxEqual(Tensor.like(d2_d1).fill(1f)) + d2_d2 should approxEqual(Tensor.like(d2_d2).fill(0f)) + + describe("jacobian of structures that are not tensors or plain tuples"): + + val params = JacParams( + Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)), + Tensor1(Axis[B]).fromArray(Array(3.0f, 4.0f)) + ) + + it("differentiates a case class tree into a named tuple of its fields"): + def f(p: JacParams): JacParams = p + val jf = Autodiff.jacobian(f) + val jac = jf(params) + + jac.w.w should approxEqual(Tensor2.eye(params.w.extent(Axis[A]))) + jac.w.b should approxEqual(Tensor.like(jac.w.b).fill(0f)) + jac.b.w should approxEqual(Tensor.like(jac.b.w).fill(0f)) + jac.b.b should approxEqual(Tensor2.eye(params.b.extent(Axis[B]))) + + it("takes the hessian of a scalar loss over a case class tree"): + def loss(p: JacParams): Tensor0[Float32] = (p.w * p.w).sum + (p.b * p.b).sum + val hf = Autodiff.hessian(loss) + val hess = hf(params) + + hess.w.w should approxEqual(Tensor2.eye(params.w.extent(Axis[A])) *! 2.0f) + hess.w.b should approxEqual(Tensor.like(hess.w.b).fill(0f)) + hess.b.b should approxEqual(Tensor2.eye(params.b.extent(Axis[B])) *! 2.0f) + + it("differentiates a function returning a named tuple"): + def f(x: Tensor1[A, Float32]): (u: Tensor1[A, Float32], v: Tensor1[A, Float32]) = + (u = x *! 2.0f, v = x *! 3.0f) + val jf = Autodiff.jacobian(f) + + val x = Tensor1(Axis[A]).fromArray(Array(1.0f, 1.0f)) + val jac = jf(x) + jac.u should approxEqual(Tensor2.eye(x.extent(Axis[A])) *! 2.0f) + jac.v should approxEqual(Tensor2.eye(x.extent(Axis[A])) *! 3.0f) + describe("Complex application"): it("case class support"): case class Params(w: Tensor1[A, Float32], b: Tensor0[Float32]) diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala index 0278b1f..09bceb5 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala @@ -29,6 +29,28 @@ class TensorOpsFunctionalSuite extends DimwitTest: y1 shouldEqual (t +! 5f) y2 shouldEqual (t -! 5f) + it("vmap return named tuple"): + val t = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3)).fill(0f) + val res = t.vmap(Axis[A]): x => + (first = x +! 5f, second = x -! 5f) + res.first shouldEqual (t +! 5f) + res.second shouldEqual (t -! 5f) + + it("vmap return nested named and plain tuples"): + val t = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3)).fill(0f) + val res = t.vmap(Axis[A]): x => + (nested = (x +! 5f, (inner = x -! 5f)), plain = x) + res.nested._1 shouldEqual (t +! 5f) + res.nested._2.inner shouldEqual (t -! 5f) + res.plain shouldEqual t + + it("vmap return named tuple inside a plain tuple"): + val t = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3)).fill(0f) + val (y1, named) = t.vmap(Axis[A]): x => + (x +! 5f, (inner = x -! 5f)) + y1 shouldEqual (t +! 5f) + named.inner shouldEqual (t -! 5f) + it("vmap over Axis B (columns)"): val res = t2.vmap(Axis[B])(_.sum) res shouldEqual Tensor1(Axis[B]).fromArray(Array(4.0f, 6.0f)) @@ -75,6 +97,15 @@ class TensorOpsFunctionalSuite extends DimwitTest: y1 shouldEqual (t1 + t2) y2 shouldEqual (t1 - t2) + it("zipvmap2 return named tuple"): + val t1 = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3)).fill(0f) + val t2 = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3)).fill(1f) + val res = zipvmap(Axis[A])(t1, t2): + case (x1, x2) => + (sum = x1 + x2, diff = x1 - x2) + res.sum shouldEqual (t1 + t2) + res.diff shouldEqual (t1 - t2) + describe("vapply (Axis-wise application)"): def l2[L: Label](v1: Tensor1[L, Float32], v2: Tensor1[L, Float32]): Tensor0[Float32] = (v1 - v2).pow(2.0f).sum.sqrt diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala index 338ac28..6b84bb2 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala @@ -1,7 +1,6 @@ package dimwit.tensor import dimwit.* -import dimwit.tensor.Labels.concat import scala.compiletime.testing.typeCheckErrors class TensorOpsStructureSuite extends DimwitTest: @@ -214,6 +213,42 @@ class TensorOpsStructureSuite extends DimwitTest: t.shape(Axis[B]) shouldBe (3) t.shape(Axis[C]) shouldBe (4) + describe("swap function"): + + it("swaps two axes that both sit behind the head"): + val swapped = t3.swap(Axis[B], Axis[C]) + swapped.axes shouldBe List("A", "C", "B") + swapped should approxEqual(t3.transpose(Axis[A], Axis[C], Axis[B])) + + it("swaps the head axis with a later one"): + val swapped = t3.swap(Axis[A], Axis[B]) + swapped.axes shouldBe List("B", "A", "C") + swapped should approxEqual(t3.transpose(Axis[B], Axis[A], Axis[C])) + + it("swaps the outermost and innermost axis"): + val swapped = t3.swap(Axis[A], Axis[C]) + swapped.axes shouldBe List("C", "B", "A") + swapped should approxEqual(t3.transpose(Axis[C], Axis[B], Axis[A])) + + it("swapping an axis with itself is the identity"): + val same = t3.swap(Axis[A], Axis[A]) + same.axes shouldBe List("A", "B", "C") + same should approxEqual(t3) + + describe("dropPrimes"): + + it("drops the prime from a primed trailing axis"): + val primed = Tensor(Shape(Axis[A] -> 2, Axis[Prime[B]] -> 3)).fill(1f) + val dropped = primed.dropPrimes + dropped.axes shouldBe List("A", "B") + dropped should approxEqual(Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3)).fill(1f)) + + it("drops primes from a mix of primed and unprimed axes"): + val primed = Tensor(Shape(Axis[Prime[A]] -> 2, Axis[B] -> 3, Axis[Prime[C]] -> 4)).fill(1f) + val dropped = primed.dropPrimes + dropped.axes shouldBe List("A", "B", "C") + dropped.shape(Axis[C]) shouldBe 4 + describe("Dimension manipulation"): it("squeeze axis of size 1"): @@ -573,3 +608,23 @@ class TensorOpsStructureSuite extends DimwitTest: ) ) )) + + it("stacks after an axis that is not the head"): + val t1 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) + val t2 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(5.0f, 6.0f), Array(7.0f, 8.0f))) + + val stacked = stack(Seq(t1, t2), newAxis = Axis[C], afterAxis = Axis[B]) + + stacked.axes shouldBe List("A", "B", "C") + stacked should approxEqual(Tensor3(Axis[A], Axis[B], Axis[C]).fromArray( + Array( + Array(Array(1.0f, 5.0f), Array(2.0f, 6.0f)), + Array(Array(3.0f, 7.0f), Array(4.0f, 8.0f)) + ) + )) + + it("stacks after the middle axis of a 3D tensor"): + val stacked = stack(Seq(t3, t3), newAxis = Axis[D], afterAxis = Axis[B]) + + stacked.axes shouldBe List("A", "B", "D", "C") + stacked.shape(Axis[D]) shouldBe 2 diff --git a/examples/src/main/scala/dimwit/dataset/MNISTLoader.scala b/examples/src/main/scala/dimwit/dataset/MNISTLoader.scala index fef739f..138c512 100644 --- a/examples/src/main/scala/dimwit/dataset/MNISTLoader.scala +++ b/examples/src/main/scala/dimwit/dataset/MNISTLoader.scala @@ -61,7 +61,7 @@ object MNISTLoader: finally file.close() - private def createDataset[S <: Sample: Label](imagesFile: String, labelsFile: String, maxSamples: Option[Int] = None): Try[Tuple2[Tensor[(S, Height, Width), Float32], Tensor1[S, Int8]]] = + private def createDataset[S <: Sample: Label](imagesFile: String, labelsFile: String, maxSamples: Option[Int] = None): Try[Tuple2[Tensor3[S, Height, Width, Float32], Tensor1[S, Int8]]] = Try: val images = loadImages[S](imagesFile, maxSamples) val labels = loadLabels[S](labelsFile, maxSamples) @@ -69,12 +69,12 @@ object MNISTLoader: val imagesFloat = images.asFloat32 /! 255.0f (imagesFloat, labels) - def createTrainingDataset(dataDir: String = "data", maxSamples: Option[Int] = None): Try[Tuple2[Tensor[(TrainSample, Height, Width), Float32], Tensor1[TrainSample, Int8]]] = + def createTrainingDataset(dataDir: String = "data", maxSamples: Option[Int] = None): Try[Tuple2[Tensor3[TrainSample, Height, Width, Float32], Tensor1[TrainSample, Int8]]] = val imagesFile = s"$dataDir/train-images-idx3-ubyte" val labelsFile = s"$dataDir/train-labels-idx1-ubyte" createDataset[TrainSample](imagesFile, labelsFile, maxSamples) - def createTestDataset(dataDir: String = "data", maxSamples: Option[Int] = None): Try[Tuple2[Tensor[(TestSample, Height, Width), Float32], Tensor1[TestSample, Int8]]] = + def createTestDataset(dataDir: String = "data", maxSamples: Option[Int] = None): Try[Tuple2[Tensor3[TestSample, Height, Width, Float32], Tensor1[TestSample, Int8]]] = val imagesFile = s"$dataDir/t10k-images-idx3-ubyte" val labelsFile = s"$dataDir/t10k-labels-idx1-ubyte" createDataset[TestSample](imagesFile, labelsFile, maxSamples)