diff --git a/AGENTS.md b/AGENTS.md index a5cb9ada..118e2c22 100644 --- a/AGENTS.md +++ b/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) ``` diff --git a/DOCKER.md b/DOCKER.md index 64aa4f23..6f148944 100644 --- a/DOCKER.md +++ b/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 diff --git a/Dockerfile b/Dockerfile index bb3210a9..e209e527 100644 --- a/Dockerfile +++ b/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 diff --git a/Dockerfile.ci b/Dockerfile.ci index 08071fdd..2ca3aae7 100644 --- a/Dockerfile.ci +++ b/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 diff --git a/README.md b/README.md index 70f3aac1..87e6cdd5 100644 --- a/README.md +++ b/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", diff --git a/core/src/main/scala/dimwit/jax/Einops.scala b/core/src/main/scala/dimwit/jax/Einops.scala deleted file mode 100644 index 18d8b607..00000000 --- a/core/src/main/scala/dimwit/jax/Einops.scala +++ /dev/null @@ -1,23 +0,0 @@ -package dimwit.jax - -import me.shadaj.scalapy.py -import me.shadaj.scalapy.py.SeqConverters - -object Einops: - - type PyAny = py.Any - type PyDynamic = py.Dynamic - - private val kwargUnpacker = py.eval("lambda f, x, p, k: f(x, p, **k)") - - def rearrange( - x: PyDynamic, - pattern: String, - kwargsMap: Map[String, Int] = Map() - ): PyDynamic = - val einops = py.module("einops") - val convertedSeq = kwargsMap.map { case (k, v: Int) => - k -> py.Any.from(v) - }.toSeq - val pyKwargs = py.Dynamic.global.dict(convertedSeq.toPythonProxy) - kwargUnpacker(einops.rearrange, x, pattern, pyKwargs).as[py.Dynamic] diff --git a/core/src/main/scala/dimwit/python/PythonSetup.scala b/core/src/main/scala/dimwit/python/PythonSetup.scala index cf19e6c6..c1a58ae7 100644 --- a/core/src/main/scala/dimwit/python/PythonSetup.scala +++ b/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 diff --git a/core/src/main/scala/dimwit/tensor/TensorEvidence.scala b/core/src/main/scala/dimwit/tensor/TensorEvidence.scala index 59f652d6..91803994 100644 --- a/core/src/main/scala/dimwit/tensor/TensorEvidence.scala +++ b/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 - ) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala index eaa96f2a..abad3d60 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala +++ b/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] diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala index 190bd437..a084c72a 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala @@ -13,105 +13,6 @@ class TensorOpsStructureSuite extends DimwitTest: ) ) - describe("rearrange (Einops-style)"): - - it("transpose: a b c -> c a b"): - val res = t3.rearrange((Axis[C], Axis[A], Axis[B])) - res.axes shouldBe List("C", "A", "B") - - it("merging axes: a b c -> (a b) c"): - // Merges 2x2 into 4. Result shape (4, 1) - val res = t3.rearrange((Axis[A |*| B], Axis[C])) - res.axes shouldBe List("A*B", "C") - res should approxEqual( - Tensor2(Axis[A |*| B], Axis[C]).fromArray( - Array(Array(1.0f), Array(2.0f), Array(3.0f), Array(4.0f)) - ) - ) - - describe("should have good errors"): - it("non-existing axis: a b c -> d a b"): - val wrongDimCode = "t3.rearrange((Axis[D], Axis[A], Axis[B]))" - val wrongDimErrors = typeCheckErrors(wrongDimCode) - wrongDimErrors should have size 1 - wrongDimErrors.head.message should include("Missing Axis: 'dimwit.D'") // Axis D does not exist in t3 - - it("unflatten without size fails: (a b) c -> a b c"): - val flattened = t3.rearrange((Axis[A |*| B], Axis[C])) - - val noDimsCode = "flattened.rearrange((Axis[A], Axis[B], Axis[C]))" - val noDimsErrors = typeCheckErrors(noDimsCode) - noDimsErrors should have size 1 - noDimsErrors.head.message should include("Missing Axis: 'dimwit.A'") // First Axis A is missing - - val aDimCode = "flattened.rearrange((Axis[A], Axis[B], Axis[C]), Axis[A] -> 2)" - val aDimErrors = typeCheckErrors(aDimCode) - aDimErrors should have size 1 - aDimErrors.head.message should include("Missing Axis: 'dimwit.B'") // Given A, next Axis B is missing - - "flattened.rearrange((Axis[A], Axis[B], Axis[C]), Axis[A] -> 2, Axis[B] -> 2)" should compile - - it("complex reshape fails incrementally: (a b) (c d) -> (a c) (b d)"): - val t = Tensor2(Axis[A |*| B], Axis[C |*| D]).fromArray( - Array( - Array(1.0f, 2.0f, 3.0f, 4.0f), - Array(5.0f, 6.0f, 7.0f, 8.0f), - Array(9.0f, 10.0f, 11.0f, 12.0f), - Array(13.0f, 14.0f, 15.0f, 16.0f) - ) - ) - - val noDimsCode = "t.rearrange((Axis[A |*| C], Axis[B |*| D]))" - val noDimsErrors = typeCheckErrors(noDimsCode) - noDimsErrors should have size 1 - noDimsErrors.head.message should include("Missing Axis: 'dimwit.A'") - - val aDimCode = "t.rearrange((Axis[A |*| C], Axis[B |*| D]), Axis[A] -> 2)" - val aDimErrors = typeCheckErrors(aDimCode) - aDimErrors should have size 1 - aDimErrors.head.message should include("Missing Axis: 'dimwit.C'") - - val acDimCode = "t.rearrange((Axis[A |*| C], Axis[B |*| D]), Axis[A] -> 2, Axis[C] -> 2)" - val acDimErrors = typeCheckErrors(acDimCode) - acDimErrors should have size 1 - acDimErrors.head.message should include("Missing Axis: 'dimwit.B'") - - val acbDimCode = "t.rearrange((Axis[A |*| C], Axis[B |*| D]), Axis[A] -> 2, Axis[C] -> 2, Axis[B] -> 2)" - val acbDimErrors = typeCheckErrors(acbDimCode) - acbDimErrors should have size 1 - acbDimErrors.head.message should include("Missing Axis: 'dimwit.D'") - - "t.rearrange((Axis[A |*| C], Axis[B |*| D]), Axis[A] -> 2, Axis[C] -> 2, Axis[B] -> 2, Axis[D] -> 2)" should compile - - it("complex rearrange: (a b) (c d) -> (a c) (b d)"): - val t = Tensor2(Axis[A |*| B], Axis[C |*| D]).fromArray( - Array( - Array(1.0f, 2.0f, 3.0f, 4.0f), - Array(5.0f, 6.0f, 7.0f, 8.0f), - Array(9.0f, 10.0f, 11.0f, 12.0f), - Array(13.0f, 14.0f, 15.0f, 16.0f) - ) - ) - t.axes shouldBe List("A*B", "C*D") - val res = t.rearrange((Axis[A |*| C], Axis[B |*| D]), Axis[A] -> 2, Axis[B] -> 2, Axis[C] -> 2, Axis[D] -> 2) - res.axes shouldBe List("A*C", "B*D") - res should approxEqual(Tensor2(Axis[A |*| C], Axis[B |*| D]).fromArray( - Array( - Array(1.0f, 2.0f, 5.0f, 6.0f), - Array(3.0f, 4.0f, 7.0f, 8.0f), - Array(9.0f, 10.0f, 13.0f, 14.0f), - Array(11.0f, 12.0f, 15.0f, 16.0f) - ) - )) - - it("reshaping axes: (a b) c -> a b c"): - val flattened = t3.rearrange((Axis[A |*| B], Axis[C])) - val res = flattened.rearrange( - (Axis[A], Axis[B], Axis[C]), - (Axis[A] -> 2, Axis[B] -> 2) - ) - res should approxEqual(t3) - describe("unflatten function"): it("unflatten axis into two axes"): @@ -402,25 +303,6 @@ class TensorOpsStructureSuite extends DimwitTest: describe("Concatenation"): - it("Prime axes are rearrangable"): - // As rearrange uses einops the "+" om the derived label for B |+| C must be handled in the rearrange operation to not trigger error - val t = Tensor2(Axis[A], Axis[Prime[B] |*| B]).fromArray( - Array(Array(1.0f, 2.0f, 3.0f, 4.0f), Array(5.0f, 6.0f, 7.0f, 8.0f)) - ) - val tRearranged = t.rearrange( - (Axis[B], Axis[Prime[B]], Axis[A]), - (Axis[B] -> 2, Axis[Prime[B]] -> 2) - ) - tRearranged.axes shouldBe List("B", "B'", "A") - - it("|+| axes are rearrangable"): - // As rearrange uses einops the "+" om the derived label for B |+| C must be handled in the rearrange operation to not trigger error - val t = Tensor2(Axis[A], Axis[B |+| C]).fromArray( - Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f)) - ) - val tRearranged = t.rearrange((Axis[B |+| C], Axis[A])) - tRearranged.axes shouldBe List("B+C", "A") - it("concatenate2 same axes"): val part1 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f))) val part2 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(3.0f, 4.0f))) diff --git a/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala b/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala index 1b912c4a..54ac9c3d 100644 --- a/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala +++ b/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala @@ -245,6 +245,17 @@ object VariationalAutoencoderExample: plt.imshow(toPyTensor(img2d), cmap = "gray") plt.show() + /** Arranges the images as a grid with 8 images per row, i.e. one image of rows * Height x 8 * Width pixels. */ + def toImageGrid(images: Tensor2[Prime[Height] |*| Prime[Width], Height |*| Width, Float32]): Tensor2[Prime[Height] |*| Height, Prime[Width] |*| Width, Float32] = + val columns = 8 + val numImages = images.shape(Axis[Prime[Height] |*| Prime[Width]]) + require(numImages % columns == 0, s"The number of images ($numImages) must be a multiple of $columns") + images + .unflatten(Axis[Prime[Height] |*| Prime[Width]], Shape(Axis[Prime[Height]] -> numImages / columns, Axis[Prime[Width]] -> columns)) + .unflatten(Axis[Height |*| Width], Shape(heightDim, widthDim)) + .flatten((Axis[Prime[Height]], Axis[Height])) + .flatten((Axis[Prime[Width]], Axis[Width])) + val vae = VariationalAutoencoder(trainedParams) /* Reconstructing images */ @@ -255,13 +266,7 @@ object VariationalAutoencoderExample: vae.decoder(latent) .relabel(Axis[TestSample].as(Axis[Prime[Height] |*| Prime[Width]])) - plotImg( - reconstructed - .rearrange( - (Axis[Prime[Height] |*| Height], Axis[Prime[Width] |*| Width]), - (Axis[Prime[Height]] -> 8, Axis[Prime[Width]] -> 8, heightDim, widthDim) - ) - ) + plotImg(toImageGrid(reconstructed)) /* Sampling from the latent space */ val stdNormal = Normal.standardNormal(Shape1(latentDim)) @@ -269,9 +274,4 @@ object VariationalAutoencoderExample: val z = stdNormal.sample(key) vae.decoder(z) - plotImg( - sampled.rearrange( - (Axis[Prime[Height] |*| Height], Axis[Prime[Width] |*| Width]), - (Axis[Prime[Height]] -> 8, Axis[Prime[Width]] -> 8, heightDim, widthDim) - ) - ) + plotImg(toImageGrid(sampled)) diff --git a/mdocs/README.md b/mdocs/README.md index b3fac12b..817d9551 100644 --- a/mdocs/README.md +++ b/mdocs/README.md @@ -117,7 +117,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`): @@ -127,7 +127,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",