Remove covariance from tensor. Fix issues related to this.

core/src/main/scala/dimwit/stats/IndependentDistributions.scala

@@ -8,7 +8,7 @@ import me.shadaj.scalapy.py
 import me.shadaj.scalapy.py.SeqConverters
 import dimwit.random.Random
 
-class Normal[LocT <: T, ScaleT <: T, T <: Tuple: Labels](
+class Normal[T <: Tuple: Labels, LocT <: T, ScaleT <: T](
     val loc: Tensor[LocT, Float],
     val scale: Tensor[ScaleT, Float]
 ) extends IndependentDistribution[T, Float]:

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

@@ -28,7 +28,7 @@ object Device:
       require(devices.nonEmpty, s"No JAX devices found for platform: ${device.platform}")
       devices.head
 
-class Tensor[+T <: Tuple: Labels, V] private[tensor] (
+class Tensor[T <: Tuple: Labels, V] private[tensor] (
     val jaxValue: Jax.PyDynamic
 ):
 

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

@@ -87,26 +87,29 @@ object TensorOps:
     // IsNumber operations (IsFloat or IsInt)
     // ---------------------------------------------------------
 
-    def add[T <: Tuple: Labels, V: IsNumber](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.add(t1.jaxValue, t2.jaxValue))
+    def add[T <: Tuple: Labels, T1 <: T, T2 <: T, V: IsNumber](t1: Tensor[T1, V], t2: Tensor[T2, V]): Tensor[T, V] = Tensor(Jax.jnp.add(t1.jaxValue, t2.jaxValue))
     def addScalar[T <: Tuple: Labels, V: IsNumber](t1: Tensor[T, V], t2: Tensor0[V]): Tensor[T, V] = Tensor(Jax.jnp.add(t1.jaxValue, t2.jaxValue))
 
     def negate[T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.negative(t.jaxValue))
-    def subtract[T <: Tuple: Labels, V: IsNumber](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.subtract(t1.jaxValue, t2.jaxValue))
+    def subtract[T <: Tuple: Labels, T1 <: T, T2 <: T, V: IsNumber](t1: Tensor[T1, V], t2: Tensor[T2, V]): Tensor[T, V] = Tensor(Jax.jnp.subtract(t1.jaxValue, t2.jaxValue))
     def subtractScalar[T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V], t2: Tensor0[V]): Tensor[T, V] = Tensor(Jax.jnp.subtract(t.jaxValue, t2.jaxValue))
 
-    def multiply[T <: Tuple: Labels, V: IsNumber](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.multiply(t1.jaxValue, t2.jaxValue))
+    def multiply[T <: Tuple: Labels, T1 <: T, T2 <: T, V: IsNumber](t1: Tensor[T1, V], t2: Tensor[T2, V]): Tensor[T, V] = Tensor(Jax.jnp.multiply(t1.jaxValue, t2.jaxValue))
     def multiplyScalar[T <: Tuple: Labels, V: IsNumber](t1: Tensor[T, V], t2: Tensor0[V]): Tensor[T, V] = Tensor(Jax.jnp.multiply(t1.jaxValue, t2.jaxValue))
 
+    extension [T <: Tuple: Labels, T1 <: T, T2 <: T, V: IsNumber](t: Tensor[T1, V])
+
+      def +(other: Tensor[T2, V]): Tensor[T, V] = add(t, other)
+      def -(other: Tensor[T2, V]): Tensor[T, V] = subtract(t, other)
+      def *(other: Tensor[T2, V]): Tensor[T, V] = multiply(t, other)
+
     extension [T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V])
 
-      def +(other: Tensor[T, V]): Tensor[T, V] = add(t, other)
       def +![O <: Tuple](other: Tensor[O, V])(using bc: Broadcast[T, O, V]): Tensor[bc.Out, V] = bc.applyTo(t, other)(add)
 
       def unary_- : Tensor[T, V] = negate(t)
-      def -(other: Tensor[T, V]): Tensor[T, V] = subtract(t, other)
       def -![O <: Tuple](other: Tensor[O, V])(using bc: Broadcast[T, O, V]): Tensor[bc.Out, V] = bc.applyTo(t, other)(subtract)
 
