skainet-nn-dsl

Building blocks for neural networks: the sequential builder for layered models, the DAG builder for graphs with arbitrary wiring, and the rule for choosing between them. This skill is a cheatsheet — DSL skills teach usage rather than enforce constraints.

When to use

• Defining a model with stacked layers (MLP, CNN, simple encoder/decoder).
• Defining a model with skip connections, multiple inputs, or multiple outputs (ResNet block, YOLO head, multi-task heads).
• Reusing a sub-graph as a DagModule.
• Annotating a layer block (conv2d(...) { inChannels = ... }).

When NOT to use

• Building input tensors or preprocessing pipelines for the model — skainet-data-dsl.
• Asserting model output values in tests — skainet-testing.
• Exposing a model builder to Java callers — skainet-java-interop.
• Configuring KSP for tracing wrappers — gradle-multimodule + kmp.

Cheatsheet

Sequential — linear stack of layers

val model = sequential<FP32, Float> {
    input(28 * 28)
    dense(128)
    activation { tensor -> tensor.relu() }
    dense(10)
}
// from: SKaiNET/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/readme/ReadmeSnippetsTest.kt:38-43

Sequential CNN

val cnn = sequential<FP32, Float> {
    input(intArrayOf(1, 28, 28))                                 // C, H, W per sample
    conv2d(outChannels = 16, kernelSize = 3 to 3, stride = 1 to 1, padding = 0 to 0)
    activation { it.relu() }
    maxPool2d(kernelSize = 2 to 2, stride = 2 to 2)
    flatten()
    dense(64)
    activation { it.relu() }
    dense(10)
    softmax(dim = -1)
}

Layer signatures (from NetworkBuilder.kt:113-467):

input(inputSize: Int, id: String = "", requiresGrad: Boolean = false)
input(inputShape: IntArray, id: String = "", requiresGrad: Boolean = false)
dense(outputDimension: Int, id: String = "", content: DENSE<T, V>.() -> Unit = {})
flatten(id: String = "", content: FLATTEN<T, V>.() -> Unit = {})
activation(id: String = "", activation: (Tensor<T, V>) -> Tensor<T, V>)
softmax(dim: Int = -1, id: String = "")
gru(hiddenSize: Int, id: String = "", content: GRU<T, V>.() -> Unit = {})
conv1d(outChannels: Int, kernelSize: Int, stride: Int = 1, padding: Int = 0, dilation: Int = 1, groups: Int = 1, bias: Boolean = true, id: String = "", content: CONV1D<T, V>.() -> Unit = {})
conv2d(outChannels: Int, kernelSize: Pair<Int, Int>, stride: Pair<Int, Int> = 1 to 1, padding: Pair<Int, Int> = 0 to 0, dilation: Pair<Int, Int> = 1 to 1, groups: Int = 1, bias: Boolean = true, id: String = "", content: CONV2D<T, V>.() -> Unit = {})
conv3d(outChannels: Int, kernelSize: Triple<Int, Int, Int>, ...)
maxPool2d(kernelSize: Pair<Int, Int>, stride: Pair<Int, Int> = kernelSize, padding: Pair<Int, Int> = 0 to 0, id: String = "")
avgPool2d(kernelSize: Pair<Int, Int>, stride: Pair<Int, Int> = kernelSize, padding: Pair<Int, Int> = 0 to 0, countIncludePad: Boolean = true, id: String = "")
upsample2d(scale: Pair<Int, Int> = 2 to 2, mode: UpsampleMode = UpsampleMode.Nearest, alignCorners: Boolean = false, id: String = "")
batchNorm(numFeatures: Int, eps: Double = 1e-5, momentum: Double = 0.1, affine: Boolean = true, id: String = "")
groupNorm(numGroups: Int, numChannels: Int, eps: Double = 1e-5, affine: Boolean = true, id: String = "")
layerNorm(normalizedShape: IntArray, eps: Double = 1e-5, elementwiseAffine: Boolean = true, id: String = "")
sequential(content: NeuralNetworkDsl<T, V>.() -> Unit)   // nested grouping block
stage(id: String, content: NeuralNetworkDsl<T, V>.() -> Unit)
stage<TStage : DType>(id: String, content: NeuralNetworkDsl<TStage, V>.() -> Unit)  // precision-scoped stage

conv2d, maxPool2d and upsample2d have a "content-block-only" overload (conv2d(id) { outChannels = ...; kernelSize(5); stride(1); padding(2) }) — conv1d, conv3d and avgPool2d do not. Use it when configuration is verbose enough to warrant block syntax.

LLM/transformer layer builders (embedding, rmsNorm, multiHeadAttention, swiGluFFN, residual, xielu) are NOT in this DSL — they moved to the SKaiNET-transformers repository's llm-core module (NetworkBuilder.kt:422-423 notes the move). The fused attention op stayed in the engine: ops.scaledDotProductAttention(query, key, value, mask = null, scale = 0f, causal = false) (TensorOps.kt:389-397) with native grouped-query attention — K/V come as [batch, nKVHeads, kvLen, headDim], nKVHeads must divide the query's nHeads, and K/V are never tiled up to the query head count.

Activation as a separate layer vs function

// As a layer with id (preferred when traceability matters)
activation(id = "post_dense_1") { it.relu() }

// As an activation function inline (not a separate layer in the topology)
val out = ctx.relu(dense_out)  // operator on a Tensor

DAG — arbitrary wiring

val program = dag {
    val x = input<FP32>("x", TensorSpec("x", listOf(1, 4), "FP32"))
    val w = parameter<FP32, Float>("w") { shape(4, 4) { ones() } }
    val mm = matmul(x, w)
    val y = relu(mm)
    output(y)
}
// from: SKaiNET/skainet-lang/skainet-lang-dag/src/commonMain/kotlin/sk/ainet/lang/dag/GraphDsl.kt:43-70

