From 6961cd1146ea474cfb640d334d95c41740c34d38 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Tue, 8 Sep 2026 14:17:06 +0200 Subject: [PATCH 1/4] Add sharding support for data-parallelism --- build.sbt | 2 + core/src/main/resources/python/jax_helper.py | 17 ++ core/src/main/scala/dimwit/package.scala | 5 + .../src/main/scala/dimwit/sharding/Mesh.scala | 65 +++++++ .../main/scala/dimwit/sharding/MeshAxis.scala | 31 ++++ .../scala/dimwit/sharding/MeshLabel.scala | 60 +++++++ .../main/scala/dimwit/sharding/Sharded.scala | 17 ++ .../scala/dimwit/sharding/ShardingOps.scala | 40 +++++ .../scala/dimwit/sharding/ShardingSuite.scala | 165 ++++++++++++++++++ .../basic/ShardedLinearRegression.scala | 85 +++++++++ .../main/scala/dimwit/basic/ShardedSum.scala | 60 +++++++ 11 files changed, 547 insertions(+) create mode 100644 core/src/main/scala/dimwit/sharding/Mesh.scala create mode 100644 core/src/main/scala/dimwit/sharding/MeshAxis.scala create mode 100644 core/src/main/scala/dimwit/sharding/MeshLabel.scala create mode 100644 core/src/main/scala/dimwit/sharding/Sharded.scala create mode 100644 core/src/main/scala/dimwit/sharding/ShardingOps.scala create mode 100644 core/src/test/scala/dimwit/sharding/ShardingSuite.scala create mode 100644 examples/src/main/scala/dimwit/basic/ShardedLinearRegression.scala create mode 100644 examples/src/main/scala/dimwit/basic/ShardedSum.scala diff --git a/build.sbt b/build.sbt index fb813c7..fcb3d61 100644 --- a/build.sbt +++ b/build.sbt @@ -84,6 +84,8 @@ lazy val core = (project in file("core")) fork := true, javaOptions ++= scalapyJavaOptions, Test / envVars += "DIMWIT_SKIP_SYNC" -> "true", + // The sharding tests need a mesh of devices; a default JAX process reports a single CPU device. + Test / envVars += "XLA_FLAGS" -> "--xla_force_host_platform_device_count=8", coverageMinimumStmtTotal := 80, coverageFailOnMinimum := false, coverageHighlighting := true, diff --git a/core/src/main/resources/python/jax_helper.py b/core/src/main/resources/python/jax_helper.py index fe1b761..f3b4136 100644 --- a/core/src/main/resources/python/jax_helper.py +++ b/core/src/main/resources/python/jax_helper.py @@ -85,3 +85,20 @@ def jit(f): def jit_fn(f, jit_kwargs=None): return wrap(jax.jit, f, kwargs=jit_kwargs) + +# --- sharding ------------------------------------------------------------- + +def device_mesh(devices, shape, axis_names): + """Builds a jax.sharding.Mesh from a flat, row-major list of devices.""" + import numpy as np + grid = np.array(list(devices), dtype=object).reshape(tuple(shape)) + return jax.sharding.Mesh(grid, axis_names=tuple(axis_names)) + +def named_sharding(mesh, rank, axis_index, mesh_axis_name): + """NamedSharding splitting axis `axis_index` of a rank-`rank` array over `mesh_axis_name`. + + Every other axis is replicated. + """ + spec = [None] * rank + spec[axis_index] = mesh_axis_name + return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*spec)) diff --git a/core/src/main/scala/dimwit/package.scala b/core/src/main/scala/dimwit/package.scala index 555ecde..c937b04 100644 --- a/core/src/main/scala/dimwit/package.scala +++ b/core/src/main/scala/dimwit/package.scala @@ -62,6 +62,11 @@ package object dimwit: // Export the Prime axis marker and the type classes that manipulate it export dimwit.prime.{Prime, PrimeRemover, PrimeRest, PrimeConcat} + // Export the sharding types: mesh labels, meshes, and the sharded axis marker + export dimwit.sharding.{Mesh, Mesh1, Mesh2, Mesh3, MeshAxis, MeshAxisExtent, MeshAxisIndex, MeshLabel, MeshLabels} + export dimwit.sharding.`|@|` + export dimwit.sharding.ShardingOps.* + // Export operations export dimwit.tensor.TensorOps.* export dimwit.linalg.LinearAlgebra.{VectorNormType, MatrixNormType, QRMode} diff --git a/core/src/main/scala/dimwit/sharding/Mesh.scala b/core/src/main/scala/dimwit/sharding/Mesh.scala new file mode 100644 index 0000000..fc254b1 --- /dev/null +++ b/core/src/main/scala/dimwit/sharding/Mesh.scala @@ -0,0 +1,65 @@ +package dimwit.sharding + +import dimwit.hardware.Device +import dimwit.jax.Jax +import me.shadaj.scalapy.py.SeqConverters + +/** A grid of devices, binding each mesh axis label of `M` to a number of devices. */ +final class Mesh[M <: Tuple: MeshLabels] private[sharding] ( + val axisSizes: List[Int], + val devices: Seq[Device] +): + + lazy val axisNames: List[String] = MeshLabels[M].names + + def size: Int = axisSizes.product + + def sizeOf[A](meshAxis: MeshAxis[A])(using ev: MeshAxisIndex[M, A]): Int = axisSizes(ev.index) + + private[dimwit] lazy val jaxMesh: Jax.PyDynamic = + Jax.jax_helper + .device_mesh(devices.map(_.toJaxDevice).toPythonProxy, axisSizes.toPythonProxy, axisNames.toPythonProxy) + .as[Jax.PyDynamic] + + override def toString: String = + axisNames.zip(axisSizes).map((name, size) => s"$name -> $size").mkString("Mesh(", ", ", ")") + +object Mesh: + + private[sharding] type ExtractLabels[Extents <: Tuple] <: Tuple = Extents match + case EmptyTuple => EmptyTuple + case MeshAxisExtent[a] *: tail => a *: ExtractLabels[tail] + + def apply[A: MeshLabel](extent: MeshAxisExtent[A]): Mesh[Tuple1[A]] = fromTuple(Tuple1(extent)) + + def apply[Extents <: Tuple](extents: Extents)(using MeshLabels[ExtractLabels[Extents]]): Mesh[ExtractLabels[Extents]] = + fromTuple(extents) + + def fromTuple[Extents <: Tuple](extents: Extents)(using labels: MeshLabels[ExtractLabels[Extents]]): Mesh[ExtractLabels[Extents]] = + val sizes = extents.toList.collect: + case extent: MeshAxisExtent[?] => extent.size + val required = sizes.product + val available = Jax.devices + require( + available.size >= required, + s"Mesh ${labels.names.zip(sizes).map((name, size) => s"$name -> $size").mkString("(", ", ", ")")} requires $required devices, but only ${available.size} are available" + ) + new Mesh(sizes, available.take(required)) + +type Mesh1[A] = Mesh[Tuple1[A]] +type Mesh2[A, B] = Mesh[(A, B)] +type Mesh3[A, B, C] = Mesh[(A, B, C)] + +object Mesh1: + def apply[A: MeshLabel](extent: MeshAxisExtent[A]): Mesh1[A] = Mesh(extent) + +object Mesh2: + def apply[A: MeshLabel, B: MeshLabel](extent1: MeshAxisExtent[A], extent2: MeshAxisExtent[B]): Mesh2[A, B] = + Mesh.fromTuple((extent1, extent2)) + +object Mesh3: + def apply[A: MeshLabel, B: MeshLabel, C: MeshLabel]( + extent1: MeshAxisExtent[A], + extent2: MeshAxisExtent[B], + extent3: MeshAxisExtent[C] + ): Mesh3[A, B, C] = Mesh.fromTuple((extent1, extent2, extent3)) diff --git a/core/src/main/scala/dimwit/sharding/MeshAxis.scala b/core/src/main/scala/dimwit/sharding/MeshAxis.scala new file mode 100644 index 0000000..fef34bb --- /dev/null +++ b/core/src/main/scala/dimwit/sharding/MeshAxis.scala @@ -0,0 +1,31 @@ +package dimwit.sharding + +import scala.annotation.implicitNotFound + +/** Represents an axis of a device [[Mesh]]. The mesh-side counterpart of [[dimwit.tensor.Axis]]. */ +final class MeshAxis[A: MeshLabel]: + + def name: String = summon[MeshLabel[A]].name + + def extent(size: Int): MeshAxisExtent[A] = MeshAxisExtent(this, size) + def ->(size: Int): MeshAxisExtent[A] = this.extent(size) + + override def toString: String = s"MeshAxis($name)" + +/** A mesh axis together with the number of devices along it. */ +case class MeshAxisExtent[A: MeshLabel](axis: MeshAxis[A], size: Int) + +/** Finds the position of a mesh axis in a mesh. */ +@implicitNotFound("MeshAxis[${A}] not found in Mesh[${M}]") +trait MeshAxisIndex[M <: Tuple, A]: + def index: Int + +object MeshAxisIndex: + + def apply[M <: Tuple, A](using idx: MeshAxisIndex[M, A]): Int = idx.index + + given found[A, Tail <: Tuple]: MeshAxisIndex[A *: Tail, A] with + val index = 0 + + given search[H, T <: Tuple, A](using next: MeshAxisIndex[T, A]): MeshAxisIndex[H *: T, A] with + val index = 1 + next.index diff --git a/core/src/main/scala/dimwit/sharding/MeshLabel.scala b/core/src/main/scala/dimwit/sharding/MeshLabel.scala new file mode 100644 index 0000000..6df7736 --- /dev/null +++ b/core/src/main/scala/dimwit/sharding/MeshLabel.scala @@ -0,0 +1,60 @@ +package dimwit.sharding + +import scala.quoted.* + +/** A label for an axis of a device [[Mesh]]. + * + * Mesh labels are a separate kind from the data axis labels carried by + * `dimwit.tensor.Label`: a type that `derives Label` has no `MeshLabel` instance and + * vice versa, so a data axis label cannot be used where a mesh label is required. + * + * {{{ + * trait X derives MeshLabel + * }}} + */ +@scala.annotation.implicitNotFound(""" +A mesh axis label ${T} was given or inferred, which does not have a MeshLabel instance. +Mesh axis labels are a different kind than data axis labels: a type declared with +'derives Label' cannot be used as a mesh axis label. +Ensure that all mesh axis types ${T} are defined with 'derives MeshLabel' (e.g. 'trait X derives MeshLabel') +""") +trait MeshLabel[T]: + def name: String + +object MeshLabel: + + def apply[T](using meshLabel: MeshLabel[T]): MeshLabel[T] = meshLabel + + inline def derived[T]: MeshLabel[T] = ${ derivedMacro[T] } + + private def derivedMacro[T: Type](using Quotes): Expr[MeshLabel[T]] = + import quotes.reflect.* + val tpe = TypeRepr.of[T] + val simpleName = tpe.typeSymbol.name + '{ + new MeshLabel[T]: + def name: String = ${ Expr(simpleName) } + } + +@scala.annotation.implicitNotFound(""" +A tuple of mesh axis labels ${T} was given or inferred that does not have a valid MeshLabels instance. + +Ensure that all of the types in the tuple have a 'derives MeshLabel' clause. +""") +trait MeshLabels[T]: + def names: List[String] + +private class MeshLabelsImpl[T](val names: List[String]) extends MeshLabels[T] + +object MeshLabels: + + def apply[T](using labels: MeshLabels[T]): MeshLabels[T] = labels + + given emptyTuple: MeshLabels[EmptyTuple] = new MeshLabelsImpl[EmptyTuple](Nil) + + given lift[A](using v: MeshLabel[A]): MeshLabels[A] = new MeshLabelsImpl[A](List(v.name)) + + given consTuple[H, T <: Tuple](using + head: MeshLabel[H], + tail: MeshLabels[T] + ): MeshLabels[H *: T] = new MeshLabelsImpl[H *: T](head.name :: tail.names) diff --git a/core/src/main/scala/dimwit/sharding/Sharded.scala b/core/src/main/scala/dimwit/sharding/Sharded.scala new file mode 100644 index 0000000..5828b91 --- /dev/null +++ b/core/src/main/scala/dimwit/sharding/Sharded.scala @@ -0,0 +1,17 @@ +package dimwit.sharding + +import dimwit.tensor.Label + +/** Marks the data axis `A` as sharded over the mesh axis `M`, as in `Tensor2[Batch |@| X, Feature, Float32]`. + * + * `Batch |@| X` is just another axis label, so every operation applies to a sharded tensor + * unchanged: reducing it all-reduces, reducing or mapping any other axis stays local, and it + * cannot be combined with an unsharded `Batch` because the names differ. + * + * Spelled `|@|` alongside [[dimwit.|*|]] and [[dimwit.|+|]]; a bare `@` is Scala's annotation syntax. + */ +infix trait |@|[A, M] + +object `|@|`: + given [A, M](using label: Label[A], meshLabel: MeshLabel[M]): Label[A |@| M] with + val name: String = s"${label.name}@${meshLabel.name}" diff --git a/core/src/main/scala/dimwit/sharding/ShardingOps.scala b/core/src/main/scala/dimwit/sharding/ShardingOps.scala new file mode 100644 index 0000000..6bf036b --- /dev/null +++ b/core/src/main/scala/dimwit/sharding/ShardingOps.scala @@ -0,0 +1,40 @@ +package dimwit.sharding + +import dimwit.jax.Jax +import dimwit.tensor.Axis +import dimwit.tensor.Label +import dimwit.tensor.Labels +import dimwit.tensor.ShapeTypeHelpers.AxisReplacer +import dimwit.tensor.Tensor + +object ShardingOps: + + extension [T <: Tuple: Labels, V](t: Tensor[T, V]) + + /** Splits `t` across `mesh` along one axis, rewriting that axis from `L` to `L |@| A`. + * + * The shards are placed one per device with `jax.device_put` under a `NamedSharding`. + * + * {{{ + * val mesh = Mesh(MeshAxis[X] -> 4) + * val sharded: Tensor2[Batch |@| X, Feature, Float32] = t.shard(mesh, Axis[Batch] -> MeshAxis[X]) + * }}} + * + * @throws IllegalArgumentException if the extent of `L` is not divisible by the mesh axis size. + */ + def shard[M <: Tuple, L: Label, A: MeshLabel](mesh: Mesh[M], mapping: (Axis[L], MeshAxis[A]))(using + replacer: AxisReplacer[T, L, L |@| A], + meshIndex: MeshAxisIndex[M, A], + labels: Labels[replacer.NewShape] + ): Tensor[replacer.NewShape, V] = + val axisName = summon[Label[L]].name + val meshAxisName = mapping._2.name + val extent = t.shape.dimensions(replacer.index) + val meshSize = mesh.axisSizes(meshIndex.index) + if extent % meshSize != 0 then + throw new IllegalArgumentException( + s"Cannot shard axis $axisName of extent $extent over mesh axis $meshAxisName of size $meshSize: " + + s"$extent is not divisible by $meshSize." + ) + val sharding = Jax.jax_helper.named_sharding(mesh.jaxMesh, t.shape.rank, replacer.index, meshAxisName) + Tensor[replacer.NewShape, V](Jax.device_put(t.jaxValue, sharding.as[Jax.PyDynamic])) diff --git a/core/src/test/scala/dimwit/sharding/ShardingSuite.scala b/core/src/test/scala/dimwit/sharding/ShardingSuite.scala new file mode 100644 index 0000000..0a736ab --- /dev/null +++ b/core/src/test/scala/dimwit/sharding/ShardingSuite.scala @@ -0,0 +1,165 @@ +package dimwit.sharding + +import dimwit.* +import dimwit.hardware.DeviceBackend +import me.shadaj.scalapy.py + +import scala.compiletime.testing.typeCheckErrors + +/** Mesh axis labels: a different kind than the data axis labels `A`, `B`, ... in the test package. */ +trait X derives MeshLabel +trait Y derives MeshLabel + +class ShardingSuite extends DimwitTest: + + private val BatchExtent = 8 + private val FeatureExtent = 4 + private val MeshExtent = 4 + + private def tensorOf(batch: Int, feature: Int = FeatureExtent): Tensor2[A, B, Float32] = + Tensor2(Axis[A], Axis[B]).fromArray(Array.tabulate(batch, feature)((i, j) => (i * feature + j).toFloat)) + + private val t: Tensor2[A, B, Float32] = tensorOf(BatchExtent) + + private def enoughDevices(): Unit = + assume( + DeviceBackend.CPU.devices.size >= MeshExtent, + s"needs at least $MeshExtent devices, found ${DeviceBackend.CPU.devices.size}" + ) + + private def isFullyReplicated[T <: Tuple: Labels, V](tensor: Tensor[T, V]): Boolean = + tensor.jaxValue.sharding.is_fully_replicated.as[Boolean] + + private def deviceCount[T <: Tuple: Labels, V](tensor: Tensor[T, V]): Int = + tensor.jaxValue.sharding.num_devices.as[Int] + + private def mesh1 = Mesh1(MeshAxis[X] -> MeshExtent) + + private def sharded: Tensor2[A |@| X, B, Float32] = t.shard(mesh1, Axis[A] -> MeshAxis[X]) + + describe("Mesh"): + + it("binds mesh axis labels to a grid of devices"): + enoughDevices() + val mesh = mesh1 + mesh.size shouldBe MeshExtent + mesh.axisNames shouldBe List("X") + mesh.axisSizes shouldBe List(MeshExtent) + mesh.sizeOf(MeshAxis[X]) shouldBe MeshExtent + mesh.devices.map(_.id).distinct should have size MeshExtent + + it("is built from a tuple of extents, with Mesh1/Mesh2 as special cases"): + enoughDevices() + val single: Mesh[Tuple1[X]] = Mesh(MeshAxis[X] -> MeshExtent) + single.axisNames shouldBe List("X") + single.axisSizes shouldBe List(MeshExtent) + + val grid: Mesh[(X, Y)] = Mesh((MeshAxis[X] -> 2, MeshAxis[Y] -> 2)) + grid.axisNames shouldBe List("X", "Y") + grid.size shouldBe 4 + grid.sizeOf(MeshAxis[Y]) shouldBe 2 + + Mesh1(MeshAxis[X] -> MeshExtent).axisSizes shouldBe single.axisSizes + Mesh2(MeshAxis[X] -> 2, MeshAxis[Y] -> 2).axisNames shouldBe grid.axisNames + + it("fails with a clear error when there are not enough devices"): + val error = intercept[IllegalArgumentException](Mesh1(MeshAxis[X] -> 1000000)) + error.getMessage should include("requires 1000000 devices") + + describe("shard"): + + it("rewrites the axis type and places the shards on the mesh devices"): + enoughDevices() + val result: Tensor2[A |@| X, B, Float32] = sharded + result.axes shouldBe List("A@X", "B") + result.shape.dimensions shouldBe List(BatchExtent, FeatureExtent) + deviceCount(result) shouldBe MeshExtent + isFullyReplicated(result) shouldBe false + result shouldEqual t + + it("rejects an axis whose extent does not divide over the mesh axis"): + enoughDevices() + val error = intercept[IllegalArgumentException](tensorOf(7).shard(mesh1, Axis[A] -> MeshAxis[X])) + error.getMessage shouldBe "Cannot shard axis A of extent 7 over mesh axis X of size 4: 7 is not divisible by 4." + + describe("a sharded axis is just another axis"): + + it("reducing the sharded axis all-reduces and equals the unsharded sum, bit for bit"): + enoughDevices() + val result: Tensor1[B, Float32] = sharded.sum(Axis[A |@| X]) + result shouldEqual t.sum(Axis[A]) + result.axes shouldBe List("B") + isFullyReplicated(result) shouldBe true + + it("reducing every axis works on both, and is how one loss function serves both runs"): + enoughDevices() + sharded.sum shouldEqual t.sum + sharded.mean shouldEqual t.mean + + it("reducing another axis is local and keeps the mesh annotation"): + enoughDevices() + val result: Tensor1[A |@| X, Float32] = sharded.sum(Axis[B]) + result shouldEqual t.sum(Axis[B]) + result.axes shouldBe List("A@X") + isFullyReplicated(result) shouldBe false + deviceCount(result) shouldBe MeshExtent + + it("other reductions over the sharded axis work too, with no extra support"): + enoughDevices() + sharded.mean(Axis[A |@| X]) shouldEqual t.mean(Axis[A]) + sharded.max(Axis[A |@| X]) shouldEqual t.max(Axis[A]) + sharded.min(Axis[A |@| X]) shouldEqual t.min(Axis[A]) + + it("vmap maps over the sharded axis, which stays sharded"): + enoughDevices() + val result: Tensor2[A |@| X, B, Float32] = sharded.vmap(Axis[A |@| X])(row => row *! Tensor0(2.0f)) + result shouldEqual t.vmap(Axis[A])(row => row *! Tensor0(2.0f)) + result.axes shouldBe List("A@X", "B") + deviceCount(result) shouldBe MeshExtent + + it("contracts a replicated tensor against the unsharded axis, staying sharded"): + enoughDevices() + val weights = Tensor1(Axis[B]).fromArray(Array.fill(FeatureExtent)(2.0f)) + val result: Tensor1[A |@| X, Float32] = sharded.dot(Axis[B])(weights) + result shouldEqual t.dot(Axis[B])(weights) + deviceCount(result) shouldBe MeshExtent + + it("broadcasts a replicated tensor over the axes it does share"): + enoughDevices() + val bias = Tensor1(Axis[B]).fromArray(Array.fill(FeatureExtent)(1.0f)) + val result: Tensor2[A |@| X, B, Float32] = sharded -! bias + result shouldEqual (t -! bias) + deviceCount(result) shouldBe MeshExtent + + describe("sharded and unsharded tensors do not mix"): + + it("a sharded tensor cannot be combined with an unsharded one"): + enoughDevices() + val errors = typeCheckErrors("sharded + t") + errors should not be empty + + it("the sharded axis cannot be named without its mesh axis"): + enoughDevices() + val errors = typeCheckErrors("sharded.sum(Axis[A])") + errors should not be empty + errors.head.message should include("not found in Tensor") + + describe("mesh labels are a separate kind from data axis labels"): + + it("a data axis label cannot be used as a mesh axis label"): + val errors = typeCheckErrors("MeshAxis[A]") + errors should not be empty + errors.head.message should include("MeshLabel") + + it("a data axis label cannot be used to build a mesh"): + val errors = typeCheckErrors("Mesh1(MeshAxis[A] -> 4)") + errors should not be empty + + it("a mesh axis label cannot be used as a data axis label"): + val errors = typeCheckErrors("Axis[X]") + errors should not be empty + errors.head.message should include("Label") + + it("a data axis label cannot stand on the mesh side of |@|"): + val errors = typeCheckErrors("summon[Label[A |@| B]]") + errors should not be empty diff --git a/examples/src/main/scala/dimwit/basic/ShardedLinearRegression.scala b/examples/src/main/scala/dimwit/basic/ShardedLinearRegression.scala new file mode 100644 index 0000000..801bcd5 --- /dev/null +++ b/examples/src/main/scala/dimwit/basic/ShardedLinearRegression.scala @@ -0,0 +1,85 @@ +package dimwit.examples.basic + +import dimwit.Conversions.given +import dimwit.* +import dimwit.autodiff.* +import dimwit.jax.Jax +import dimwit.optimizer.GradientDescent +import dimwit.random.Random +import dimwit.stats.Normal + +object ShardedLinearRegression: + + val Devices = 4 + val BatchSize = 8192 // must be divisible by Devices + val FeatureSize = 64 + val Steps = 200 + val LearningRate = 0.05f + + trait Batch derives Label + trait Feature derives Label + + case class Params(weights: Tensor1[Feature, Float32], bias: Tensor0[Float32]) + object Params: + def init: Params = Params(Tensor1(Axis[Feature]).fromArray(Array.fill(FeatureSize)(0f)), Tensor0(0f)) + + class LinearRegression(params: Params) extends (Tensor1[Feature, Float32] => Tensor0[Float32]): + def apply(x: Tensor1[Feature, Float32]): Tensor0[Float32] = x.dot(Axis[Feature])(params.weights) + params.bias + + @main def runSingleDevice(): Unit = + dimwit.initialize() + + val (x, y) = dataset(Random.Key(0)) + + def loss(params: Params): Tensor0[Float32] = + val model = LinearRegression(params) + val predictions = x.vmap(Axis[Batch])(model) + // `sum` names no axis, so this line is identical in the sharded run. + (predictions - y).pow(2f).sum / Tensor0(BatchSize.toFloat) + + val descent = GradientDescent(Tensor0(LearningRate)).iterate(Params.init)(grad(jit(loss))) + descent.next().bias.item // warm up JIT + val start = System.nanoTime() + val params = descent.drop(Steps - 1).next() + params.bias.item + report("single device", params, (System.nanoTime() - start) / 1e9) + + @main def runSharded(): Unit = + dimwit.initialize() + Jax.jax.config.update("jax_num_cpu_devices", Devices) + + val (x, y) = dataset(Random.Key(0)) + + trait X derives MeshLabel + val mesh = Mesh1(MeshAxis[X] -> Devices) + + // Inputs and targets are split the same way; the parameters stay replicated. + val shardedX: Tensor2[Batch |@| X, Feature, Float32] = x.shard(mesh, Axis[Batch] -> MeshAxis[X]) + val shardedY: Tensor1[Batch |@| X, Float32] = y.shard(mesh, Axis[Batch] -> MeshAxis[X]) + + println(s"$mesh over devices ${mesh.devices.map(_.id).mkString(", ")}") + println(s"batch ${shardedX.axes.mkString(" x ")}, ${BatchSize / Devices} rows per device") + + def loss(params: Params): Tensor0[Float32] = + val model = LinearRegression(params) + // The batch axis is called `Batch |@| X` now; `vmap` treats it like any other. + val predictions = shardedX.vmap(Axis[Batch |@| X])(model) + (predictions - shardedY).pow(2f).sum / Tensor0(BatchSize.toFloat) + + val descent = GradientDescent(Tensor0(LearningRate)).iterate(Params.init)(grad(jit(loss))) + descent.next().bias.item // warm up JIT + val start = System.nanoTime() + val params = descent.drop(Steps - 1).next() + params.bias.item + report("sharded", params, (System.nanoTime() - start) / 1e9) + + private def dataset(key: Random.Key): (Tensor2[Batch, Feature, Float32], Tensor1[Batch, Float32]) = + val (xKey, noiseKey) = key.split2() + val x = Normal.standardNormal(Shape(Axis[Batch] -> BatchSize, Axis[Feature] -> FeatureSize)).sample(xKey) + val trueWeights = Tensor1(Axis[Feature]).fromArray(Array.tabulate(FeatureSize)(i => (i % 5).toFloat - 2f)) + val noise = Normal.standardNormal(Shape(Axis[Batch] -> BatchSize)).sample(noiseKey) *! Tensor0(0.1f) + (x, x.dot(Axis[Feature])(trueWeights) + noise) + + private def report(setting: String, params: Params, seconds: Double): Unit = + println(f"$setting%-14s $Steps steps in $seconds%6.3f s (${seconds / Steps * 1000}%5.2f ms/step)") + println(f"${" "}%-14s final bias ${params.bias.item}%+.4f, weights[0] ${params.weights.slice(Axis[Feature].at(0)).item}%+.4f") diff --git a/examples/src/main/scala/dimwit/basic/ShardedSum.scala b/examples/src/main/scala/dimwit/basic/ShardedSum.scala new file mode 100644 index 0000000..17ae013 --- /dev/null +++ b/examples/src/main/scala/dimwit/basic/ShardedSum.scala @@ -0,0 +1,60 @@ +package dimwit.examples.basic + +import dimwit.* +import dimwit.jax.Jax + +/** Summing a tensor whose `Batch` axis is split across several devices. */ +object ShardedSum: + + trait Batch derives Label + trait Feature derives Label + trait X derives MeshLabel + + /** How many CPU devices to run on. `BatchSize` has to be divisible by this. */ + val Devices = 4 + + val BatchSize = 8 + + val FeatureSize = 4 + + @main def runShardedSum(): Unit = + dimwit.initialize() + Jax.jax.config.update("jax_num_cpu_devices", Devices) + + val data: Tensor2[Batch, Feature, Float32] = + Tensor2(Axis[Batch], Axis[Feature]).fromArray( + Array.tabulate(BatchSize, FeatureSize)((batch, feature) => (batch * FeatureSize + feature).toFloat) + ) + val reference: Tensor1[Feature, Float32] = data.sum(Axis[Batch]) + + val mesh = Mesh1(MeshAxis[X] -> Devices) + println(s"$mesh on devices ${mesh.devices.map(_.id).mkString(", ")}") + + val sharded = data.shard(mesh, Axis[Batch] -> MeshAxis[X]) + println(s"sharded axes: ${sharded.axes.mkString(", ")} (${BatchSize / Devices} rows per device)") + + // Reducing the sharded axis is the ordinary `sum`; XLA makes it an all-reduce. + val total: Tensor1[Feature, Float32] = sharded.sum(Axis[Batch |@| X]) + + println(s"unsharded sum: $reference") + println(s"sharded sum: $total") + println(s"identical: ${total == reference}") + + // Naming no axis works on both, so one function can serve a sharded and an unsharded run. + println(s"total of everything, either way: ${data.sum} / ${sharded.sum}") + + // Reducing an unsharded axis needs no communication; Batch stays spread over X. + val perRow: Tensor1[Batch |@| X, Float32] = sharded.sum(Axis[Feature]) + println(s"per-row sums, still sharded over X: $perRow") + + val doubled: Tensor2[Batch |@| X, Feature, Float32] = + sharded.vmap(Axis[Batch |@| X])(row => row *! Tensor0(2.0f)) + println(s"each row doubled, still sharded over X: ${doubled.axes.mkString(", ")}") + println(s"mean over the sharded axis: ${sharded.mean(Axis[Batch |@| X])}") + + // Things to try, each of which should fail: + // sharded.sum(Axis[Batch]) - the sharded axis is called Batch |@| X now (compile error) + // sharded + data - sharded and unsharded tensors do not mix (compile error) + // MeshAxis[Batch] - a data axis label is not a mesh label (compile error) + // Axis[X] - and a mesh label is not a data axis label (compile error) + // BatchSize = 7 - 7 rows do not divide over 4 devices (runtime error) From 4357e9361864054185432abdd7a07e6c5466b43a Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Tue, 8 Sep 2026 14:37:59 +0200 Subject: [PATCH 2/4] Add unshard option --- core/src/main/resources/python/jax_helper.py | 7 ++++ .../scala/dimwit/sharding/ShardingOps.scala | 14 ++++++++ .../scala/dimwit/sharding/ShardingSuite.scala | 35 +++++++++++++++++++ 3 files changed, 56 insertions(+) diff --git a/core/src/main/resources/python/jax_helper.py b/core/src/main/resources/python/jax_helper.py index f3b4136..7a73731 100644 --- a/core/src/main/resources/python/jax_helper.py +++ b/core/src/main/resources/python/jax_helper.py @@ -102,3 +102,10 @@ def named_sharding(mesh, rank, axis_index, mesh_axis_name): spec = [None] * rank spec[axis_index] = mesh_axis_name return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*spec)) + +def replicate(x): + """Gathers a sharded array so that every device of its mesh holds the whole thing.""" + sharding = x.sharding + if not hasattr(sharding, "mesh"): + return x + return jax.device_put(x, jax.sharding.NamedSharding(sharding.mesh, jax.sharding.PartitionSpec())) diff --git a/core/src/main/scala/dimwit/sharding/ShardingOps.scala b/core/src/main/scala/dimwit/sharding/ShardingOps.scala index 6bf036b..55170b2 100644 --- a/core/src/main/scala/dimwit/sharding/ShardingOps.scala +++ b/core/src/main/scala/dimwit/sharding/ShardingOps.scala @@ -38,3 +38,17 @@ object ShardingOps: ) val sharding = Jax.jax_helper.named_sharding(mesh.jaxMesh, t.shape.rank, replacer.index, meshAxisName) Tensor[replacer.NewShape, V](Jax.device_put(t.jaxValue, sharding.as[Jax.PyDynamic])) + + /** Gathers a sharded axis back across the mesh, rewriting it from `L |@| A` to `L`. + * + * Every device ends up holding the whole tensor. Only a sharded axis can be unsharded. + * + * {{{ + * val gathered: Tensor2[Batch, Feature, Float32] = sharded.unshard(Axis[Batch |@| X]) + * }}} + */ + def unshard[L, A](axis: Axis[L |@| A])(using + replacer: AxisReplacer[T, L |@| A, L], + labels: Labels[replacer.NewShape] + ): Tensor[replacer.NewShape, V] = + Tensor[replacer.NewShape, V](Jax.jax_helper.replicate(t.jaxValue).as[Jax.PyDynamic]) diff --git a/core/src/test/scala/dimwit/sharding/ShardingSuite.scala b/core/src/test/scala/dimwit/sharding/ShardingSuite.scala index 0a736ab..610448f 100644 --- a/core/src/test/scala/dimwit/sharding/ShardingSuite.scala +++ b/core/src/test/scala/dimwit/sharding/ShardingSuite.scala @@ -82,6 +82,29 @@ class ShardingSuite extends DimwitTest: val error = intercept[IllegalArgumentException](tensorOf(7).shard(mesh1, Axis[A] -> MeshAxis[X])) error.getMessage shouldBe "Cannot shard axis A of extent 7 over mesh axis X of size 4: 7 is not divisible by 4." + describe("unshard"): + + it("gathers the axis back and drops the mesh annotation"): + enoughDevices() + val gathered: Tensor2[A, B, Float32] = sharded.unshard(Axis[A |@| X]) + gathered.axes shouldBe List("A", "B") + gathered shouldEqual t + isFullyReplicated(gathered) shouldBe true + + it("round-trips with shard"): + enoughDevices() + sharded.unshard(Axis[A |@| X]).shard(mesh1, Axis[A] -> MeshAxis[X]) shouldEqual t + + it("gives back a tensor that mixes with unsharded ones again"): + enoughDevices() + val gathered = sharded.unshard(Axis[A |@| X]) + (gathered + t) shouldEqual (t + t) + + it("cannot be applied to an axis that is not sharded"): + enoughDevices() + val errors = typeCheckErrors("sharded.unshard(Axis[B])") + errors should not be empty + describe("a sharded axis is just another axis"): it("reducing the sharded axis all-reduces and equals the unsharded sum, bit for bit"): @@ -131,6 +154,18 @@ class ShardingSuite extends DimwitTest: result shouldEqual (t -! bias) deviceCount(result) shouldBe MeshExtent + it("zips with another tensor sharded the same way"): + enoughDevices() + val labels = Tensor1(Axis[A]).fromArray(Array.tabulate(BatchExtent)(_.toFloat)) + val shardedLabels = labels.shard(mesh1, Axis[A] -> MeshAxis[X]) + val result = zipvmap(Axis[A |@| X])(sharded, shardedLabels) { case (row, label) => row.sum + label } + result.axes shouldBe List("A@X") + result shouldEqual zipvmap(Axis[A])(t, labels) { case (row, label) => row.sum + label } + + it("reads back to the host, gathering from every device"): + enoughDevices() + sharded.toArray shouldEqual t.toArray + describe("sharded and unsharded tensors do not mix"): it("a sharded tensor cannot be combined with an unsharded one"): From 31be2d0df0b6697544bf32b60131784bccf8c708 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Wed, 9 Sep 2026 08:13:24 +0200 Subject: [PATCH 3/4] Code clean up --- core/src/main/scala/dimwit/package.scala | 19 ++++++++++---- .../src/main/scala/dimwit/sharding/Mesh.scala | 26 +++++++++++++++---- .../main/scala/dimwit/sharding/MeshAxis.scala | 17 ------------ .../{MeshLabel.scala => MeshLabels.scala} | 13 ++-------- .../main/scala/dimwit/sharding/Sharded.scala | 17 ------------ .../scala/dimwit/sharding/ShardingOps.scala | 11 +++----- .../scala/dimwit/sharding/ShardingSuite.scala | 20 +++++++------- 7 files changed, 52 insertions(+), 71 deletions(-) rename core/src/main/scala/dimwit/sharding/{MeshLabel.scala => MeshLabels.scala} (76%) delete mode 100644 core/src/main/scala/dimwit/sharding/Sharded.scala diff --git a/core/src/main/scala/dimwit/package.scala b/core/src/main/scala/dimwit/package.scala index c937b04..00c4c61 100644 --- a/core/src/main/scala/dimwit/package.scala +++ b/core/src/main/scala/dimwit/package.scala @@ -39,6 +39,16 @@ package object dimwit: given [A, B](using labelA: Label[A], labelB: Label[B]): Label[A |+| B] with val name: String = s"${labelA.name}+${labelB.name}" + /** Marks the data axis `A` as sharded over the mesh axis `M`, as in `Tensor2[Batch |@| X, Feature, Float32]`. + * + * `Batch |@| X` is just another axis label, so every tensor operation applies unchanged. + */ + @targetName("Sharded") + infix trait |@|[A, M] + object `|@|`: + given [A, M](using label: Label[A], meshLabel: MeshLabel[M]): Label[A |@| M] with + val name: String = s"${label.name}@${meshLabel.name}" + // Export tensor and related types export dimwit.tensor.{Tensor, Tensor0, Tensor1, Tensor2, Tensor3, Tensor4, TypedIndex} export dimwit.tensor.{Shape, Shape0, Shape1, Shape2, Shape3} @@ -62,11 +72,6 @@ package object dimwit: // Export the Prime axis marker and the type classes that manipulate it export dimwit.prime.{Prime, PrimeRemover, PrimeRest, PrimeConcat} - // Export the sharding types: mesh labels, meshes, and the sharded axis marker - export dimwit.sharding.{Mesh, Mesh1, Mesh2, Mesh3, MeshAxis, MeshAxisExtent, MeshAxisIndex, MeshLabel, MeshLabels} - export dimwit.sharding.`|@|` - export dimwit.sharding.ShardingOps.* - // Export operations export dimwit.tensor.TensorOps.* export dimwit.linalg.LinearAlgebra.{VectorNormType, MatrixNormType, QRMode} @@ -93,6 +98,10 @@ package object dimwit: export dimwit.stats.{Prob, LogProb} export dimwit.stats.{Distribution, IndependentDistribution, MultivariateDistribution, UnivariateDistribution} + // Export sharding stuff + export dimwit.sharding.{Mesh, Mesh1, Mesh2, Mesh3, MeshAxis, MeshAxisExtent, MeshAxisIndex, MeshLabel, MeshLabels} + export dimwit.sharding.ShardingOps.* + /** Memory management helpfer making sure * all python objects allocated ar freed * after the function is executed. diff --git a/core/src/main/scala/dimwit/sharding/Mesh.scala b/core/src/main/scala/dimwit/sharding/Mesh.scala index fc254b1..29fc3a4 100644 --- a/core/src/main/scala/dimwit/sharding/Mesh.scala +++ b/core/src/main/scala/dimwit/sharding/Mesh.scala @@ -4,6 +4,8 @@ import dimwit.hardware.Device import dimwit.jax.Jax import me.shadaj.scalapy.py.SeqConverters +import scala.annotation.implicitNotFound + /** A grid of devices, binding each mesh axis label of `M` to a number of devices. */ final class Mesh[M <: Tuple: MeshLabels] private[sharding] ( val axisSizes: List[Int], @@ -26,11 +28,10 @@ final class Mesh[M <: Tuple: MeshLabels] private[sharding] ( object Mesh: - private[sharding] type ExtractLabels[Extents <: Tuple] <: Tuple = Extents match - case EmptyTuple => EmptyTuple - case MeshAxisExtent[a] *: tail => a *: ExtractLabels[tail] + private[sharding] type ExtractLabel[Extent] = Extent match + case MeshAxisExtent[a] => a - def apply[A: MeshLabel](extent: MeshAxisExtent[A]): Mesh[Tuple1[A]] = fromTuple(Tuple1(extent)) + private[sharding] type ExtractLabels[Extents <: Tuple] = Tuple.Map[Extents, ExtractLabel] def apply[Extents <: Tuple](extents: Extents)(using MeshLabels[ExtractLabels[Extents]]): Mesh[ExtractLabels[Extents]] = fromTuple(extents) @@ -51,7 +52,7 @@ type Mesh2[A, B] = Mesh[(A, B)] type Mesh3[A, B, C] = Mesh[(A, B, C)] object Mesh1: - def apply[A: MeshLabel](extent: MeshAxisExtent[A]): Mesh1[A] = Mesh(extent) + def apply[A: MeshLabel](extent: MeshAxisExtent[A]): Mesh1[A] = Mesh.fromTuple(Tuple1(extent)) object Mesh2: def apply[A: MeshLabel, B: MeshLabel](extent1: MeshAxisExtent[A], extent2: MeshAxisExtent[B]): Mesh2[A, B] = @@ -63,3 +64,18 @@ object Mesh3: extent2: MeshAxisExtent[B], extent3: MeshAxisExtent[C] ): Mesh3[A, B, C] = Mesh.fromTuple((extent1, extent2, extent3)) + +/** Finds the position of a mesh axis in a mesh. */ +@implicitNotFound("MeshAxis[${A}] not found in Mesh[${M}]") +trait MeshAxisIndex[M <: Tuple, A]: + def index: Int + +object MeshAxisIndex: + + def apply[M <: Tuple, A](using idx: MeshAxisIndex[M, A]): Int = idx.index + + given found[A, Tail <: Tuple]: MeshAxisIndex[A *: Tail, A] with + val index = 0 + + given search[H, T <: Tuple, A](using next: MeshAxisIndex[T, A]): MeshAxisIndex[H *: T, A] with + val index = 1 + next.index diff --git a/core/src/main/scala/dimwit/sharding/MeshAxis.scala b/core/src/main/scala/dimwit/sharding/MeshAxis.scala index fef34bb..2394d3e 100644 --- a/core/src/main/scala/dimwit/sharding/MeshAxis.scala +++ b/core/src/main/scala/dimwit/sharding/MeshAxis.scala @@ -1,7 +1,5 @@ package dimwit.sharding -import scala.annotation.implicitNotFound - /** Represents an axis of a device [[Mesh]]. The mesh-side counterpart of [[dimwit.tensor.Axis]]. */ final class MeshAxis[A: MeshLabel]: @@ -14,18 +12,3 @@ final class MeshAxis[A: MeshLabel]: /** A mesh axis together with the number of devices along it. */ case class MeshAxisExtent[A: MeshLabel](axis: MeshAxis[A], size: Int) - -/** Finds the position of a mesh axis in a mesh. */ -@implicitNotFound("MeshAxis[${A}] not found in Mesh[${M}]") -trait MeshAxisIndex[M <: Tuple, A]: - def index: Int - -object MeshAxisIndex: - - def apply[M <: Tuple, A](using idx: MeshAxisIndex[M, A]): Int = idx.index - - given found[A, Tail <: Tuple]: MeshAxisIndex[A *: Tail, A] with - val index = 0 - - given search[H, T <: Tuple, A](using next: MeshAxisIndex[T, A]): MeshAxisIndex[H *: T, A] with - val index = 1 + next.index diff --git a/core/src/main/scala/dimwit/sharding/MeshLabel.scala b/core/src/main/scala/dimwit/sharding/MeshLabels.scala similarity index 76% rename from core/src/main/scala/dimwit/sharding/MeshLabel.scala rename to core/src/main/scala/dimwit/sharding/MeshLabels.scala index 6df7736..77639da 100644 --- a/core/src/main/scala/dimwit/sharding/MeshLabel.scala +++ b/core/src/main/scala/dimwit/sharding/MeshLabels.scala @@ -2,20 +2,11 @@ package dimwit.sharding import scala.quoted.* -/** A label for an axis of a device [[Mesh]]. - * - * Mesh labels are a separate kind from the data axis labels carried by - * `dimwit.tensor.Label`: a type that `derives Label` has no `MeshLabel` instance and - * vice versa, so a data axis label cannot be used where a mesh label is required. - * - * {{{ - * trait X derives MeshLabel - * }}} +/** A label for an axis of a device [[Mesh]], declared with `trait X derives MeshLabel`. + * The mesh-side counterpart of [[dimwit.tensor.Label]]. */ @scala.annotation.implicitNotFound(""" A mesh axis label ${T} was given or inferred, which does not have a MeshLabel instance. -Mesh axis labels are a different kind than data axis labels: a type declared with -'derives Label' cannot be used as a mesh axis label. Ensure that all mesh axis types ${T} are defined with 'derives MeshLabel' (e.g. 'trait X derives MeshLabel') """) trait MeshLabel[T]: diff --git a/core/src/main/scala/dimwit/sharding/Sharded.scala b/core/src/main/scala/dimwit/sharding/Sharded.scala deleted file mode 100644 index 5828b91..0000000 --- a/core/src/main/scala/dimwit/sharding/Sharded.scala +++ /dev/null @@ -1,17 +0,0 @@ -package dimwit.sharding - -import dimwit.tensor.Label - -/** Marks the data axis `A` as sharded over the mesh axis `M`, as in `Tensor2[Batch |@| X, Feature, Float32]`. - * - * `Batch |@| X` is just another axis label, so every operation applies to a sharded tensor - * unchanged: reducing it all-reduces, reducing or mapping any other axis stays local, and it - * cannot be combined with an unsharded `Batch` because the names differ. - * - * Spelled `|@|` alongside [[dimwit.|*|]] and [[dimwit.|+|]]; a bare `@` is Scala's annotation syntax. - */ -infix trait |@|[A, M] - -object `|@|`: - given [A, M](using label: Label[A], meshLabel: MeshLabel[M]): Label[A |@| M] with - val name: String = s"${label.name}@${meshLabel.name}" diff --git a/core/src/main/scala/dimwit/sharding/ShardingOps.scala b/core/src/main/scala/dimwit/sharding/ShardingOps.scala index 55170b2..a00d391 100644 --- a/core/src/main/scala/dimwit/sharding/ShardingOps.scala +++ b/core/src/main/scala/dimwit/sharding/ShardingOps.scala @@ -1,6 +1,7 @@ package dimwit.sharding import dimwit.jax.Jax +import dimwit.|@| import dimwit.tensor.Axis import dimwit.tensor.Label import dimwit.tensor.Labels @@ -11,12 +12,10 @@ object ShardingOps: extension [T <: Tuple: Labels, V](t: Tensor[T, V]) - /** Splits `t` across `mesh` along one axis, rewriting that axis from `L` to `L |@| A`. - * - * The shards are placed one per device with `jax.device_put` under a `NamedSharding`. + /** Splits `t` across `mesh` along one axis, placing one shard per device and rewriting that axis from `L` to `L |@| A`. * * {{{ - * val mesh = Mesh(MeshAxis[X] -> 4) + * val mesh = Mesh1(MeshAxis[X] -> 4) * val sharded: Tensor2[Batch |@| X, Feature, Float32] = t.shard(mesh, Axis[Batch] -> MeshAxis[X]) * }}} * @@ -39,9 +38,7 @@ object ShardingOps: val sharding = Jax.jax_helper.named_sharding(mesh.jaxMesh, t.shape.rank, replacer.index, meshAxisName) Tensor[replacer.NewShape, V](Jax.device_put(t.jaxValue, sharding.as[Jax.PyDynamic])) - /** Gathers a sharded axis back across the mesh, rewriting it from `L |@| A` to `L`. - * - * Every device ends up holding the whole tensor. Only a sharded axis can be unsharded. + /** Gathers a sharded axis back across the mesh, rewriting it from `L |@| A` to `L` so that every device holds the whole tensor. * * {{{ * val gathered: Tensor2[Batch, Feature, Float32] = sharded.unshard(Axis[Batch |@| X]) diff --git a/core/src/test/scala/dimwit/sharding/ShardingSuite.scala b/core/src/test/scala/dimwit/sharding/ShardingSuite.scala index 610448f..480ebfd 100644 --- a/core/src/test/scala/dimwit/sharding/ShardingSuite.scala +++ b/core/src/test/scala/dimwit/sharding/ShardingSuite.scala @@ -50,11 +50,11 @@ class ShardingSuite extends DimwitTest: it("is built from a tuple of extents, with Mesh1/Mesh2 as special cases"): enoughDevices() - val single: Mesh[Tuple1[X]] = Mesh(MeshAxis[X] -> MeshExtent) + val single = Mesh1(MeshAxis[X] -> MeshExtent) single.axisNames shouldBe List("X") single.axisSizes shouldBe List(MeshExtent) - val grid: Mesh[(X, Y)] = Mesh((MeshAxis[X] -> 2, MeshAxis[Y] -> 2)) + val grid = Mesh((MeshAxis[X] -> 2, MeshAxis[Y] -> 2)) grid.axisNames shouldBe List("X", "Y") grid.size shouldBe 4 grid.sizeOf(MeshAxis[Y]) shouldBe 2 @@ -70,7 +70,7 @@ class ShardingSuite extends DimwitTest: it("rewrites the axis type and places the shards on the mesh devices"): enoughDevices() - val result: Tensor2[A |@| X, B, Float32] = sharded + val result = sharded result.axes shouldBe List("A@X", "B") result.shape.dimensions shouldBe List(BatchExtent, FeatureExtent) deviceCount(result) shouldBe MeshExtent @@ -86,7 +86,7 @@ class ShardingSuite extends DimwitTest: it("gathers the axis back and drops the mesh annotation"): enoughDevices() - val gathered: Tensor2[A, B, Float32] = sharded.unshard(Axis[A |@| X]) + val gathered = sharded.unshard(Axis[A |@| X]) gathered.axes shouldBe List("A", "B") gathered shouldEqual t isFullyReplicated(gathered) shouldBe true @@ -109,7 +109,7 @@ class ShardingSuite extends DimwitTest: it("reducing the sharded axis all-reduces and equals the unsharded sum, bit for bit"): enoughDevices() - val result: Tensor1[B, Float32] = sharded.sum(Axis[A |@| X]) + val result = sharded.sum(Axis[A |@| X]) result shouldEqual t.sum(Axis[A]) result.axes shouldBe List("B") isFullyReplicated(result) shouldBe true @@ -121,7 +121,7 @@ class ShardingSuite extends DimwitTest: it("reducing another axis is local and keeps the mesh annotation"): enoughDevices() - val result: Tensor1[A |@| X, Float32] = sharded.sum(Axis[B]) + val result = sharded.sum(Axis[B]) result shouldEqual t.sum(Axis[B]) result.axes shouldBe List("A@X") isFullyReplicated(result) shouldBe false @@ -135,7 +135,7 @@ class ShardingSuite extends DimwitTest: it("vmap maps over the sharded axis, which stays sharded"): enoughDevices() - val result: Tensor2[A |@| X, B, Float32] = sharded.vmap(Axis[A |@| X])(row => row *! Tensor0(2.0f)) + val result = sharded.vmap(Axis[A |@| X])(row => row *! Tensor0(2.0f)) result shouldEqual t.vmap(Axis[A])(row => row *! Tensor0(2.0f)) result.axes shouldBe List("A@X", "B") deviceCount(result) shouldBe MeshExtent @@ -143,15 +143,17 @@ class ShardingSuite extends DimwitTest: it("contracts a replicated tensor against the unsharded axis, staying sharded"): enoughDevices() val weights = Tensor1(Axis[B]).fromArray(Array.fill(FeatureExtent)(2.0f)) - val result: Tensor1[A |@| X, Float32] = sharded.dot(Axis[B])(weights) + val result = sharded.dot(Axis[B])(weights) result shouldEqual t.dot(Axis[B])(weights) + result.axes shouldBe List("A@X") deviceCount(result) shouldBe MeshExtent it("broadcasts a replicated tensor over the axes it does share"): enoughDevices() val bias = Tensor1(Axis[B]).fromArray(Array.fill(FeatureExtent)(1.0f)) - val result: Tensor2[A |@| X, B, Float32] = sharded -! bias + val result = sharded -! bias result shouldEqual (t -! bias) + result.axes shouldBe List("A@X", "B") deviceCount(result) shouldBe MeshExtent it("zips with another tensor sharded the same way"): From 425999c76cb564c2633e8193a9b38a7aea98d04c Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Wed, 9 Sep 2026 17:44:23 +0200 Subject: [PATCH 4/4] Add Mesh constructor to not require Mesh1 --- core/src/main/scala/dimwit/sharding/Mesh.scala | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/core/src/main/scala/dimwit/sharding/Mesh.scala b/core/src/main/scala/dimwit/sharding/Mesh.scala index 29fc3a4..dab3ec5 100644 --- a/core/src/main/scala/dimwit/sharding/Mesh.scala +++ b/core/src/main/scala/dimwit/sharding/Mesh.scala @@ -33,8 +33,10 @@ object Mesh: private[sharding] type ExtractLabels[Extents <: Tuple] = Tuple.Map[Extents, ExtractLabel] - def apply[Extents <: Tuple](extents: Extents)(using MeshLabels[ExtractLabels[Extents]]): Mesh[ExtractLabels[Extents]] = - fromTuple(extents) + def apply[A: MeshLabel](extent: MeshAxisExtent[A]): Mesh1[A] = Mesh1(extent) + def apply[A: MeshLabel, B: MeshLabel](extent1: MeshAxisExtent[A], extent2: MeshAxisExtent[B]): Mesh2[A, B] = Mesh2(extent1, extent2) + def apply[A: MeshLabel, B: MeshLabel, C: MeshLabel](extent1: MeshAxisExtent[A], extent2: MeshAxisExtent[B], extent3: MeshAxisExtent[C]): Mesh3[A, B, C] = Mesh3(extent1, extent2, extent3) + def apply[Extents <: Tuple](extents: Extents)(using MeshLabels[ExtractLabels[Extents]]): Mesh[ExtractLabels[Extents]] = fromTuple(extents) def fromTuple[Extents <: Tuple](extents: Extents)(using labels: MeshLabels[ExtractLabels[Extents]]): Mesh[ExtractLabels[Extents]] = val sizes = extents.toList.collect: