Simplify Module forward signature to be just function, to an extension function.

michalharakal · michalharakal · commit 4186c4f27e01 · 2025-09-25T21:42:15.000+02:00
Relate-To: #117, #104
diff --git a/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/Flatten.kt b/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/Flatten.kt
@@ -16,7 +16,9 @@ public class Flatten<T : DType, V>(
     override val modules: List<Module<T, V>>
         get() = emptyList()
 
-    override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> {
-        return input.flatten(startDim, endDim)
+    override fun forward(input: Tensor<T, V>): Tensor<T, V> {
+        return with(input) {
+            flatten(startDim, endDim)
+        }
     }
 }
diff --git a/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/Input.kt b/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/Input.kt
@@ -12,7 +12,7 @@ public class Input<T : DType, V>(private val inputShape: Shape, override val nam
         get() = emptyList()
 
 
-    override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> {
+    override fun forward(input: Tensor<T, V>): Tensor<T, V> {
         return input
     }
 }
diff --git a/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/Linear.kt b/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/Linear.kt
@@ -37,13 +37,15 @@ public class Linear<T : DType, V>(
     override val modules: List<Module<T, V>>
         get() = emptyList()
 
-    override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> {
+    override fun forward(input: Tensor<T, V>): Tensor<T, V> {
         val weight = params.weights().value
         val bias = params.bias().value
 
         // Use TensorOps context operations
-        val weightTransposed = weight.t()
-        val matmulResult = matmul(input, weightTransposed)
-        return matmulResult + bias
+        with(input) {
+            val weightTransposed = weight.t()
+            val matmulResult = matmul(input, weightTransposed)
+            return matmulResult + bias
+        }
     }
 }
diff --git a/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/Module.kt b/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/Module.kt
@@ -11,9 +11,9 @@ public abstract class Module<T : DType, V> {
 
     public abstract val modules: List<Module<T, V>>
 
-    public abstract fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V>
+    public abstract fun forward(input: Tensor<T, V>): Tensor<T, V>
 
-    public operator fun Tensor<T, V>.invoke(input: Tensor<T, V>): Tensor<T, V> {
+    public operator fun invoke(input: Tensor<T, V>): Tensor<T, V> {
         return forward(input)
     }
 }
diff --git a/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/activations/ActivationsWrapperModule.kt b/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/activations/ActivationsWrapperModule.kt
@@ -14,7 +14,7 @@ public class ActivationsWrapperModule<T : DType, V>(
     override val modules: List<Module<T, V>>
         get() = emptyList()
 
-    override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> {
+    override fun forward(input: Tensor<T, V>): Tensor<T, V> {
         return activationHandler(input)
     }
 }
diff --git a/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/activations/Softmax.kt b/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/activations/Softmax.kt
@@ -9,8 +9,10 @@ public class Softmax<T : DType, V>(private val dimension: Int, override val name
     override val modules: List<Module<T, V>>
         get() = emptyList()
 
-    override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> {
-        return input.softmax(dimension)
+    override fun forward(input: Tensor<T, V>): Tensor<T, V> {
+        return with(input) {
+            softmax(dimension)
+        }
     }
 }
 
diff --git a/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/activations/relu.kt b/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/activations/relu.kt
@@ -9,6 +9,6 @@ public class ReLU<T : DType, V>(override val name: String = "ReLU") : Module<T,
     override val modules: List<Module<T, V>>
         get() = emptyList()
 
-    override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> = input.relu()
+    override fun forward(input: Tensor<T, V>): Tensor<T, V> = with(input) { relu() }
 }
 
diff --git a/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/topology/MLP.kt b/skainet-nn/skainet-nn-api/src/commonMain/kotlin/sk/ainet/nn/topology/MLP.kt
@@ -12,10 +12,10 @@ public class MLP<T : DType, V>(vararg modules: Module<T, V>, override val name:
     override val modules: List<Module<T, V>>
         get() = modulesList
 
-    override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> {
+    override fun forward(input: Tensor<T, V>): Tensor<T, V> {
         var tmp = input
         modulesList.forEach { module ->
-            tmp = with(module) { this@forward.forward(tmp) }
+            tmp = module.forward(tmp)
         }
         return tmp
     }
diff --git a/skainet-nn/skainet-nn-api/src/commonTest/kotlin/sk/ainet/nn/dsl/NetworkBuilderGenericTest.kt b/skainet-nn/skainet-nn-api/src/commonTest/kotlin/sk/ainet/nn/dsl/NetworkBuilderGenericTest.kt
@@ -9,6 +9,8 @@ import sk.ainet.core.tensor.Int8
 import sk.ainet.core.tensor.Int32
 import sk.ainet.core.tensor.backend.CpuBackend
 import sk.ainet.core.tensor.backend.CpuBackendInt8
+import sk.ainet.nn.Linear
+import sk.ainet.nn.Module
 import kotlin.test.*
 
 /**
@@ -20,7 +22,7 @@ class NetworkBuilderGenericTest {
     @Test
     fun testGenericNetworkFP32() {
         // Test FP32/Float combination - basic functionality
-        val network = network<FP32, Float> {
+        val network: Module<FP32, Float> = network<FP32, Float> {
             input(2)
             dense(3) {
                 weights { shape -> CpuTensorFP32.ones(shape) }
@@ -181,12 +183,12 @@ class NetworkBuilderGenericTest {
         // Test NetworkBuilder class directly with generic types
         val builder = NetworkBuilder<FP32, Float>()
         
-        val linear1 = sk.ainet.nn.Linear(2, 4, "layer1", 
+        val linear1 = Linear(2, 4, "layer1",
             CpuTensorFP32.ones(Shape(4, 2)), 
             CpuTensorFP32.zeros(Shape(4))
         )
         
-        val linear2 = sk.ainet.nn.Linear(4, 1, "layer2",
+        val linear2 = Linear(4, 1, "layer2",
             CpuTensorFP32.ones(Shape(1, 4)),
             CpuTensorFP32.zeros(Shape(1))
         )
diff --git a/skainet-nn/skainet-nn-api/src/commonTest/kotlin/sk/ainet/nn/dsl/NetworkBuilderTest.kt b/skainet-nn/skainet-nn-api/src/commonTest/kotlin/sk/ainet/nn/dsl/NetworkBuilderTest.kt
@@ -38,7 +38,7 @@ class NetworkBuilderTest {
         
         // Test forward pass
         val input = CpuTensorFP32.fromArray(Shape(1, 2), floatArrayOf(1.0f, 2.0f))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
     }
 
@@ -63,7 +63,7 @@ class NetworkBuilderTest {
         
         // Test forward pass
         val input = CpuTensorInt8.fromArray(Shape(1, 2), byteArrayOf(1, 2))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
     }
 
@@ -82,7 +82,7 @@ class NetworkBuilderTest {
         
         // Test forward pass
         val input = CpuTensorInt32.fromArray(Shape(1, 3), intArrayOf(1, 2, 3))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 2), output.shape)
     }
 
@@ -105,7 +105,7 @@ class NetworkBuilderTest {
         
         // Verify it's using FP32/Float types
         val input = CpuTensorFP32.fromArray(Shape(1, 2), floatArrayOf(1.0f, 2.0f))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
     }
 
@@ -122,7 +122,7 @@ class NetworkBuilderTest {
         
         assertNotNull(network)
         val input = CpuTensorFP32.fromArray(Shape(1, 2), floatArrayOf(1.0f, 2.0f))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
     }
 
@@ -174,7 +174,7 @@ class NetworkBuilderTest {
         
         assertNotNull(network)
         val input = CpuTensorFP32.fromArray(Shape(1, 2), floatArrayOf(1.0f, 2.0f))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
     }
 
@@ -203,7 +203,7 @@ class NetworkBuilderTest {
         
         assertNotNull(network)
         val input = CpuTensorFP32.fromArray(Shape(1, 3), floatArrayOf(1.0f, 2.0f, 3.0f))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 2), output.shape)
     }
 
@@ -225,7 +225,7 @@ class NetworkBuilderTest {
         
         assertNotNull(network)
         val input = CpuTensorFP32.fromArray(Shape(1, 2), floatArrayOf(1.0f, -1.0f))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
     }
 
@@ -250,7 +250,7 @@ class NetworkBuilderTest {
         assertEquals("MLP", network.name)
         
         val input = CpuTensorFP32.fromArray(Shape(1, 2), floatArrayOf(1.0f, 2.0f))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
     }
 
@@ -274,7 +274,7 @@ class NetworkBuilderTest {
         
         // Test forward pass
         val input = CpuTensorFP32.fromArray(Shape(1, 2), floatArrayOf(1.0f, 2.0f))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
     }
 
@@ -297,7 +297,7 @@ class NetworkBuilderTest {
         
         // Test forward pass
         val input = CpuTensorInt32.fromArray(Shape(1, 3), intArrayOf(1, 2, 3))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
     }
 
@@ -321,7 +321,7 @@ class NetworkBuilderTest {
         assertNotNull(network)
         
         val input = CpuTensorFP32.fromArray(Shape(1, 4), floatArrayOf(1.0f, 2.0f, 3.0f, 4.0f))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
     }
 
@@ -358,8 +358,8 @@ class NetworkBuilderTest {
         
         // Test that both networks produce the same output for the same input (deterministic)
         val input = CpuTensorFP32.fromArray(Shape(1, 3), floatArrayOf(1.0f, 2.0f, 3.0f))
-        val output1 = with(network1) { input.forward(input) }
-        val output2 = with(network2) { input.forward(input) }
+        val output1 = network1.forward(input)
+        val output2 = network2.forward(input)
         
         assertEquals(output1.shape, output2.shape)
         // Note: We can't easily test exact equality due to potential floating point differences
@@ -392,7 +392,7 @@ class NetworkBuilderTest {
         
         // Test forward pass
         val input = CpuTensorFP32.fromArray(Shape(1, 5), floatArrayOf(1.0f, 2.0f, 3.0f, 4.0f, 5.0f))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
     }
 
@@ -416,7 +416,7 @@ class NetworkBuilderTest {
         
         // Test forward pass
         val input = CpuTensorFP32.fromArray(Shape(1, 3), floatArrayOf(1.0f, 2.0f, 3.0f))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 2), output.shape)
     }
 
@@ -439,7 +439,7 @@ class NetworkBuilderTest {
         
         // Test forward pass
         val input = CpuTensorInt32.fromArray(Shape(1, 2), intArrayOf(10, 20))
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
     }
 }
diff --git a/skainet-nn/skainet-nn-api/src/commonTest/kotlin/sk/ainet/nn/dsl/NetworkIntegrationTest.kt b/skainet-nn/skainet-nn-api/src/commonTest/kotlin/sk/ainet/nn/dsl/NetworkIntegrationTest.kt
@@ -59,7 +59,7 @@ class NetworkIntegrationTest {
             )
         )
         
-        val output = with(network) { batchInput.forward(batchInput) }
+        val output = network.forward(batchInput)
         assertEquals(Shape(3, 2), output.shape)
     }
 
@@ -95,7 +95,7 @@ class NetworkIntegrationTest {
             byteArrayOf(1, 2, 3, -1, 0, 1)
         )
         
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(2, 2), output.shape)
 
          */
@@ -141,7 +141,7 @@ class NetworkIntegrationTest {
             intArrayOf(10, 20)
         )
         
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(1, 1), output.shape)
 
          */
@@ -238,7 +238,7 @@ class NetworkIntegrationTest {
             )
         )
         
-        val output = with(network) { input.forward(input) }
+        val output = network.forward(input)
         assertEquals(Shape(2, 1), output.shape)
 
          */
diff --git a/skainet-nn/skainet-nn-api/src/commonTest/kotlin/sk/ainet/nn/topology/MLPTest.kt b/skainet-nn/skainet-nn-api/src/commonTest/kotlin/sk/ainet/nn/topology/MLPTest.kt
@@ -36,19 +36,15 @@ class MLPTest {
         val input = CpuTensorFP32.fromArray(Shape(1, 1), floatArrayOf(2.0f))
         
         // Perform inference
-        val output = with(mlp) { 
-            input.forward(input)
-        }
+        val output = mlp.forward(input)
         
         // Verify the output shape and that inference ran without errors
         assertEquals(Shape(1, 1), output.shape)
         assertTrue(output[0, 0] > 0.0f, "Output should be positive due to ReLU activations and positive weights")
         
         // Test with different input
         val input2 = CpuTensorFP32.fromArray(Shape(1, 1), floatArrayOf(-1.0f))
-        val output2 = with(mlp) {
-            input2.forward(input2)
-        }
+        val output2 = mlp.forward(input2)
         
         assertEquals(Shape(1, 1), output2.shape)
         // With negative input and ReLU, intermediate values may be zero, but final output depends on the architecture

Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,9 @@ public class Flatten<T : DType, V>(`
`16`	`16`	`override val modules: List<Module<T, V>>`
`17`	`17`	`get() = emptyList()`
`18`	`18`
`19`		`- override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> {`
`20`		`- return input.flatten(startDim, endDim)`
	`19`	`+ override fun forward(input: Tensor<T, V>): Tensor<T, V> {`
	`20`	`+ return with(input) {`
	`21`	`+ flatten(startDim, endDim)`
	`22`	`+ }`
`21`	`23`	`}`
`22`	`24`	`}`
Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,7 @@ public class Input<T : DType, V>(private val inputShape: Shape, override val nam`
`12`	`12`	`get() = emptyList()`
`13`	`13`
`14`	`14`
`15`		`- override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> {`
	`15`	`+ override fun forward(input: Tensor<T, V>): Tensor<T, V> {`
`16`	`16`	`return input`
`17`	`17`	`}`
`18`	`18`	`}`
Original file line number	Diff line number	Diff line change
`@@ -11,9 +11,9 @@ public abstract class Module<T : DType, V> {`
`11`	`11`
`12`	`12`	`public abstract val modules: List<Module<T, V>>`
`13`	`13`
`14`		`- public abstract fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V>`
	`14`	`+ public abstract fun forward(input: Tensor<T, V>): Tensor<T, V>`
`15`	`15`
`16`		`- public operator fun Tensor<T, V>.invoke(input: Tensor<T, V>): Tensor<T, V> {`
	`16`	`+ public operator fun invoke(input: Tensor<T, V>): Tensor<T, V> {`
`17`	`17`	`return forward(input)`
`18`	`18`	`}`
`19`	`19`	`}`
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@ public class ActivationsWrapperModule<T : DType, V>(`
`14`	`14`	`override val modules: List<Module<T, V>>`
`15`	`15`	`get() = emptyList()`
`16`	`16`
`17`		`- override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> {`
	`17`	`+ override fun forward(input: Tensor<T, V>): Tensor<T, V> {`
`18`	`18`	`return activationHandler(input)`
`19`	`19`	`}`
`20`	`20`	`}`
Original file line number	Diff line number	Diff line change
`@@ -9,8 +9,10 @@ public class Softmax<T : DType, V>(private val dimension: Int, override val name`
`9`	`9`	`override val modules: List<Module<T, V>>`
`10`	`10`	`get() = emptyList()`
`11`	`11`
`12`		`- override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> {`
`13`		`- return input.softmax(dimension)`
	`12`	`+ override fun forward(input: Tensor<T, V>): Tensor<T, V> {`
	`13`	`+ return with(input) {`
	`14`	`+ softmax(dimension)`
	`15`	`+ }`
`14`	`16`	`}`
`15`	`17`	`}`
`16`	`18`
Original file line number	Diff line number	Diff line change
`@@ -9,6 +9,6 @@ public class ReLU<T : DType, V>(override val name: String = "ReLU") : Module<T,`
`9`	`9`	`override val modules: List<Module<T, V>>`
`10`	`10`	`get() = emptyList()`
`11`	`11`
`12`		`- override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> = input.relu()`
	`12`	`+ override fun forward(input: Tensor<T, V>): Tensor<T, V> = with(input) { relu() }`
`13`	`13`	`}`
`14`	`14`
Original file line number	Diff line number	Diff line change
`@@ -12,10 +12,10 @@ public class MLP<T : DType, V>(vararg modules: Module<T, V>, override val name:`
`12`	`12`	`override val modules: List<Module<T, V>>`
`13`	`13`	`get() = modulesList`
`14`	`14`
`15`		`- override fun Tensor<T, V>.forward(input: Tensor<T, V>): Tensor<T, V> {`
	`15`	`+ override fun forward(input: Tensor<T, V>): Tensor<T, V> {`
`16`	`16`	`var tmp = input`
`17`	`17`	`modulesList.forEach { module ->`
`18`		`- tmp = with(module) { this@forward.forward(tmp) }`
	`18`	`+ tmp = module.forward(tmp)`
`19`	`19`	`}`
`20`	`20`	`return tmp`
`21`	`21`	`}`