DAG node helpers (from GraphDsl.kt):

input<T : DType>(name: String, spec: TensorSpec = TensorSpec(name, null, "unknown")): GraphValue<T>
parameter<T : DType>(name: String, spec: TensorSpec): GraphValue<T>
parameter<reified T : DType, V>(name: String, builder: SymbolicTensorBuilder<T>.() -> TensorSpec): GraphValue<T>
constant<T : DType>(name: String, spec: TensorSpec): GraphValue<T>
constant<reified T : DType, V>(name: String, builder: SymbolicTensorBuilder<T>.() -> TensorSpec): GraphValue<T>
op(operation: Operation, inputs: List<GraphValue<*>>, id: String = "", attributes: Map<String, Any?> = emptyMap()): List<GraphValue<*>>
output(vararg values: GraphValue<*>)

The DAG builder is definition-only — no tensors are allocated. The returned GraphProgram is consumed by skainet-compile-dag to produce a ComputeGraph.

Schedule and dtype hints on DAG nodes

Two annotation DSLs ride a node's attributes without changing what the graph computes:

dag {
    schedule(parallel("heads")) {                      // every op recorded inside carries the hint
        op(sdpa, listOf(q, k, v))
    }
    op(matmul, listOf(x, w), schedule = parallel("rows", parallelism = 8))   // single op
    op(matmul, listOf(x, w), dtypePolicy = DTypePolicy.Require(BF16))        // dtype constraint
}
// from: SKaiNET/skainet-lang/skainet-lang-dag/src/commonMain/kotlin/sk/ainet/lang/dag/ScheduleDsl.kt:25-49
// from: SKaiNET/skainet-lang/skainet-lang-dag/src/commonMain/kotlin/sk/ainet/lang/dag/DtypePolicyDsl.kt:1-64

parallel(vararg dims, parallelism = null) builds a ScheduleHint (sk.ainet.context.schedule.ScheduleHint). ScheduleAnnotationPass (skainet-compile-opt) validates hints per op, and the StableHLO export emits them as the skainet.schedule module attribute. Dtype policies feed DTypeConstraintResolutionPass.

Reusable sub-graphs — dagModule

val residualBlock = dagModule { inputs ->
    val x = inputs[0]
    val w1 = parameter<FP32, Float>("w1") { shape(64, 64) { randn(0f, 0.02f) } }
    val w2 = parameter<FP32, Float>("w2") { shape(64, 64) { randn(0f, 0.02f) } }
    val h = relu(matmul(x, w1))
    val y = matmul(h, w2)
    listOf(/* residual add */ add(x, y))
}

val program = dag {
    val x = input<FP32>("x", TensorSpec("x", listOf(1, 64), "FP32"))
    val out = module(residualBlock, listOf(x))
    output(out[0])
}
// from: SKaiNET/skainet-lang/skainet-lang-dag/src/commonMain/kotlin/sk/ainet/lang/dag/GraphDsl.kt:419-457

Decision rule — sequential or DAG?

sequential produces a Module<T, V> you can call forward(x, ctx) on. dag produces a GraphProgram for downstream compilation.

Workflow

1. Pick the builder via the table above.
2. Choose <T, V> from the dtype/value-type table (FP32 ↔ Float, Int8 ↔ Byte, …).
3. For sequential: declare input(...) first; chain layers; finish with the prediction layer (often dense(numClasses) followed by softmax(dim = -1)).
4. For DAG: declare every input and parameter upfront, build intermediate values, mark outputs with output(...).
5. Run a forward pass in a test using skainet-testing patterns to verify shapes and tolerances.

Related skills

• The input tensor for forward(x, ctx) and parameter tensors used in parameter { shape(…) { … } } come from the data DSL — see ../skainet-data-dsl/SKILL.md.
• Loading pre-trained weights into a parameter-style graph — see ../skainet-model-loading/SKILL.md.
• Running the forward pass: ExecutionContext, batching, threading — see ../skainet-inference/SKILL.md.
• Calling a network from Java — see ../skainet-java-consumer/SKILL.md.
• (In-repo only) Asserting forward-pass output values in tests — see the contributor skainet-testing skill.

Anti-patterns

// WRONG — calling relu() as a layer (no such layer; relu is an activation function)
sequential<FP32, Float> {
    input(28 * 28)
    dense(128)
    relu()              // does not exist on NeuralNetworkDsl
}

// RIGHT — wrap with `activation { ... }` or call .relu() inline on a tensor
sequential<FP32, Float> {
    input(28 * 28)
    dense(128)
    activation { it.relu() }
}

// WRONG — using sequential for a residual block
sequential<FP32, Float> {
    input(intArrayOf(64))
    dense(64)
    activation { it.relu() }
    // ... how do I add the input back to this output? sequential cannot.
}

// RIGHT — DAG for skip connections
dag {
    val x = input<FP32>("x", TensorSpec("x", listOf(1, 64), "FP32"))
    val w = parameter<FP32, Float>("w") { shape(64, 64) { randn(0f, 0.02f) } }
    val h = relu(matmul(x, w))
    val y = add(x, h)              // residual
    output(y)
}

// WRONG — placing the input declaration after a layer
sequential<FP32, Float> {
    dense(128)         // there is no inferred input shape yet
    input(28 * 28)
}

// RIGHT — input first
sequential<FP32, Float> {
    input(28 * 28)
    dense(128)
}

References

• references/layers.md — every public layer signature on NeuralNetworkDsl<T, V>, with full default values.
• references/dag-nodes.md — every public method on DagBuilder, with shape-inference notes.
• references/sequential-vs-dag.md — extended decision tree, including the cases where you start with sequential and migrate to DAG mid-development.