Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
66 changes: 64 additions & 2 deletions skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
}

Expand Down Expand Up @@ -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 <init> (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
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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;
Expand Down Expand Up @@ -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 <init> (Ljava/util/List;Ljava/lang/String;Ljava/lang/String;)V
public synthetic fun <init> (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 <init> ()V
public final fun segment (Lkotlin/jvm/functions/Function1;)V
Expand All @@ -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;
Expand Down Expand Up @@ -2909,19 +2967,21 @@ 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 <init> (Lsk/ainet/lang/tensor/data/TensorData;Lkotlin/reflect/KClass;Lsk/ainet/lang/tensor/GradState;)V
public synthetic fun <init> (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
public fun getData ()Lsk/ainet/lang/tensor/data/TensorData;
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
}

Expand Down Expand Up @@ -4123,20 +4183,22 @@ 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 <init> (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/ops/TensorOps;)V
public fun accumulateGrad (Lsk/ainet/lang/tensor/Tensor;)V
public fun getData ()Lsk/ainet/lang/tensor/data/TensorData;
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
}

Expand Down
Original file line number Diff line number Diff line change
@@ -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(", ", "[", "]")
Original file line number Diff line number Diff line change
@@ -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<TensorId, Tensor<*, *>>,
val notCarried: List<TensorId>,
) {
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 `"<module>.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<TensorId, Tensor<*, *>>()
val notCarried = ArrayList<TensorId>()
fun visit(node: ModuleNode, segments: List<String>) {
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)
}

/** `"<moduleName>.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
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,13 @@ public interface Tensor<T : DType, V> {
*/
public val dtype: KClass<T>

/**
* 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.
*/
Expand Down
Original file line number Diff line number Diff line change
@@ -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<String>,
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?
}
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,9 @@ public class VoidOpsTensor<T : DType, V>(
override val data: TensorData<T, V>,
override val dtype: KClass<T>,
override val gradState: GradState<T, V> = GradState()
) : Tensor<T, V> {
) : Tensor<T, V>, TensorIdBearer {
override val ops: TensorOps
get() = VoidTensorOps()

override var id: TensorId? = null
}
Loading
Loading