Remove rearrange and einops dependency

AGENTS.md

@@ -750,8 +750,8 @@ val gridWrongValue = grid.updated(Axis[A].at(0))(grid)
 grid.select(Axis[A].at(3, 2))
 // java.lang.IllegalArgumentException: requirement failed: Window 3 until 5 is out of bounds for axis A of size 4
 // 	at scala.Predef$.require(Predef.scala:389)
-// 	at dimwit.tensor.tensorops.StructuralExtensions$.windowExtents(StructuralOps.scala:742)
-// 	at dimwit.tensor.tensorops.StructuralExtensions$.select(StructuralOps.scala:605)
+// 	at dimwit.tensor.tensorops.StructuralExtensions$.windowExtents(StructuralOps.scala:737)
+// 	at dimwit.tensor.tensorops.StructuralExtensions$.select(StructuralOps.scala:600)
 // 	at dimwit.package$.select(package.scala:80)
 // 	at repl.MdocSession$MdocApp17.$init$$$anonfun$3(AGENTS.md:723)
 ```

DOCKER.md

@@ -249,7 +249,6 @@ Runs on every push and pull request:
 **Python Packages**:
 - JAX (CPU version)
 - NumPy
-- Einops
 - Matplotlib, Pandas, Scikit-learn
 - Jupyter
 

Dockerfile

@@ -37,7 +37,7 @@ COPY . /workspace/
 
 # Create venv inheriting JAX from the base image, add extra packages
 RUN uv venv .venv --system-site-packages && \
-    uv pip install --python .venv/bin/python einops matplotlib pandas
+    uv pip install --python .venv/bin/python matplotlib pandas
 
 # Skip uv sync — JAX is already available via system-site-packages
 ENV DIMWIT_SKIP_SYNC=true

Dockerfile.ci

@@ -27,7 +27,7 @@ COPY project/ /workspace/project/
 
 # Set up Python venv with CPU JAX (pyproject.toml specifies cuda12; override here)
 RUN uv venv .venv --python python3.11 && \
-    uv pip install --python .venv/bin/python "jax[cpu]>=0.4" "einops>=0.8"
+    uv pip install --python .venv/bin/python "jax[cpu]>=0.4"
 
 # Skip uv sync at runtime and point directly at the venv Python
 # This also prevents uv from being called at sbt load time

README.md

@@ -100,7 +100,7 @@ lazy val myProject = (project in file("."))
 
 ### Python Environment Setup
 
-DimWit requires **Python 3.9+** and **JAX** since it uses JAX as the backend for tensor operations via ScalaPy. It also relies on **Einops** for tensor reshaping and manipulation.
+DimWit requires **Python 3.9+** and **JAX** since it uses JAX as the backend for tensor operations via ScalaPy.
 
 The easiest way to set up the Python environment is to install the [uv package manager](https://github.com/astral-sh/uv) and add a `pyproject.toml` to your **project root** (the directory where you run `sbt`):
 
@@ -110,7 +110,6 @@ name = "my-project-python-env"
 version = "0.1.0"
 requires-python = "==3.13.*"
 dependencies = [
-    "einops>=0.8.1",
     # Use jax[cpu] for CPU-only environments (e.g. development/CI without a GPU)
     # Use jax[cuda12] for NVIDIA GPU support
     "jax[cpu]>=0.8.2",

core/src/main/scala/dimwit/python/PythonSetup.scala

@@ -35,7 +35,6 @@ object PythonSetup:
             |  requires-python = ">=3.11"
             |  dependencies = [
             |      "jax[cpu]>=0.4",      # or jax[cuda12], jax[tpu]
-            |      "einops>=0.8",
             |  ]
             |
             |Set DIMWIT_SKIP_SYNC=true to manage the Python environment yourself.""".stripMargin

core/src/main/scala/dimwit/tensor/TensorEvidence.scala

@@ -1,41 +1,14 @@
 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.
+  * actually has: is the new order a permutation of the old one?
   */
 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]
 
@@ -46,74 +19,3 @@ object TensorEvidence:
         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
-      )

core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala

@@ -1,7 +1,6 @@
 package dimwit.tensor.tensorops
 
 import dimwit.python.PyIndex.itemAt
-import dimwit.jax.Einops
 import dimwit.jax.Jax
 import dimwit.tensor.Axis
 import dimwit.tensor.AxisAtIndex
@@ -25,15 +24,11 @@ import dimwit.tensor.ShapeTypeHelpers.AxisReplacerAll
 import dimwit.tensor.ShapeTypeHelpers.DimExtractor
 import dimwit.tensor.ShapeTypeHelpers.MergeLabels
 import dimwit.tensor.ShapeTypeHelpers.UnwrapAxes
-import dimwit.tensor.ShapeTypeHelpers.UnwrapDims
 import dimwit.tensor.Tensor
 import dimwit.tensor.Tensor1
 import dimwit.tensor.TupleHelpers
 import dimwit.tensor.TupleHelpers.StrictSubset
-import dimwit.tensor.TensorEvidence.CheckValid
-import dimwit.tensor.TensorEvidence.ComputeMissing
 import dimwit.tensor.TensorEvidence.IsPermutation
-import dimwit.tensor.TensorEvidence.ValidationResult
 import dimwit.tensor.Broadcast3
 import dimwit.|+|
 import me.shadaj.scalapy.py
@@ -320,7 +315,7 @@ private[dimwit] object StructuralOps:
           val headTensor = Tensor[NewShape, V](currentArr)(using newLabelsWitness)
           headTensor *: tailMaker(arrays.tail, compLabels.tail, originalLabels, splitIndex)
 
-/** Extension methods for structural operations, e.g. `t.transpose`, `t.select(...)` or `t.rearrange(...)`. */
+/** Extension methods for structural operations, e.g. `t.transpose`, `t.select(...)` or `t.flatten(...)`. */
 private[dimwit] object StructuralExtensions:
 
   import StructuralOps.*
@@ -467,7 +462,7 @@ private[dimwit] object StructuralExtensions:
       * @param NewOrder A tuple representing the new order of axes for the tensor.
       * @return A new tensor with the axes transposed according to the specified order.
       */
-    def transpose[NewOrder <: Tuple, Status <: ValidationResult](newOrder: NewOrder)(using
+    def transpose[NewOrder <: Tuple](newOrder: NewOrder)(using
         ev: AxisIndices[T, UnwrapAxes[NewOrder]],
         newLabels: Labels[UnwrapAxes[NewOrder]]
     )(using
@@ -751,81 +746,6 @@ private[dimwit] object StructuralExtensions:
       val starts = Seq.tabulate[py.Any](tensor.shape.rank)(i => if i == dimIndex then start else py.Any.from(0))
       Jax.Dynamic.global.tuple(starts.toPythonProxy)
 
-    def rearrange[Axes <: Tuple, Status <: ValidationResult](newOrder: Axes)(using
-        Labels[UnwrapAxes[Axes]]
-    )(using
-        computer: ComputeMissing[UnwrapAxes[Axes], T, EmptyTuple, Status],
-        guard: CheckValid[Status]
-    ): Tensor[UnwrapAxes[Axes], V] =
-      rearrange[Axes, EmptyTuple, Status](newOrder, EmptyTuple)
-
-    // Convenience overload for 1 dims (to support error messages with single axis)
-    inline def rearrange[Axes <: Tuple, L1, Status <: ValidationResult](newOrder: Axes, d1: AxisExtent[L1])(using computer: ComputeMissing[UnwrapAxes[Axes], T, UnwrapDims[Tuple1[AxisExtent[L1]]], Status], guard: CheckValid[Status])(using newLabels: Labels[UnwrapAxes[Axes]], extractor: DimExtractor[Tuple1[AxisExtent[L1]]]): Tensor[UnwrapAxes[Axes], V] =
-      rearrange(newOrder, Tuple1(d1))
-
-    // Convenience overload for 2 dims
-    inline def rearrange[Axes <: Tuple, L1, L2, Status <: ValidationResult](newOrder: Axes, d1: AxisExtent[L1], d2: AxisExtent[L2])(using computer: ComputeMissing[UnwrapAxes[Axes], T, UnwrapDims[(AxisExtent[L1], AxisExtent[L2])], Status], guard: CheckValid[Status])(using newLabels: Labels[UnwrapAxes[Axes]], extractor: DimExtractor[(AxisExtent[L1], AxisExtent[L2])]): Tensor[UnwrapAxes[Axes], V] =
-      rearrange(newOrder, (d1, d2))
-
-    // Convenience overload for 3 dims
-    inline def rearrange[Axes <: Tuple, L1, L2, L3, Status <: ValidationResult](newOrder: Axes, d1: AxisExtent[L1], d2: AxisExtent[L2], d3: AxisExtent[L3])(using computer: ComputeMissing[UnwrapAxes[Axes], T, UnwrapDims[(AxisExtent[L1], AxisExtent[L2], AxisExtent[L3])], Status], guard: CheckValid[Status])(using newLabels: Labels[UnwrapAxes[Axes]], extractor: DimExtractor[(AxisExtent[L1], AxisExtent[L2], AxisExtent[L3])]): Tensor[UnwrapAxes[Axes], V] =
-      rearrange(newOrder, (d1, d2, d3))
-
-    // Convenience overload for 4 dims
-    inline def rearrange[Axes <: Tuple, L1, L2, L3, L4, Status <: ValidationResult](newOrder: Axes, d1: AxisExtent[L1], d2: AxisExtent[L2], d3: AxisExtent[L3], d4: AxisExtent[L4])(using computer: ComputeMissing[UnwrapAxes[Axes], T, UnwrapDims[(AxisExtent[L1], AxisExtent[L2], AxisExtent[L3], AxisExtent[L4])], Status], guard: CheckValid[Status])(using newLabels: Labels[UnwrapAxes[Axes]], extractor: DimExtractor[(AxisExtent[L1], AxisExtent[L2], AxisExtent[L3], AxisExtent[L4])]): Tensor[UnwrapAxes[Axes], V] =
-      rearrange(newOrder, (d1, d2, d3, d4))
-
-    def rearrange[Axes <: Tuple, Dims <: Tuple, Status <: ValidationResult](
-        newOrder: Axes,
-        dims: Dims
-    )(using
-        computer: ComputeMissing[UnwrapAxes[Axes], T, UnwrapDims[Dims], Status],
-        guard: CheckValid[Status]
-    )(using
-        newLabels: Labels[UnwrapAxes[Axes]],
-        extractor: DimExtractor[Dims]
-    ): Tensor[UnwrapAxes[Axes], V] =
-      def cleanPatternPrime(pattern: String): String =
-        // Support dimwit.Prime by replacing ' with "Prime"
-        pattern.replaceAll(
-          "'",
-          "Prime"
-        )
-      def createEinopsPattern(fromPattern: String, toPattern: String): String =
-        def cleanPatternStar(pattern: String): String =
-          // to replace all a*b*c in pattern with (a b c), example:
-          // "a*b*c d e f*g h" -> "(a b c) d e (f g) h"
-          val regex = raw"([a-zA-Z0-9_]+(\*[a-zA-Z0-9_]+)+)".r
-          regex.replaceAllIn(
-            pattern,
-            _.group(1).split("\\*").mkString("(", " ", ")")
-          )
-        def cleanPatternPlus(pattern: String): String =
-          // Support dimwit.|+| by replacing + with underlines
-          val regex = raw"([a-zA-Z0-9_]+(\+[a-zA-Z0-9_]+)+)".r
-          regex.replaceAllIn(
-            pattern,
-            _.group(1).replace("+", "_")
-          )
-        def cleanPattern(pattern: String): String =
-          cleanPatternPlus(cleanPatternStar(cleanPatternPrime(pattern)))
-        s"${cleanPattern(fromPattern)} -> ${cleanPattern(toPattern)}"
-      val fromPattern = tensor.shape.labels.mkString(" ")
-      val toPattern = newLabels.names.mkString(" ")
-      val pattern = createEinopsPattern(fromPattern, toPattern)
-      val dimSizesMap = extractor.extract(dims)
-      val cleanDimSizesMap = dimSizesMap.map { case (k, v) =>
-        val newKey = cleanPatternPrime(k)
-        (newKey, v)
-      }
-      Tensor(
-        Einops.rearrange(
-          tensor.jaxValue,
-          pattern,
-          kwargsMap = cleanDimSizesMap
-        )
-      )
-
     def broadcastTo[O <: Tuple: Labels](newShape: Shape[O])(using
         labels: Labels[T],
         ev: StrictSubset[T, O]