Carriage 1: TensorRef carries identity and encoding to the graph

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt

@@ -5,6 +5,8 @@ import sk.ainet.lang.tensor.Tensor
 import sk.ainet.lang.tensor.withRequiresGrad
 import sk.ainet.lang.tensor.ops.Operation
 import sk.ainet.lang.tensor.ops.TensorSpec
+import sk.ainet.lang.tensor.ops.withTensorEncoding
+import sk.ainet.lang.tensor.ops.withTensorId
 import sk.ainet.lang.types.DType
 import sk.ainet.lang.trace.OpTrace
 import sk.ainet.lang.trace.TraceToGraphBuilder
@@ -80,15 +82,15 @@ public open class DefaultExecutionTape(
                     name = ref.id,
                     shape = inputShapes?.getOrNull(i) ?: ref.shape.dimensions.toList(),
                     dtype = inputDTypes?.getOrNull(i) ?: ref.dtype.name,
-                )
+                ).withTensorEncoding(ref.encoding).withTensorId(ref.tensorId)
             }
             val outputs = List(trace.outputs.size) { i ->
                 val ref = trace.outputs[i]
                 TensorSpec(
                     name = ref.id,
                     shape = outputShapes?.getOrNull(i) ?: ref.shape.dimensions.toList(),
                     dtype = outputDTypes?.getOrNull(i) ?: ref.dtype.name,
-                )
+                ).withTensorEncoding(ref.encoding).withTensorId(ref.tensorId)
             }
 
             val op = object : sk.ainet.lang.tensor.ops.Operation {
@@ -126,7 +128,7 @@ public open class DefaultExecutionTape(
                 shape = tensor.shape.dimensions.toList(),
                 dtype = tensor.dtype.toString(),
                 requiresGrad = tensor.requiresGrad
-            )
+            ).withTensorEncoding(ref.encoding).withTensorId(ref.tensorId)
         }
 
         val outputSpecs = outputs.map { tensor ->
@@ -136,7 +138,7 @@ public open class DefaultExecutionTape(
                 shape = tensor.shape.dimensions.toList(),
                 dtype = tensor.dtype.toString(),
                 requiresGrad = tensor.requiresGrad
-            )
+            ).withTensorEncoding(ref.encoding).withTensorId(ref.tensorId)
         }
 
         val recordedOp = RecordedOperation(

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt

@@ -7,7 +7,9 @@ import sk.ainet.lang.tensor.ops.Operation
 import sk.ainet.lang.tensor.ops.TensorSpec
 import sk.ainet.lang.tensor.ops.ValidationResult
 import sk.ainet.lang.tensor.ops.inferTensorEncoding
+import sk.ainet.lang.tensor.ops.tensorId
 import sk.ainet.lang.tensor.ops.withTensorEncoding
+import sk.ainet.lang.tensor.ops.withTensorId
 
 /**
  * Shared builder to convert OpTrace streams into a ComputeGraph.
@@ -291,6 +293,7 @@ public class TraceToGraphBuilder(
                     shape = weightShape,
                     dtype = weightDtype
                 ).withTensorEncoding(encoding)
+                    .withTensorId(refs.firstNotNullOfOrNull { it.spec.tensorId })
                 syntheticNode = GraphNode(
                     id = nodeId,
                     operation = op,
@@ -346,6 +349,8 @@ public class TraceToGraphBuilder(
             val shape = shapes?.getOrNull(i) ?: effectiveInputs[i].shape.dimensions.toList()
             val dtype = dtypes?.getOrNull(i) ?: effectiveInputs[i].dtype::class.simpleName ?: "unknown"
             TensorSpec(name = name, shape = shape, dtype = dtype)
+                .withTensorEncoding(effectiveInputs[i].encoding)
+                .withTensorId(effectiveInputs[i].tensorId)
         }
     }
 
@@ -358,6 +363,8 @@ public class TraceToGraphBuilder(
             val shape = shapes?.getOrNull(i) ?: trace.outputs[i].shape.dimensions.toList()
             val dtype = dtypes?.getOrNull(i) ?: trace.outputs[i].dtype::class.simpleName ?: "unknown"
             TensorSpec(name = name, shape = shape, dtype = dtype)
+                .withTensorEncoding(trace.outputs[i].encoding)
+                .withTensorId(trace.outputs[i].tensorId)
         }
     }
 

skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/TraceIdentityToGraphTest.kt

@@ -0,0 +1,67 @@
+package sk.ainet.compile.graph
+
+import sk.ainet.context.Phase
+import sk.ainet.exec.tensor.ops.DefaultCpuOps
+import sk.ainet.lang.graph.DefaultExecutionTape
+import sk.ainet.lang.graph.DefaultGradientTape
+import sk.ainet.lang.graph.DefaultGraphExecutionContext
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.TensorId
+import sk.ainet.lang.tensor.data.DenseTensorDataFactory
+import sk.ainet.lang.tensor.data.Q8_0BlockTensorData
+import sk.ainet.lang.tensor.ops.AddOperation
+import sk.ainet.lang.tensor.ops.tensorEncoding
+import sk.ainet.lang.tensor.ops.tensorId
+import sk.ainet.lang.tensor.storage.TensorEncoding
+import sk.ainet.lang.types.FP32
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+/**
+ * #1178: what the tape captured survives to the graph. The trace→graph handoff used to drop all
+ * metadata — a packed weight arrived downstream as a name string, a parameter as `t7`. Now the
+ * `TensorSpec`s carry the encoding *object* and the module-path identity through the same
+ * untyped-metadata mechanism `tensorEncoding` proved.
+ */
+class TraceIdentityToGraphTest {
+
+    private fun ctx(): DefaultGraphExecutionContext {
+        val dataFactory = DenseTensorDataFactory()
+        return DefaultGraphExecutionContext(
+            baseOps = DefaultCpuOps(dataFactory),
+            phase = Phase.TRAIN,
+            tensorDataFactory = dataFactory,
+            createTapeFactory = { _ -> DefaultGradientTape(true) },
+        )
+    }
+
+    @Test
+    fun identityAndEncodingSurviveToTheGraphSpecs() {
+        val trainCtx = ctx()
+        val id = TensorId(listOf("model", "blk[0]"), "ffn.weight")
+
+        val packed = Q8_0BlockTensorData.fromRawBytes(Shape(32), ByteArray(34))
+        @Suppress("UNCHECKED_CAST")
+        val weight = trainCtx.fromData(packed as sk.ainet.lang.tensor.data.TensorData<FP32, Float>, FP32::class)
+        trainCtx.session.identify(weight, id)
+
+        val dense = trainCtx.fromFloatArray<FP32, Float>(Shape(32), FP32::class, FloatArray(32) { it * 0.5f })
+
+        val tape = DefaultExecutionTape(trainCtx.session)
+        tape.startRecording()
+        tape.recordOperation(AddOperation<FP32, Float>(), listOf(weight, dense), listOf(dense))
+        tape.stopRecording()
+
+        val graph = tape.toComputeGraph(synthesizeExternalInputs = true)
+        val allSpecs = graph.nodes.flatMap { it.inputs + it.outputs } + graph.edges.map { it.tensorSpec }
+
+        val identified = allSpecs.filter { it.tensorId == id }
+        assertTrue(identified.isNotEmpty(), "the weight's TensorId must reach the graph, got specs: ${allSpecs.map { it.name }}")
+        assertEquals(
+            TensorEncoding.Q8_0,
+            identified.first { it.tensorEncoding != null }.tensorEncoding,
+            "the encoding object — block size intact — must ride along",
+        )
+    }
+}

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt

@@ -62,6 +62,14 @@ public object HloGenerator {
         ctx: DefaultGraphExecutionContext
     ) {
         val module = model.create(ctx)
+        // Register each parameter's module-path identity on the trace session before recording
+        // (#1178): the tensor itself does not know which parameter it is, and refs are immutable
+        // once created, so this must happen before the forward pass touches them. Unparsable
+        // names stay unidentified rather than guessed.
+        for (parameter in module.trainableParameters()) {
+            val parsed = runCatching { sk.ainet.lang.tensor.TensorId.parse(parameter.name) }.getOrNull()
+            if (parsed != null) ctx.session.identify(parameter.value, parsed)
+        }
         model.calculate(
             module = module,
             inputValue = sampleInput,

skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api

@@ -6563,6 +6563,12 @@ public final class sk/ainet/lang/tensor/ops/TensorSpec {
 	public fun toString ()Ljava/lang/String;
 }
 
+public final class sk/ainet/lang/tensor/ops/TensorSpecIdentities {
+	public static final field TENSOR_ID_METADATA_KEY Ljava/lang/String;
+	public static final fun getTensorId (Lsk/ainet/lang/tensor/ops/TensorSpec;)Lsk/ainet/lang/tensor/TensorId;
+	public static final fun withTensorId (Lsk/ainet/lang/tensor/ops/TensorSpec;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/tensor/ops/TensorSpec;
+}
+
 public final class sk/ainet/lang/tensor/ops/TensorSpecs {
 	public static final field TENSOR_ENCODING_METADATA_KEY Ljava/lang/String;
 	public static final fun getTensorEncoding (Lsk/ainet/lang/tensor/ops/TensorSpec;)Lsk/ainet/lang/tensor/storage/TensorEncoding;
@@ -7796,23 +7802,29 @@ public final class sk/ainet/lang/trace/OpTrace {
 }
 
 public final class sk/ainet/lang/trace/TensorRef {
-	public fun <init> (Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;)V
+	public fun <init> (Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/tensor/storage/TensorEncoding;)V
+	public synthetic fun <init> (Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/tensor/storage/TensorEncoding;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun component1 ()Ljava/lang/String;
 	public final fun component2 ()Lsk/ainet/lang/tensor/Shape;
 	public final fun component3 ()Lsk/ainet/lang/types/DType;
-	public final fun copy (Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/trace/TensorRef;
-	public static synthetic fun copy$default (Lsk/ainet/lang/trace/TensorRef;Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;ILjava/lang/Object;)Lsk/ainet/lang/trace/TensorRef;
+	public final fun component4 ()Lsk/ainet/lang/tensor/TensorId;
+	public final fun component5 ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
+	public final fun copy (Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/tensor/storage/TensorEncoding;)Lsk/ainet/lang/trace/TensorRef;
+	public static synthetic fun copy$default (Lsk/ainet/lang/trace/TensorRef;Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/tensor/storage/TensorEncoding;ILjava/lang/Object;)Lsk/ainet/lang/trace/TensorRef;
 	public fun equals (Ljava/lang/Object;)Z
 	public final fun getDtype ()Lsk/ainet/lang/types/DType;
+	public final fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
 	public final fun getId ()Ljava/lang/String;
 	public final fun getShape ()Lsk/ainet/lang/tensor/Shape;
+	public final fun getTensorId ()Lsk/ainet/lang/tensor/TensorId;
 	public fun hashCode ()I
 	public fun toString ()Ljava/lang/String;
 }
 
 public class sk/ainet/lang/trace/TraceSession {
 	public fun <init> ()V
 	public final fun clear ()V
+	public fun identify (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/TensorId;)V
 	public fun refOf (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/trace/TensorRef;
 	public final fun refsOf (Ljava/util/List;)Ljava/util/List;
 	public fun resolve (Ljava/lang/String;)Lsk/ainet/lang/tensor/Tensor;

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorSpecIdentity.kt

@@ -0,0 +1,36 @@
+@file:JvmName("TensorSpecIdentities")
+
+package sk.ainet.lang.tensor.ops
+
+import sk.ainet.lang.tensor.TensorId
+import kotlin.jvm.JvmName
+
+/**
+ * Metadata key used to carry a [TensorId] on a [TensorSpec] (#1178).
+ *
+ * Same contract as [TENSOR_ENCODING_METADATA_KEY]: an untyped metadata entry with typed
+ * accessors, so the compile pipeline carries the fact without importing storage-model types —
+ * the `TensorSpecEncoding` precedent, and the standing rule for everything the tape captures.
+ */
+public const val TENSOR_ID_METADATA_KEY: String = "tensorId"
+
+/**
+ * The module-path identity carried on this spec (`model.layers[3].attn.q_proj.weight`), or
+ * `null` if no producer knew it. `null` means "unidentified", not "anonymous by design" —
+ * consumers keying per-weight policy on identity should skip unidentified tensors.
+ */
+public val TensorSpec.tensorId: TensorId?
+    get() = metadata[TENSOR_ID_METADATA_KEY] as? TensorId
+
+/**
+ * Return a copy of this spec with [id] stored in its metadata map. Passing `null` removes the
+ * entry; a non-null value adds or replaces it, leaving all other metadata untouched.
+ */
+public fun TensorSpec.withTensorId(id: TensorId?): TensorSpec {
+    val newMetadata: Map<String, Any> = if (id == null) {
+        metadata - TENSOR_ID_METADATA_KEY
+    } else {
+        metadata + (TENSOR_ID_METADATA_KEY to id)
+    }
+    return copy(metadata = newMetadata)
+}

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TensorRef.kt

@@ -1,14 +1,24 @@
 package sk.ainet.lang.trace
 
 import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.TensorId
+import sk.ainet.lang.tensor.storage.TensorEncoding
 import sk.ainet.lang.types.DType
 
 /**
  * TensorRef represents a reference to a tensor in the computation graph.
  * It captures the essential metadata needed for graph construction and optimization.
+ *
+ * [tensorId] and [encoding] carry what only the tape can see (#1178): the live tensor's
+ * module-path identity (when someone who knows it called [TraceSession.identify]) and its
+ * physical storage encoding as an *object* — block size intact, unlike the display-name
+ * strings downstream stages used to be left with. Both default to `null`, so every existing
+ * construction and destructuring compiles unchanged.
  */
 public data class TensorRef(
     val id: String,
     val shape: Shape,
-    val dtype: DType
-)
+    val dtype: DType,
+    val tensorId: TensorId? = null,
+    val encoding: TensorEncoding? = null,
+)

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt

@@ -1,6 +1,8 @@
 package sk.ainet.lang.trace
 
 import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.tensor.TensorId
+import sk.ainet.lang.tensor.ops.inferTensorEncoding
 import sk.ainet.lang.types.*
 
 /**
@@ -11,6 +13,19 @@ public open class TraceSession {
     private var nextId = 0
     private val tensorToRef = mutableMapOf<Any, TensorRef>()
     private val refToId = mutableMapOf<String, Tensor<*, *>>()
+    private val identities = mutableMapOf<Any, TensorId>()
+
+    /**
+     * Register [tensor]'s module-path identity so its [TensorRef] carries it (#1178).
+     *
+     * A tensor does not know which parameter it is — the module that owns it does. Whoever
+     * holds that knowledge (e.g. the HLO generator walking `trainableParameters()` before
+     * recording) calls this *before* the tensor's first [refOf]; refs are immutable and cached,
+     * so identities registered later do not retrofit existing refs.
+     */
+    public open fun identify(tensor: Tensor<*, *>, id: TensorId) {
+        identities[unwrap(tensor)] = id
+    }
 
     /**
      * Get or create a TensorRef for the given tensor.
@@ -35,7 +50,9 @@ public open class TraceSession {
             val ref = TensorRef(
                 id = "t${nextId++}",
                 shape = tensor.shape,
-                dtype = dtypeInstance
+                dtype = dtypeInstance,
+                tensorId = identities[key],
+                encoding = tensor.data.inferTensorEncoding(),
             )
             refToId[ref.id] = tensor
             ref

skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/trace/TraceSessionIdentityTest.kt

@@ -0,0 +1,57 @@
+package sk.ainet.lang.trace
+
+import sk.ainet.context.DefaultDataExecutionContext
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.TensorId
+import sk.ainet.lang.tensor.data.Q8_0BlockTensorData
+import sk.ainet.lang.tensor.storage.TensorEncoding
+import sk.ainet.lang.types.FP32
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNull
+
+/**
+ * #1178: the tape captures what only it can see — the live tensor's storage encoding as an
+ * object (block size intact) and, when someone who knows it registered one, its module-path
+ * identity. Refs are immutable and cached, so identity must be registered before first use.
+ */
+class TraceSessionIdentityTest {
+
+    private val ctx = DefaultDataExecutionContext()
+
+    @Test
+    fun refCapturesThePackedEncodingObject() {
+        val packed = Q8_0BlockTensorData.fromRawBytes(Shape(32), ByteArray(34))
+        @Suppress("UNCHECKED_CAST")
+        val tensor = ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData<FP32, Float>, FP32::class)
+        val ref = TraceSession().refOf(tensor)
+        assertEquals(TensorEncoding.Q8_0, ref.encoding, "the encoding object, not a display name")
+    }
+
+    @Test
+    fun denseTensorHasNoEncodingAndNoIdentityByDefault() {
+        val tensor = ctx.fromFloatArray<FP32, Float>(Shape(4), FP32::class, floatArrayOf(1f, 2f, 3f, 4f))
+        val ref = TraceSession().refOf(tensor)
+        assertNull(ref.encoding)
+        assertNull(ref.tensorId)
+    }
+
+    @Test
+    fun identityRegisteredBeforeFirstUseIsCarried() {
+        val session = TraceSession()
+        val weight = ctx.fromFloatArray<FP32, Float>(Shape(2, 2), FP32::class, FloatArray(4))
+        val id = TensorId(listOf("model", "layers[3]", "attn"), "q_proj.weight")
+        session.identify(weight, id)
+        assertEquals(id, session.refOf(weight).tensorId)
+    }
+
+    @Test
+    fun identityRegisteredAfterFirstUseDoesNotRetrofitTheCachedRef() {
+        val session = TraceSession()
+        val weight = ctx.fromFloatArray<FP32, Float>(Shape(2), FP32::class, FloatArray(2))
+        val before = session.refOf(weight)
+        assertNull(before.tensorId)
+        session.identify(weight, TensorId(listOf("model"), "weight"))
+        assertNull(session.refOf(weight).tensorId, "refs are immutable and cached — identify before first refOf")
+    }
+}