diff --git a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api index 6425c257a..a912cfe79 100644 --- a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api +++ b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api @@ -314,6 +314,7 @@ protected final class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase$CpuTensor : sk/ public fun getDtype ()Lkotlin/reflect/KClass; public fun getGrad ()Lsk/ainet/lang/tensor/Tensor; public fun getGradState ()Lsk/ainet/lang/tensor/GradState; + public fun getId ()Lsk/ainet/lang/tensor/TensorId; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; public fun getRank ()I public fun getRequiresGrad ()Z diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 29be4c0d9..80757e161 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -665,6 +665,12 @@ public final class sk/ainet/lang/memory/AllocationSpec$Companion { public static synthetic fun of$default (Lsk/ainet/lang/memory/AllocationSpec$Companion;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/memory/ScopeKind;ZIILjava/lang/Object;)Lsk/ainet/lang/memory/AllocationSpec; } +public final class sk/ainet/lang/memory/DescribeKt { + public static final fun describe (Lsk/ainet/lang/tensor/Tensor;)Ljava/lang/String; + public static final fun describe (Lsk/ainet/lang/tensor/storage/TensorStorage;Lsk/ainet/lang/tensor/TensorId;)Ljava/lang/String; + public static synthetic fun describe$default (Lsk/ainet/lang/tensor/storage/TensorStorage;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Ljava/lang/String; +} + public abstract interface annotation class sk/ainet/lang/memory/ExperimentalMemoryApi : java/lang/annotation/Annotation { } @@ -2368,6 +2374,26 @@ public final class sk/ainet/lang/nn/topology/Parameter { public fun toString ()Ljava/lang/String; } +public final class sk/ainet/lang/nn/topology/TensorIdAssignment { + public fun (Ljava/util/Map;Ljava/util/List;)V + public final fun component1 ()Ljava/util/Map; + public final fun component2 ()Ljava/util/List; + public final fun copy (Ljava/util/Map;Ljava/util/List;)Lsk/ainet/lang/nn/topology/TensorIdAssignment; + public static synthetic fun copy$default (Lsk/ainet/lang/nn/topology/TensorIdAssignment;Ljava/util/Map;Ljava/util/List;ILjava/lang/Object;)Lsk/ainet/lang/nn/topology/TensorIdAssignment; + public fun equals (Ljava/lang/Object;)Z + public final fun get (Ljava/lang/String;)Lsk/ainet/lang/tensor/Tensor; + public final fun get (Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/tensor/Tensor; + public final fun getNotCarried ()Ljava/util/List; + public final fun getTensors ()Ljava/util/Map; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/nn/topology/TensorIdsKt { + public static final fun assignTensorIds (Lsk/ainet/lang/nn/topology/ModuleNode;Ljava/lang/String;)Lsk/ainet/lang/nn/topology/TensorIdAssignment; + public static synthetic fun assignTensorIds$default (Lsk/ainet/lang/nn/topology/ModuleNode;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/nn/topology/TensorIdAssignment; +} + public final class sk/ainet/lang/nn/topology/TraversalKt { public static final fun bindPaths (Lsk/ainet/lang/nn/topology/ModuleNode;Ljava/lang/String;Ljava/lang/String;)V public static synthetic fun bindPaths$default (Lsk/ainet/lang/nn/topology/ModuleNode;Ljava/lang/String;Ljava/lang/String;ILjava/lang/Object;)V @@ -2699,6 +2725,7 @@ public final class sk/ainet/lang/tensor/SlicedTensorView : sk/ainet/lang/tensor/ public fun getDtype ()Lkotlin/reflect/KClass; public fun getGrad ()Lsk/ainet/lang/tensor/Tensor; public fun getGradState ()Lsk/ainet/lang/tensor/GradState; + public fun getId ()Lsk/ainet/lang/tensor/TensorId; public fun getIndexMapping ()Lsk/ainet/lang/tensor/IndexMapper; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; public fun getParentTensor ()Lsk/ainet/lang/tensor/Tensor; @@ -2726,6 +2753,7 @@ public abstract interface class sk/ainet/lang/tensor/Tensor { public abstract fun getDtype ()Lkotlin/reflect/KClass; public fun getGrad ()Lsk/ainet/lang/tensor/Tensor; public abstract fun getGradState ()Lsk/ainet/lang/tensor/GradState; + public fun getId ()Lsk/ainet/lang/tensor/TensorId; public abstract fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; public fun getRank ()I public fun getRequiresGrad ()Z @@ -2737,6 +2765,7 @@ public abstract interface class sk/ainet/lang/tensor/Tensor { public final class sk/ainet/lang/tensor/Tensor$DefaultImpls { public static fun accumulateGrad (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V public static fun getGrad (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public static fun getId (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/TensorId; public static fun getRank (Lsk/ainet/lang/tensor/Tensor;)I public static fun getRequiresGrad (Lsk/ainet/lang/tensor/Tensor;)Z public static fun getShape (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Shape; @@ -2827,6 +2856,34 @@ public abstract interface class sk/ainet/lang/tensor/TensorFactory { public abstract fun zeros (Lsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/Tensor; } +public final class sk/ainet/lang/tensor/TensorId { + public static final field Companion Lsk/ainet/lang/tensor/TensorId$Companion; + public fun (Ljava/util/List;Ljava/lang/String;Ljava/lang/String;)V + public synthetic fun (Ljava/util/List;Ljava/lang/String;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun equals (Ljava/lang/Object;)Z + public final fun getCanonical ()Ljava/lang/String; + public final fun getDiscriminator ()Ljava/lang/String; + public final fun getModulePath ()Ljava/util/List; + public final fun getParameter ()Ljava/lang/String; + public fun hashCode ()I + public final fun legacyPath (Ljava/lang/String;)Ljava/lang/String; + public static synthetic fun legacyPath$default (Lsk/ainet/lang/tensor/TensorId;Ljava/lang/String;ILjava/lang/Object;)Ljava/lang/String; + public fun toString ()Ljava/lang/String; + public final fun view (Ljava/lang/String;)Lsk/ainet/lang/tensor/TensorId; + public final fun withDiscriminator (Ljava/lang/String;)Lsk/ainet/lang/tensor/TensorId; +} + +public final class sk/ainet/lang/tensor/TensorId$Companion { + public final fun fromLegacyPath (Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Lsk/ainet/lang/tensor/TensorId; + public static synthetic fun fromLegacyPath$default (Lsk/ainet/lang/tensor/TensorId$Companion;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/tensor/TensorId; + public final fun parse (Ljava/lang/String;)Lsk/ainet/lang/tensor/TensorId; +} + +public abstract interface class sk/ainet/lang/tensor/TensorIdBearer { + public abstract fun getId ()Lsk/ainet/lang/tensor/TensorId; + public abstract fun setId (Lsk/ainet/lang/tensor/TensorId;)V +} + public final class sk/ainet/lang/tensor/TensorSliceBuilder { public fun ()V public final fun segment (Lkotlin/jvm/functions/Function1;)V @@ -2853,6 +2910,7 @@ public abstract interface class sk/ainet/lang/tensor/TensorView : sk/ainet/lang/ public final class sk/ainet/lang/tensor/TensorView$DefaultImpls { public static fun accumulateGrad (Lsk/ainet/lang/tensor/TensorView;Lsk/ainet/lang/tensor/Tensor;)V public static fun getGrad (Lsk/ainet/lang/tensor/TensorView;)Lsk/ainet/lang/tensor/Tensor; + public static fun getId (Lsk/ainet/lang/tensor/TensorView;)Lsk/ainet/lang/tensor/TensorId; public static fun getRank (Lsk/ainet/lang/tensor/TensorView;)I public static fun getRequiresGrad (Lsk/ainet/lang/tensor/TensorView;)Z public static fun getShape (Lsk/ainet/lang/tensor/TensorView;)Lsk/ainet/lang/tensor/Shape; @@ -2909,7 +2967,7 @@ public final class sk/ainet/lang/tensor/ViewManagementExtensionsKt { public static final fun isContiguous (Lsk/ainet/lang/tensor/TensorView;)Z } -public final class sk/ainet/lang/tensor/VoidOpsTensor : sk/ainet/lang/tensor/Tensor { +public final class sk/ainet/lang/tensor/VoidOpsTensor : sk/ainet/lang/tensor/Tensor, sk/ainet/lang/tensor/TensorIdBearer { public fun (Lsk/ainet/lang/tensor/data/TensorData;Lkotlin/reflect/KClass;Lsk/ainet/lang/tensor/GradState;)V public synthetic fun (Lsk/ainet/lang/tensor/data/TensorData;Lkotlin/reflect/KClass;Lsk/ainet/lang/tensor/GradState;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun accumulateGrad (Lsk/ainet/lang/tensor/Tensor;)V @@ -2917,11 +2975,13 @@ public final class sk/ainet/lang/tensor/VoidOpsTensor : sk/ainet/lang/tensor/Ten public fun getDtype ()Lkotlin/reflect/KClass; public fun getGrad ()Lsk/ainet/lang/tensor/Tensor; public fun getGradState ()Lsk/ainet/lang/tensor/GradState; + public fun getId ()Lsk/ainet/lang/tensor/TensorId; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; public fun getRank ()I public fun getRequiresGrad ()Z public fun getShape ()Lsk/ainet/lang/tensor/Shape; public fun getVolume ()I + public fun setId (Lsk/ainet/lang/tensor/TensorId;)V public fun zeroGrad ()V } @@ -4123,7 +4183,7 @@ public final class sk/ainet/lang/tensor/operators/OperatorsKt { public static final fun withOps (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/ops/TensorOps;)Lsk/ainet/lang/tensor/Tensor; } -public final class sk/ainet/lang/tensor/operators/OpsBoundTensor : sk/ainet/lang/tensor/Tensor { +public final class sk/ainet/lang/tensor/operators/OpsBoundTensor : sk/ainet/lang/tensor/Tensor, sk/ainet/lang/tensor/TensorIdBearer { public static final field Companion Lsk/ainet/lang/tensor/operators/OpsBoundTensor$Companion; public fun (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/ops/TensorOps;)V public fun accumulateGrad (Lsk/ainet/lang/tensor/Tensor;)V @@ -4131,12 +4191,14 @@ public final class sk/ainet/lang/tensor/operators/OpsBoundTensor : sk/ainet/lang public fun getDtype ()Lkotlin/reflect/KClass; public fun getGrad ()Lsk/ainet/lang/tensor/Tensor; public fun getGradState ()Lsk/ainet/lang/tensor/GradState; + public fun getId ()Lsk/ainet/lang/tensor/TensorId; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; public final fun getOrigin ()Lsk/ainet/lang/tensor/Tensor; public fun getRank ()I public fun getRequiresGrad ()Z public fun getShape ()Lsk/ainet/lang/tensor/Shape; public fun getVolume ()I + public fun setId (Lsk/ainet/lang/tensor/TensorId;)V public fun zeroGrad ()V } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Describe.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Describe.kt new file mode 100644 index 000000000..63494d4f2 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Describe.kt @@ -0,0 +1,52 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.storage.BufferHandle +import sk.ainet.lang.tensor.storage.TensorStorage + +/** + * One-line debugger rendering of a tensor (SKEEP-003 §4.7, PRD M0-F5): + * `TensorId · Format · Shape · storage kind · origin · scope · StorageId`. + * Fields that only exist from milestone M1 on (scope, storage id) print as `—`. + * + * Example: `model.layers.blk.3.attn.q_proj.weight · Float32/Q4_K · [2048, 2048] · Q4_KBlockTensorData · — · scope — · storage —` + */ +@ExperimentalMemoryApi +public fun Tensor<*, *>.describe(): String = buildString { + append(id?.canonical ?: "—"); append(SEP) + append(formatOrNull?.toString() ?: "?"); append(SEP) + append(shapeText(data.shape.dimensions)); append(SEP) + append(data::class.simpleName ?: "data"); append(SEP) + append("—"); append(SEP) // origin (file/offset) — storage-backed tensors only, M1 + append("scope —"); append(SEP) + append("storage —") +} + +/** Same rendering for a storage descriptor; origin is the file for file-backed buffers. */ +@ExperimentalMemoryApi +public fun TensorStorage.describe(id: sk.ainet.lang.tensor.TensorId? = null): String = buildString { + append(id?.canonical ?: "—"); append(SEP) + append(format.toString()); append(SEP) + append(shapeText(shape.dimensions)); append(SEP) + append( + when (val b = buffer) { + is BufferHandle.Owned -> "Owned" + is BufferHandle.Borrowed -> "Borrowed" + is BufferHandle.Aliased -> "Aliased" + is BufferHandle.FileBacked -> "Mapped" + is BufferHandle.DeviceResident -> "Device(${b.deviceId})" + } + ); append(SEP) + append( + when (val b = buffer) { + is BufferHandle.FileBacked -> "${b.path} @0x${b.fileOffset.toString(16)}" + else -> "—" + } + ); append(SEP) + append("scope —"); append(SEP) + append("storage —") +} + +private const val SEP = " · " + +private fun shapeText(dims: IntArray): String = dims.joinToString(", ", "[", "]") diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/topology/TensorIds.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/topology/TensorIds.kt new file mode 100644 index 000000000..03df49214 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/topology/TensorIds.kt @@ -0,0 +1,61 @@ +package sk.ainet.lang.nn.topology + +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.TensorIdBearer + +/** + * Result of [assignTensorIds]: every parameter tensor of the tree keyed by its [TensorId], and the + * ids whose tensor could not carry the id (not a [TensorIdBearer]) — they are still in [tensors]. + */ +public data class TensorIdAssignment( + val tensors: Map>, + val notCarried: List, +) { + public operator fun get(id: TensorId): Tensor<*, *>? = tensors[id] + public operator fun get(canonical: String): Tensor<*, *>? = tensors[TensorId.parse(canonical)] +} + +/** + * Assign a [TensorId] to every parameter tensor in this module tree, derived from the module + * structure: the id's module path is the chain of module names from [root] down (the same segments + * [bindPaths] uses, so `id.legacyPath()` equals the node's `path`), and its parameter is the + * parameter's short name (`weight` for a parameter registered as `".weight"`). + * + * Idempotent — running it twice yields the same ids — and free of side effects on tensors that + * cannot carry an id (reported in [TensorIdAssignment.notCarried]). No runtime code path changes: + * an id is metadata on the parameter tensor. + * + * ``` + * val ids = model.assignTensorIds("model") + * ids["model.layers.blk.0.attn.weight"] // the tensor + * ``` + */ +public fun ModuleNode.assignTensorIds(root: String = name): TensorIdAssignment { + val tensors = LinkedHashMap>() + val notCarried = ArrayList() + fun visit(node: ModuleNode, segments: List) { + for (p in node.params) { + val id = TensorId(segments, parameterShortName(node, p.name)) + val t = p.value + tensors[id] = t + if (t is TensorIdBearer) t.id = id else notCarried += id + } + for (child in node.children) { + val seg = child.name.ifEmpty { child.id } + visit(child, if (seg.isEmpty()) segments else segments + seg) + } + } + visit(this, if (root.isEmpty()) emptyList() else listOf(root)) + return TensorIdAssignment(tensors, notCarried) +} + +/** `".weight"` → `weight`; a name without the module prefix is returned unchanged. */ +internal fun parameterShortName(node: ModuleNode, parameterName: String): String { + val prefix = node.name + "." + return when { + node.name.isNotEmpty() && parameterName.startsWith(prefix) && parameterName.length > prefix.length -> + parameterName.substring(prefix.length) + else -> parameterName + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Tensor.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Tensor.kt index e4a16af7d..7f51808a4 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Tensor.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Tensor.kt @@ -81,6 +81,13 @@ public interface Tensor { */ public val dtype: KClass + /** + * Stable identity of this tensor within its model (`model.layers[3].attn.q_proj.weight`), or + * `null` for an anonymous tensor. Assigned from the module tree by `assignTensorIds`; carried by + * implementations of [TensorIdBearer]. Default `null` so existing implementations are unaffected. + */ + public val id: TensorId? get() = null + /** * Gradient state tied to this tensor instance. */ diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/TensorId.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/TensorId.kt new file mode 100644 index 000000000..7c1e11e19 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/TensorId.kt @@ -0,0 +1,88 @@ +package sk.ainet.lang.tensor + +/** + * Stable, human-readable identity of a tensor derived from the model structure (SKEEP-003 §4.7): + * the module path inside the `network { }` / `dag { }` tree, the parameter name, and an optional + * discriminator for activations (`#step=17`). The *same* id names the tensor whether it is + * materialized or symbolic, eager or compiled, and is what loaders' `NameMap`s, graph nodes + * (`loc("…")`) and memory-debugging tools key on. + * + * Ids are optional — an anonymous `a matMul b` in a notebook has none. + * + * Two string forms: + * - [canonical] — dotted, for logs, `loc()` attributes and greps: `model.layers[3].attn.q_proj.weight`; + * - [legacyPath] — the slash-separated module path used by `ModuleNode.path` / `bindPaths` and the + * `WeightNameResolver`s (`MLP/blk.0/attn`), so existing name-based loaders keep working. + * + * [parse] inverts [canonical] segment-wise; module names that themselves contain `.` (e.g. `blk.0`) + * round-trip as a string but split into more segments than they were built from. For that reason + * **equality and hashing are defined on [canonical]**: two ids that print the same are the same id, + * however they were built (`TensorId.parse(id.canonical) == id` always holds). + * + * @property modulePath module names from the root down to the owner of the parameter + * @property parameter the parameter's short name within its module (`weight`, `bias`, `weight_ih`) + * @property discriminator optional suffix for activations / repeated instances + */ +public class TensorId( + public val modulePath: List, + public val parameter: String, + public val discriminator: String? = null, +) { + init { + require(parameter.isNotEmpty()) { "TensorId.parameter must not be empty" } + require(modulePath.none { it.isEmpty() }) { "TensorId.modulePath segments must not be empty: $modulePath" } + } + + /** `model.layers[3].attn.q_proj.weight` (+ `#discriminator`). */ + public val canonical: String + get() = buildString { + for (s in modulePath) { append(s); append('.') } + append(parameter) + if (discriminator != null) { append('#'); append(discriminator) } + } + + /** The slash-separated module path (`MLP/blk.0/attn`), empty string for a root-level parameter. */ + public fun legacyPath(separator: String = "/"): String = modulePath.joinToString(separator) + + /** The same id with a different [discriminator] (e.g. `withDiscriminator("step=17")`). */ + public fun withDiscriminator(discriminator: String?): TensorId = TensorId(modulePath, parameter, discriminator) + + /** Id of a sub-view of this tensor, e.g. `kv.layers[3].k` → `kv.layers[3].k[1024..2048]`. */ + public fun view(range: String): TensorId = TensorId(modulePath, "$parameter[$range]", discriminator) + + override fun toString(): String = canonical + + /** Equal iff the canonical strings are equal (see the class note). */ + override fun equals(other: Any?): Boolean = other is TensorId && other.canonical == canonical + + override fun hashCode(): Int = canonical.hashCode() + + public companion object { + /** + * Parse a [canonical] string: the last dotted segment is the parameter, the preceding + * segments the module path, an optional `#…` suffix the discriminator. + * @throws IllegalArgumentException for an empty or malformed string + */ + public fun parse(canonical: String): TensorId { + require(canonical.isNotBlank()) { "TensorId must not be blank" } + val hash = canonical.indexOf('#') + val body = if (hash >= 0) canonical.substring(0, hash) else canonical + val disc = if (hash >= 0) canonical.substring(hash + 1).ifEmpty { null } else null + val segments = body.split('.') + require(segments.all { it.isNotEmpty() }) { "Malformed TensorId '$canonical'" } + return TensorId(segments.dropLast(1), segments.last(), disc) + } + + /** Build from a slash-separated module path (as produced by `bindPaths`) and a parameter name. */ + public fun fromLegacyPath(path: String?, parameter: String, separator: String = "/"): TensorId = + TensorId(path?.takeIf { it.isNotEmpty() }?.split(separator) ?: emptyList(), parameter) + } +} + +/** + * A tensor that can carry a [TensorId]. The core tensor implementations implement it; `assignTensorIds` + * sets ids through this interface and reports tensors that cannot carry one. + */ +public interface TensorIdBearer { + public var id: TensorId? +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/VoidOpsTensor.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/VoidOpsTensor.kt index ff9f80906..506a24326 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/VoidOpsTensor.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/VoidOpsTensor.kt @@ -13,7 +13,9 @@ public class VoidOpsTensor( override val data: TensorData, override val dtype: KClass, override val gradState: GradState = GradState() -) : Tensor { +) : Tensor, TensorIdBearer { override val ops: TensorOps get() = VoidTensorOps() + + override var id: TensorId? = null } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/operators/Operators.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/operators/Operators.kt index 30a8cdeee..dff8ac963 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/operators/Operators.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/operators/Operators.kt @@ -2,6 +2,8 @@ package sk.ainet.lang.tensor.operators import sk.ainet.lang.tensor.GradState import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.TensorIdBearer import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.tensor.ops.TensorOps import sk.ainet.lang.types.DType @@ -14,12 +16,22 @@ import kotlin.reflect.KClass public class OpsBoundTensor( public val origin: Tensor, private val opsRef: TensorOps, -) : Tensor { +) : Tensor, TensorIdBearer { override val data: TensorData get() = origin.data override val dtype: KClass get() = origin.dtype override val gradState: GradState get() = origin.gradState override val ops: TensorOps get() = opsRef + private var localId: TensorId? = null + + /** The origin's id when it carries one (ids survive re-binding to another context), else a local one. */ + override var id: TensorId? + get() = origin.id ?: localId + set(value) { + val o = origin + if (o is TensorIdBearer) o.id = value else localId = value + } + override fun accumulateGrad(g: Tensor) { origin.accumulateGrad(g) } @@ -30,11 +42,12 @@ public class OpsBoundTensor( public companion object { public fun fromData(data: TensorData, dtype: KClass, ops: TensorOps): OpsBoundTensor { - val origin = object : Tensor { + val origin = object : Tensor, TensorIdBearer { override val data: TensorData = data override val dtype: KClass = dtype override val ops: TensorOps = ops override val gradState: GradState = GradState() + override var id: TensorId? = null } return OpsBoundTensor(origin, ops) } diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/DescribeTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/DescribeTest.kt new file mode 100644 index 000000000..61e612432 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/DescribeTest.kt @@ -0,0 +1,53 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.lang.tensor.data.DenseFloatArrayTensorData +import sk.ainet.lang.tensor.data.Q4_KBlockTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.storage.BufferHandle +import sk.ainet.lang.tensor.storage.Placement +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.tensor.storage.TensorStorage +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +@OptIn(ExperimentalMemoryApi::class) +class DescribeTest { + + @Test + fun tensorRendererPrintsIdFormatShapeKindAndPlaceholders() { + val t = VoidOpsTensor(DenseFloatArrayTensorData(Shape(2, 3), FloatArray(6)), FP32::class) + val anon = t.describe().split(" · ") + assertEquals(7, anon.size) + assertEquals("—", anon[0]); assertEquals("Float32/Dense(4B)", anon[1]); assertEquals("[2, 3]", anon[2]) + assertEquals("—", anon[4]); assertEquals("scope —", anon[5]); assertEquals("storage —", anon[6]) + + t.id = TensorId(listOf("model", "layers", "blk.3", "attn"), "q_proj.weight") + assertTrue(t.describe().startsWith("model.layers.blk.3.attn.q_proj.weight · Float32/Dense(4B) · [2, 3] · ")) + } + + @Test + fun packedWeightRendersLogicalDtypeAndEncoding() { + @Suppress("UNCHECKED_CAST") + val q = VoidOpsTensor(Q4_KBlockTensorData(Shape(1, 256), ByteArray(144)) as TensorData, FP32::class) + assertTrue(q.describe().contains(" · Float32/Q4_K · [1, 256] · "), q.describe()) + } + + @Test + fun storageRendererShowsKindAndOrigin() { + val mapped = TensorStorage( + Shape(2048, 2048), FP32, TensorEncoding.Q4_K, + BufferHandle.FileBacked("/models/model.gguf", 0x1A3F000L, 144L * 16384), Placement.MMAP_WEIGHTS, + ) + assertEquals( + "model.layers.blk.3.attn.q_proj.weight · Float32/Q4_K · [2048, 2048] · Mapped · /models/model.gguf @0x1a3f000 · scope — · storage —", + mapped.describe(TensorId.parse("model.layers.blk.3.attn.q_proj.weight")), + ) + val owned = TensorStorage(Shape(4), FP32, TensorEncoding.Dense(4), BufferHandle.Owned(ByteArray(16))) + assertEquals("— · Float32/Dense(4B) · [4] · Owned · — · scope — · storage —", owned.describe()) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/topology/AssignTensorIdsTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/topology/AssignTensorIdsTest.kt new file mode 100644 index 000000000..0b2d64ba7 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/topology/AssignTensorIdsTest.kt @@ -0,0 +1,117 @@ +package sk.ainet.lang.nn.topology + +import sk.ainet.lang.nn.Module +import sk.ainet.lang.tensor.GradState +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.lang.tensor.data.DenseFloatArrayTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.operators.OpsBoundTensor +import sk.ainet.lang.tensor.ops.TensorOps +import sk.ainet.lang.tensor.ops.VoidTensorOps +import sk.ainet.lang.types.FP32 +import kotlin.reflect.KClass +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNull +import kotlin.test.assertSame +import kotlin.test.assertTrue + +class AssignTensorIdsTest { + + private fun tensor(): VoidOpsTensor = + VoidOpsTensor(DenseFloatArrayTensorData(Shape(2, 2), FloatArray(4)), FP32::class) + + /** A tensor that cannot carry an id (not a TensorIdBearer). */ + private fun bareTensor(): Tensor = object : Tensor { + override val data: TensorData = DenseFloatArrayTensorData(Shape(1), FloatArray(1)) + override val dtype: KClass = FP32::class + override val ops: TensorOps = VoidTensorOps() + override val gradState: GradState = GradState() + } + + private class Leaf( + override val name: String, + weight: Tensor, + bias: Tensor? = null, + ) : Module(), ModuleParameters { + override val modules: List> = emptyList() + override val params: List> = buildList { + add(ModuleParameter.WeightParameter("$name.weight", weight)) + if (bias != null) add(ModuleParameter.BiasParameter("$name.bias", bias)) + } + } + + private class Block(override val name: String, override val modules: List>) : Module() + + @Test + fun idsFollowTheModuleTreeAndParameterShortNames() { + val qW = tensor(); val qB = tensor(); val mlpW = tensor() + val model = Block("model", listOf( + Block("layers", listOf( + Block("blk.0", listOf(Leaf("attn", qW, qB), Leaf("mlp", mlpW))), + )), + )) + val ids = model.assignTensorIds() + + assertEquals(3, ids.tensors.size) + assertTrue(ids.notCarried.isEmpty()) + assertSame(qW, ids["model.layers.blk.0.attn.weight"]) + assertSame(qB, ids["model.layers.blk.0.attn.bias"]) + assertSame(mlpW, ids["model.layers.blk.0.mlp.weight"]) + assertEquals(TensorId(listOf("model", "layers", "blk.0", "attn"), "weight"), qW.id) + assertEquals("model/layers/blk.0/attn", qW.id!!.legacyPath()) + + // same segments as bindPaths + bindPaths(model) + model.walkDepthFirst { node -> node.params.forEach { p -> assertEquals(node.path, p.value.id!!.legacyPath()) } } + } + + @Test + fun idempotentAndRootOverridable() { + val w = tensor() + val model = Block("net", listOf(Leaf("fc", w))) + val a = model.assignTensorIds(); val b = model.assignTensorIds() + assertEquals(a, b) + assertEquals("net.fc.weight", w.id!!.canonical) + model.assignTensorIds(root = "") + assertEquals("fc.weight", w.id!!.canonical) + assertEquals(TensorId(listOf("fc"), "weight"), model.assignTensorIds("").tensors.keys.single()) + } + + @Test + fun tensorsThatCannotCarryAnIdAreReported() { + val bare = bareTensor() + val model = Block("m", listOf(Leaf("fc", bare))) + val ids = model.assignTensorIds() + assertEquals(listOf(TensorId(listOf("m", "fc"), "weight")), ids.notCarried) + assertSame(bare, ids["m.fc.weight"]) + assertNull(bare.id) + } + + @Test + fun idsSurviveRebindingToAnotherContext() { + val w = tensor() + Block("m", listOf(Leaf("fc", w))).assignTensorIds() + val bound = OpsBoundTensor(w, VoidTensorOps()) + assertEquals(w.id, bound.id) + // setting through the bound wrapper writes through to the origin + bound.id = TensorId(listOf("m2"), "weight") + assertEquals("m2.weight", w.id!!.canonical) + // a bound wrapper over a bare tensor keeps the id locally + val wrapped = OpsBoundTensor(bareTensor(), VoidTensorOps()) + wrapped.id = TensorId(listOf("x"), "w") + assertEquals("x.w", wrapped.id!!.canonical) + } + + @Test + fun parameterShortNames() { + val node = Block("attn", emptyList()) + assertEquals("weight", parameterShortName(node, "attn.weight")) + assertEquals("q_proj.weight", parameterShortName(node, "attn.q_proj.weight")) + assertEquals("weight_ih", parameterShortName(node, "weight_ih")) + assertEquals("attn.", parameterShortName(node, "attn.")) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/TensorIdTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/TensorIdTest.kt new file mode 100644 index 000000000..a1778908f --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/TensorIdTest.kt @@ -0,0 +1,67 @@ +package sk.ainet.lang.tensor + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertNull + +class TensorIdTest { + + @Test + fun canonicalAndLegacyForms() { + val id = TensorId(listOf("model", "layers", "blk.3", "attn"), "q_proj.weight") + assertEquals("model.layers.blk.3.attn.q_proj.weight", id.canonical) + assertEquals(id.canonical, id.toString()) + assertEquals("model/layers/blk.3/attn", id.legacyPath()) + assertEquals("model|layers|blk.3|attn", id.legacyPath("|")) + assertNull(id.discriminator) + } + + @Test + fun discriminatorAndViews() { + val scores = TensorId(listOf("model", "layers[3]", "attn"), "scores", "step=17") + assertEquals("model.layers[3].attn.scores#step=17", scores.canonical) + assertEquals("model.layers[3].attn.scores", scores.withDiscriminator(null).canonical) + assertEquals("kv.layers[3].k[1024..2048]", TensorId(listOf("kv", "layers[3]"), "k").view("1024..2048").canonical) + } + + @Test + fun parseInvertsCanonical() { + val id = TensorId.parse("model.layers[3].attn.q_proj.weight#step=17") + assertEquals(listOf("model", "layers[3]", "attn", "q_proj"), id.modulePath) + assertEquals("weight", id.parameter) + assertEquals("step=17", id.discriminator) + assertEquals("model.layers[3].attn.q_proj.weight#step=17", id.canonical) + // round trip as a string for every well-formed id + for (s in listOf("weight", "a.b", "MLP.blk.0.attn.weight", "x.y#d")) assertEquals(s, TensorId.parse(s).canonical) + assertEquals(TensorId(emptyList(), "weight"), TensorId.parse("weight")) + } + + @Test + fun equalityIsByCanonicalString() { + val built = TensorId(listOf("MLP", "blk.0", "attn"), "weight") + val parsed = TensorId.parse("MLP.blk.0.attn.weight") + assertEquals(built, parsed) + assertEquals(built.hashCode(), parsed.hashCode()) + assertEquals(listOf("MLP", "blk", "0", "attn"), parsed.modulePath) // structure differs, identity does not + assertEquals("MLP/blk.0/attn", built.legacyPath()) + assertEquals(setOf(built), setOf(built, parsed)) + } + + @Test + fun fromLegacyPath() { + assertEquals(TensorId(listOf("MLP", "blk.0", "attn"), "weight"), TensorId.fromLegacyPath("MLP/blk.0/attn", "weight")) + assertEquals(TensorId(emptyList(), "bias"), TensorId.fromLegacyPath(null, "bias")) + assertEquals(TensorId(emptyList(), "bias"), TensorId.fromLegacyPath("", "bias")) + assertEquals("MLP/blk.0/attn", TensorId.fromLegacyPath("MLP/blk.0/attn", "weight").legacyPath()) + } + + @Test + fun validation() { + assertFailsWith { TensorId(listOf("a"), "") } + assertFailsWith { TensorId(listOf("a", ""), "w") } + assertFailsWith { TensorId.parse("") } + assertFailsWith { TensorId.parse("a..b") } + assertFailsWith { TensorId.parse(".w") } + } +}