Type fixes

[benikm91]

We have several problems in our type definitions that Claude detected.

This PR has 3 commits that I will squash before merging but document this problem:

1. Write tests that fail and show the issues in the current implementations.
2. The fixes that the new tests run
3. A cleanup comment improving the tests now that things compile (make "compile" tests to "run" test).

---

This is an extension of #144, as the problem I detected in zipvmap made me question some more match type implementations.

Summary from Claude

A match type only advances past a case when the scrutinee is provably disjoint
from it. "Does not match" is not enough — no match plus no disjointness proof means
reduction stops there, permanently.

Our axis labels are traits (trait A derives Label), and two traits are never
provably disjoint — nothing stops a third type from extending both. So a case of
the form case L1 *: tail can only ever decide when L1 is literally the head of
the shape. At position 0 it matches; anywhere else it is stuck.

That is the whole bug class. Swap is the clearest example:

type Swap[T <: Tuple, A, B] <: Tuple = T match
  case A *: tail => B *: Swap[tail, A, B]
  case B *: tail => A *: Swap[tail, A, B]
  case h *: tail => h *: Swap[tail, A, B]

Swap[(A, B, C), B, C] never reduces — the head is A, which the compiler can
neither match against B nor rule out. swap still compiled, because its return
type never has to reduce inside the method body; the caller just received an
unreduced Swap[...] that failed the moment it was ascribed or passed on. Combined
with swap having no test coverage at all, that hid it completely.

Note what does work: UnwrapAxes matches on Axis[a] *: tail, and Axis is a
final class, so the compiler can rule it out against a bare label. The usable rule
is therefore:

A match type may branch on structure — *:, EmptyTuple, Axis[_],
AxisExtent[_], Tensor[_, _] — but never on the identity of two labels.

[Copilot]

Pull request overview

This PR addresses Scala 3 match-type reduction “stuck” failures in DimWit’s type-level shape machinery (notably for trait-based axis labels) by replacing identity-based match types with structure-driven typeclass evidence, and adds regression tests to ensure runtime operations line up with the new compile-time types.

Changes:

• Replace match-type implementations for axis insertion/swapping, prime removal, vmap/zipvmap output axis-prepending, and autodiff gradient/hessian result typing with typeclass-based derivations.
• Add/expand tests covering swap, stack(..., afterAxis=...), named-tuple returns from vmap, and jacobian/hessian behavior when input/output axes differ and for case-class trees.
• Remove now-unused match-type aliases / imports related to the old implementations.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala	Adds regression tests for swap, dropPrimes, and stack with afterAxis not at the head.
core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala	Adds vmap tests for named tuples and nested named/plain tuple structures.
core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala	Adds broader jacobian/hessian tests (non-square, axis collisions via Prime, tuple inputs, and case-class trees).
core/src/main/scala/dimwit/tensortree/TensorTreeFormat.scala	Removes an unused import.
core/src/main/scala/dimwit/tensor/TupleHelpers.scala	Removes obsolete match-type aliases in favor of the typeclass-based PrimeConcat machinery.
core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala	Introduces AxisInserter/AxisSwapper typeclasses and wires them into stack(..., afterAxis=...) and swap.
core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala	Replaces match-type PrependAxis with a typeclass that supports tuples and named tuples for vmap/zipvmap output typing.
core/src/main/scala/dimwit/package.scala	Replaces RemovePrimes match type with PrimeRemover typeclass for dropPrimes.
core/src/main/scala/dimwit/autodiff/Autodiff.scala	Replaces match-type Gradient/Hessian with open typeclass derivations supporting tuples, named tuples, and Products (via Mirror).

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[marcelluethi]

As this adds quite some complexity in the types I want to push back a little. The problem with the disjoint classes in match types could be solved by requiring the labels to be sealed traits. We can extend the Label macro such that it checks that the trait is really sealed by the user and otherwise returns an error. This would allow us to keep the same machinery that we had and even simplify some cases further. You can find here a draft:

https://github.com/dimwit-dev/dimwit/compare/main...marcelluethi:dimwit:labels-as-sealed-traits?expand=1

This does not solve the vmap and zipvmap issue. However, I am not entirely convinced that these really need to map over case classes and named tuples.

I am not strictly against the more advanced type machinery, but think we need to be very conscious about introducing more complexity in the type system.

[benikm91]

My current opinion is that requiring "sealed" in the user code is not worth the clarity in the type signature. I don't know how much Aux pattern affects compilation time or error-message clarity. However, let's reflect on this for a few days, as it is quite an important decision.