-      def *(other: Tensor[T, V]): Tensor[T, V] = multiply(t, other)
       def *![O <: Tuple](other: Tensor[O, V])(using bc: Broadcast[T, O, V]): Tensor[bc.Out, V] = bc.applyTo(t, other)(multiply)
       def scale(other: Tensor0[V]): Tensor[T, V] = multiplyScalar(t, other)
 

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

@@ -65,8 +65,8 @@ object TupleHelpers:
 
     // 3. Found Case: H is a subtype of K
     // We explicitly return 'Tail' as the output
-    given singleFound[H, Tail <: Tuple, K](using H <:< K): Aux[H *: Tail, K *: EmptyTuple, Tail] =
-      new RemoverAll[H *: Tail, K *: EmptyTuple]:
+    given singleFound[K, Tail <: Tuple]: Aux[K *: Tail, K *: EmptyTuple, Tail] =
+      new RemoverAll[K *: Tail, K *: EmptyTuple]:
         type Out = Tail
 
   trait LowPriorityRemoverAll:

core/src/test/scala/dimwit/tensor/TensorCovarianceSuite.scala

@@ -8,20 +8,6 @@ import scala.collection.View.Empty
 
 class TensorCovarianceSuite extends AnyFunSpec with Matchers:
 
-  it("Shape type hierarchy example: Concrete function with supertype parameter"):
-    trait Parent derives Label
-    trait Child1 extends Parent derives Label
-    trait Child2 extends Parent derives Label
-    trait NoChild derives Label
-    def concreteFunction(t: Tensor1[Parent, Float]): Tensor1[Parent, Float] = t + t
-    val child1: Tensor1[Child1, Float] = Tensor(Shape1(Axis[Child1] -> 4)).fill(1f)
-    val child2: Tensor1[Child2, Float] = Tensor(Shape1(Axis[Child2] -> 4)).fill(1f)
-    val noChild: Tensor1[NoChild, Float] = Tensor(Shape1(Axis[NoChild] -> 4)).fill(1f)
-
-    "concreteFunction(child1)" should compile
-    "concreteFunction(child2)" should compile
-    "concreteFunction(noChild)" shouldNot compile
-
   it("Shape type hierarchy example: Generic function with upper-bounded type parameter"):
     trait Parent derives Label
     trait Child1 extends Parent derives Label

examples/src/main/scala/basic/Autoencoder.scala

@@ -160,17 +160,17 @@ object AutoencoderExample:
      * Training loop
      * */
 
-    def loss(trainData: Tensor3[Sample, Height, Width, Float])(params: Autoencoder.Params): Tensor0[Float] =
+    def loss[S <: Sample: Label](trainData: Tensor3[S, Height, Width, Float])(params: Autoencoder.Params): Tensor0[Float] =
       val ae = Autoencoder(params)
       trainData
-        .vmap(Axis[Sample])(sample => ae.loss(sample.ravel))
+        .vmap(Axis[S])(sample => ae.loss(sample.ravel))
         .mean
 
     val batches = trainX.chunk(Axis[TrainSample], numSamples / batchSize)
 
     val optimizer = GradientDescent(learningRate = Tensor0(learningRate))
 
-    def gradientStep(batch: Tensor3[Sample, Height, Width, Float], params: Autoencoder.Params): Autoencoder.Params =
+    def gradientStep(batch: Tensor3[TrainSample, Height, Width, Float], params: Autoencoder.Params): Autoencoder.Params =
       val grads = Autodiff.grad(loss(batch))(params)
       val (newParams, _) = optimizer.update(grads, params, ())
       newParams

examples/src/main/scala/basic/MLClassifierMNist.scala

@@ -87,11 +87,11 @@ object MLPClassifierMNist:
       Axis[Output] -> 10
     )(initKey)
 
