Skip to content

Type fixes #148

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
benikm91 merged 1 commit into from
Aug 19, 2026
Merged

Type fixes #148

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 32 additions & 26 deletions AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])
// ),
Expand Down Expand Up @@ -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]),
// ???)
Expand Down Expand Up @@ -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])
// ),
Expand Down Expand Up @@ -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])
// ),
Expand All @@ -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])
// ),
Expand Down Expand Up @@ -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]]
// )
Expand All @@ -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]]
// )
Expand Down Expand Up @@ -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])
// ),
Expand Down Expand Up @@ -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] =:=
Expand All @@ -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
// ^^
```
Expand All @@ -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])
// ),
Expand Down
124 changes: 98 additions & 26 deletions core/src/main/scala/dimwit/autodiff/Autodiff.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down Expand Up @@ -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:
Expand All @@ -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)))
Expand All @@ -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)))
Expand All @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion core/src/main/scala/dimwit/linalg/LinearAlgebra.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
23 changes: 4 additions & 19 deletions core/src/main/scala/dimwit/package.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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}
Expand All @@ -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.*
Expand Down
Loading
Loading