Add flatten to DSL

Related-To: #23

Author: michalharakal

build.gradle.kts

@@ -4,8 +4,6 @@ plugins {
     alias(libs.plugins.jetbrainsKotlinJvm) apply false
     alias(libs.plugins.binaryCompatibility) apply false
     alias(libs.plugins.modulegraph.souza) apply true
-
-
 }
 
 allprojects {

core/src/commonMain/kotlin/sk/ai/net/dsl/NetworkBuilder.kt

@@ -3,6 +3,7 @@ package sk.ai.net.dsl
 import sk.ai.net.nn.activations.ActivationsWrapperModule
 import sk.ai.net.Shape
 import sk.ai.net.Tensor
+import sk.ai.net.nn.Flatten
 import sk.ai.net.nn.Input
 import sk.ai.net.nn.Linear
 import sk.ai.net.nn.Module
@@ -25,6 +26,8 @@ interface NetworkDslItem
 interface NeuralNetworkDsl : NetworkDslItem {
     fun input(inputSize: Int, id: String = "")
 
+    fun flatten(id: String = "", content: FLATTEN.() -> Unit = {})
+
     fun dense(outputDimension: Int, id: String = "", content: DENSE.() -> Unit = {})
 }
 
@@ -35,6 +38,13 @@ interface DENSE : NetworkDslItem {
     fun bias(initBlock: (Shape) -> Tensor)
 }
 
+@NetworkDsl
+interface FLATTEN : NetworkDslItem {
+    var startDim: Int
+    var endDim: Int
+}
+
+
 private fun getDefaultName(id: String, s: String, size: Int): String {
     if (id.isNotEmpty()) return id
     return "$s-$size"
@@ -69,6 +79,16 @@ fun createLinear(
     }
 }
 
+class FlattenImpl(
+    override var startDim: Int = 1,
+    override var endDim: Int = -1,
+    private val id: String
+) : FLATTEN {
+    fun create(): Module {
+        return Flatten(startDim, endDim, id)
+    }
+}
+
 class DenseImpl(
     private val inputDimension: Int, private val outputDimension: Int, private val id: String
 ) : DENSE {
@@ -119,6 +139,14 @@ private class NeuralNetworkDslImpl : NeuralNetworkDsl {
         modules.add(Input(Shape(inputSize), name = getDefaultName(id, "Input", modules.size)))
     }
 
+    override fun flatten(id: String, content: FLATTEN.() -> Unit) {
+        val impl = FlattenImpl(
+            id = getDefaultName(id, "flatten", modules.size)
+        )
+        impl.content()
+        modules += impl.create()
+    }
+
     override fun dense(outputDimension: Int, id: String, content: DENSE.() -> Unit) {
         val inputDimension = lastDimension
         lastDimension = outputDimension

settings.gradle.kts

@@ -16,4 +16,4 @@ dependencyResolutionManagement {
 rootProject.name = "skainet"
 include(":core")
 include(":io")
-include(":gguf")
+include(":gguf")