-    def accuracy[Sample: Label](
-        predictions: Tensor1[Sample, Int],
-        targets: Tensor1[Sample, Int]
+    def accuracy[S: Label](
+        predictions: Tensor1[S, Int],
+        targets: Tensor1[S, Int]
     ): Tensor0[Float] =
-      val matches = zipvmap(Axis[Sample])(predictions, targets)(_ === _)
+      val matches = zipvmap(Axis[S])(predictions, targets)(_ === _)
       matches.asFloat.mean
 
     val optimizer = Lion(learningRate = Tensor0(learningRate), weightDecay = Tensor0(0f))
@@ -128,21 +128,21 @@ object MLPClassifierMNist:
       timed("Training"):
         dimwit.gc()
         trainMiniBatchGradientDescent(currentParams, state)
-    def evaluate(
+    def evaluate[S <: Sample: Label](
         params: MLP.Params,
-        dataX: Tensor[(Sample, Height, Width), Float],
-        dataY: Tensor1[Sample, Int]
+        dataX: Tensor3[S, Height, Width, Float],
+        dataY: Tensor1[S, Int]
     ): Tensor0[Float] =
       val model = MLP(params)
-      val predictions = dataX.vmap(Axis[Sample])(model)
+      val predictions = dataX.vmap(Axis[S])(model)
       accuracy(predictions, dataY)
-    val jitEvaluate = jit(evaluate)
+    val jitEvaluate = evaluate
     val (finalParams, finalState) = trainTrajectory.zipWithIndex
       .tapEach:
         case ((params, state), epoch) =>
           timed("Evaluation"):
-            val testAccuracy = jitEvaluate(params, testX, testY)
-            val trainAccuracy = jitEvaluate(params, trainX, trainY)
+            val testAccuracy = evaluate(params, testX, testY)
+            val trainAccuracy = evaluate(params, trainX, trainY)
             println(
               List(
                 s"Epoch $epoch",

examples/src/main/scala/complex/GPT2.scala

@@ -126,7 +126,8 @@ case class GPT2(params: GPT2Params) extends (Tensor2[Batch, Context, Int] => Ten
       val attnScores = (queries.dot(Axis[HeadQuery ~ HeadKey])(keys) /! dk)
       val attnWeights = causalMasking(attnScores)
         .vmap(Axis[Context])(attnScore => softmax(attnScore).relabelTo(Axis[AttnWeights]))
-      attnWeights.dot(Axis[AttnWeights ~ Context])(values)
+      val res = attnWeights.dot(Axis[AttnWeights ~ Context])(values)
+      res
 
   private case class LayerNorm(params: LayerNormalizationParams) extends (Tensor1[Embedding, Float] => Tensor1[Embedding, Float]):
 

examples/src/main/scala/complex/VariationalAutoencoder.scala

@@ -77,7 +77,7 @@ case class VariationalAutoencoder(params: VariationalAutoencoder.Params):
   val encoder = Encoder(params.encoderParams)
   val decoder = Decoder(params.decoderParams)
 
-  def apply(pixels: Tensor1[Pixel, Float], key: Random.Key): (Tensor1[ReconstructedPixel, Float], Tensor1[Latent, Float], Tensor1[Latent, Float]) =
+  def apply(pixels: Tensor1[Pixel, Float], key: Random.Key): (Tensor1[ReconstructedPixel, Float], Tensor1[MeanLatent, Float], Tensor1[LogVarLatent, Float]) =
     val (mean, logVar) = encoder(pixels)
     val latent = reparametrize(mean, logVar, key)
     val reconstructedPixels = decoder(latent)
@@ -170,18 +170,18 @@ object VariationalAutoencoderExample:
     /*
      * Training
      */
-    def batchLoss(key: Random.Key, trainData: Tensor3[Sample, Height, Width, Float])(params: Params): Tensor0[Float] =
+    def batchLoss[S <: Sample: Label](key: Random.Key, trainData: Tensor3[S, Height, Width, Float])(params: Params): Tensor0[Float] =
       val vae = VariationalAutoencoder(params)
-      val batchSize = trainData.shape.dim(Axis[Sample])._2
+      val batchSize = trainData.shape.dim(Axis[S])._2
       val keys = key.split(batchSize)
       val losses = (0 until batchSize).map: idx =>
-        val sample = trainData.slice(Axis[Sample] -> idx)
+        val sample = trainData.slice(Axis[S] -> idx)
         vae.loss(sample.ravel, keys(idx))
       losses.reduce(_ + _) / batchSize.toFloat
 
     val batches = trainImages.chunk(Axis[TrainSample], numSamples / batchSize)
     val optimizer = GradientDescent(learningRate = Tensor0(learningRate))
-    def trainBatch(trainKey: Random.Key, batch: Tensor3[Sample, Height, Width, Float], params: Params): Params =
+    def trainBatch(trainKey: Random.Key, batch: Tensor3[TrainSample, Height, Width, Float], params: Params): Params =
       val grads = Autodiff.grad(batchLoss(trainKey, batch))(params)
       val (newParams, _) = optimizer.update(grads, params, ())
       newParams