I tried to summarize the arguments for myself; feel free to extend. Important: at some points, an Aux pattern is necessary, and match types are possible (without downsides); I left this out of this summary. The summary is only in regard to type constructs on labels / labeled tensors. For example, we could design on a sealed trait, yet still implement zipvmap with Aux to support named tuples.

We have two solutions:

1. sealed traits with match types
2. traits with Aux pattern

---

+) clearer, more minimal type signature, not extra implicit
+) Easier mechanism
+) faster compilation (I think)
(-) requires sealed trait at user side) -> "sealed trait" are used for phantom types in Scala. Labels are Phantom types. Still required by user, yet, well motivated.
(-) sealed trait restrict possible design space, e.g., extendable label type hiearchies) -> still possible

-) more complex type signature with extra implicit
-) Harder mechanism
-) Slower compilation due to implicit resolution

---

Match types:

def swap[L1: Label, L2: Label](
    axis1: Axis[L1],
    axis2: Axis[L2]
)(using
    labels: Labels[T],
    axisIndex1: AxisIndex[T, L1],
    axisIndex2: AxisIndex[T, L2]
): Tensor[Swap[T, L1, L2], V] =  // return type clearer

// -- Mechanism --

type Swap[T <: Tuple, A, B] <: Tuple = T match
  case A *: tail => B *: Swap[tail, A, B]
  case B *: tail => A *: Swap[tail, A, B]
  case h *: tail => h *: Swap[tail, A, B]

Aux Pattern:

def swap[L1: Label, L2: Label](
      axis1: Axis[L1],
      axis2: Axis[L2]
  )(using
      labels: Labels[T],
      axisIndex1: AxisIndex[T, L1],
      axisIndex2: AxisIndex[T, L2],
      swapper: AxisSwapper[T, L1, L2] // Extra implicit 
  ): Tensor[swapper.Out, V] =

// -- Mechanism --

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

  given empty[L1, L2]: Aux[EmptyTuple, L1, L2, EmptyTuple] = instance
  
  trait AxisSwapperSecond extends AxisSwapperOther:
    given second[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 other[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
    

[marcelluethi]

yes, I think this needs some time to think about.

Just two more notes.

• Disjointness cannot only be achieved by sealed traits. Even a normal class would do it: e..g class A derives Label should be perfectly valid as well. As far as I understand it, it's just the "multiple inheritance" that causes problems
• Type hierarchies are still possible with the sealed trait approach. E.g. trait Data, sealed trait TrainingData extends Data should not cause any problems.

[benikm91]

sealed trait is also commonly used for Phantom types in Scala 3. Phantom types are types that only exist on the type-level, not on the value-level, i.e., our labels.

So actually the extra sealed is more precise, not an "extra" step...

https://www.codecentric.de/en/knowledge-hub/blog/phantom-types-scala

=> I updated the argument summary post.

=> Both disadvantages were weakened by this and our previous comment... Feels like sealed trait is the way to go to me

[benikm91]

I implemented the sealed trait variant with a push towards match types whenever possible: https://github.com/benikm91/dimwit/tree/sealed-type-fixes

However, one MAJOR problem arose with generic types:

// Limits with generics?
  def dontCompile[L1: Label, L2: Label, L3: Label](t: Tensor[(L1, L2, L3), Float32]): Tensor[(L1, L2), Float32] = t.sum(Axis[L3])

// we would have to write:
  def compilesFine[L1: Label, L2: Label, L3: Label](t: Tensor[(L1, L2, L3), Float32])(using
      axisIndex: AxisIndex[(L1, L2, L3), L3],
      labels: Labels[Remove[(L1, L2, L3), L3]]
  ): Tensor[Remove[(L1, L2, L3), L3], Float32] = t.sum(Axis[L3])

This makes writing generic functions difficult. I haven't found a solution for this so far. This is an essential feature for libraries (not examples) build on dimwit (e.g. deepwit).

Lets find a solution for this or go the type class road.

[marcelluethi]

After trying alternative routes with sealed traits and match types, I think we have to byte the bullet and embrace the Aux patterns used here.

If we anyway have this level of complexity, I am also fine with adding the type machinery that is required to extend zipvmap and vmap to work with named tuples.

I am therefore all for merging this PR and gain some experience with it. If it turns out that the type machinery is too heavy and the error messages confusing, we know at least where to improve.