From 5cac0d85a1b8bc9a433417d83eb58fcbca9d6034 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Mon, 28 Sep 2026 18:41:20 +0200 Subject: [PATCH] Update to new DimWit --- .../src/main/scala/deepwit/examples/gpt/FineWebDataset.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/src/main/scala/deepwit/examples/gpt/FineWebDataset.scala b/examples/src/main/scala/deepwit/examples/gpt/FineWebDataset.scala index e0f0924..f68cc4b 100644 --- a/examples/src/main/scala/deepwit/examples/gpt/FineWebDataset.scala +++ b/examples/src/main/scala/deepwit/examples/gpt/FineWebDataset.scala @@ -4,7 +4,7 @@ import java.io.File import dimwit.* import dimwit.python.PyBridge.liftPyTensor import me.shadaj.scalapy.py -import dimwit.stats.Uniform +import dimwit.stats.DiscreteUniform import dimwit.hardware.DeviceBackend.GPU object FineWebDataset: @@ -28,7 +28,7 @@ object FineWebDataset: val maxIdx = data.shape(Axis[Sample]) - batchSize - 1 val randomIndices = IndependentDistribution.fromUnivariate( Shape1(Axis[Sample] -> batchSize), - Uniform(Tensor0(0), Tensor0(maxIdx)) + DiscreteUniform(0, maxIdx) ).sample(key) val shiftedIndices = randomIndices +! 1 val inputs = randomIndices.vmap(Axis[Sample])(startIndex =>