feat(safetensors): loadTensorStorageMapped on the sharded reader

skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/StreamingShardedSafeTensorsReader.kt

@@ -5,6 +5,7 @@ import kotlinx.coroutines.flow.MutableSharedFlow
 import sk.ainet.io.model.LoadingProgress
 import sk.ainet.io.model.LoadingStage
 import sk.ainet.io.model.ProgressReportingLoader
+import sk.ainet.lang.tensor.storage.TensorStorage
 
 /**
  * Streaming reader for sharded/multi-file SafeTensors models.
@@ -96,6 +97,54 @@ public class StreamingShardedSafeTensorsReader private constructor(
         return reader.loadTensorData(tensor.name)
     }
 
+    /**
+     * Same shape as [loadTensorData] but returns a file-backed
+     * [TensorStorage] instead of a heap [ByteArray]. Lets callers
+     * memory-map the tensor's byte range straight from the shard file
+     * without going through a 2 GB-capped `ByteArray` round-trip.
+     *
+     * The returned [TensorStorage] holds a
+     * [sk.ainet.lang.tensor.storage.BufferHandle.FileBacked] that
+     * references the shard file by absolute path; callers (or the
+     * runtime that consumes the storage) own the mmap lifecycle.
+     *
+     * Sharded analog of
+     * [StreamingSafeTensorsReader.loadTensorStorageMapped]. The
+     * shard's file path is resolved internally from the index — the
+     * caller doesn't need to know which physical file contains the
+     * tensor.
+     *
+     * @param tensor The tensor info from [tensors].
+     * @return [TensorStorage] descriptor with a file-backed buffer
+     *   handle pointing at the shard file's tensor byte range.
+     * @throws IllegalStateException if the containing shard was not
+     *   loaded, or if the per-shard reader does not surface the
+     *   tensor (consistency check).
+     */
+    public fun loadTensorStorageMapped(tensor: ShardedTensorInfo): TensorStorage {
+        val reader = shardReaders[tensor.shardFilename]
+            ?: throw IllegalStateException("Shard not loaded: ${tensor.shardFilename}")
+        val streamingTensor = reader.tensors.firstOrNull { it.name == tensor.name }
+            ?: throw IllegalStateException(
+                "Tensor '${tensor.name}' not found in shard '${tensor.shardFilename}'",
+            )
+        val path = resolveShardPath(tensor.shardFilename)
+        return reader.loadTensorStorageMapped(streamingTensor, path)
+    }
+
+    /**
+     * Convenience overload for [loadTensorStorageMapped] that looks up
+     * the tensor by name. Mirrors the [loadTensorData] name-based
+     * overload.
+     *
+     * @throws IllegalArgumentException if no tensor matches [name].
+     */
+    public fun loadTensorStorageMapped(name: String): TensorStorage {
+        val tensor = _tensors.firstOrNull { it.name == name }
+            ?: throw IllegalArgumentException("Tensor not found: $name")
+        return loadTensorStorageMapped(tensor)
+    }
+
     override fun close() {
         shardReaders.values.forEach { it.close() }
         shardReaders.clear()

skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/StreamingShardedSafeTensorsReaderJvmTest.kt

@@ -0,0 +1,122 @@
+package sk.ainet.io.safetensors
+
+import java.nio.file.Files
+import java.nio.file.Path
+import kotlin.io.path.deleteIfExists
+import kotlin.io.path.exists
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNotNull
+import kotlin.test.assertTrue
+import kotlinx.coroutines.runBlocking
+import kotlinx.io.buffered
+import kotlinx.io.files.Path as IoPath
+import kotlinx.io.files.SystemFileSystem
+import sk.ainet.lang.tensor.storage.BufferHandle
+
+/**
+ * JVM-only end-to-end coverage of [StreamingShardedSafeTensorsReader],
+ * focused on the `loadTensorStorageMapped` entry points added to support
+ * file-backed `TensorStorage` reads on tensors that exceed the JVM
+ * `ByteArray` limit (used by the Gemma 4 PLE token-embedding table).
+ *
+ * Builds a real tiny single-shard SafeTensors file via [SafeTensorsWriter]
+ * and a hand-crafted index JSON, opens via the sharded reader, and
+ * verifies the returned [sk.ainet.lang.tensor.storage.TensorStorage]
+ * carries the right metadata and a [BufferHandle.FileBacked] handle.
+ */
+class StreamingShardedSafeTensorsReaderJvmTest {
+
+    @Test
+    fun `loadTensorStorageMapped returns file-backed storage with correct metadata`() {
+        val tempDir = Files.createTempDirectory("sharded-mapped-")
+        try {
+            val shardName = "model.safetensors"
+            val shardPath = tempDir.resolve(shardName)
+            val indexPath = tempDir.resolve("model.safetensors.index.json")
+
+            // Build a 2x2 F32 tensor and write it as a single-shard SafeTensors file.
+            val data = floatArrayOf(1f, 2f, 3f, 4f)
+            writeSingleTensor(shardPath, name = "test.weight", shape = listOf(2L, 2L), data = data)
+
+            // Hand-craft the index. The sharded reader treats this as a 1-shard
+            // sharded model — equivalent to how a single-file gemma-4-e2b-it
+            // checkpoint gets routed through the sharded loader.
+            val totalSize = Files.size(shardPath)
+            val indexJson = """
+                {
+                  "metadata": {"total_size": $totalSize},
+                  "weight_map": {"test.weight": "$shardName"}
+                }
+            """.trimIndent()
+            Files.writeString(indexPath, indexJson)
+
+            runBlocking {
+                StreamingShardedSafeTensorsReader.openFromIndex(indexPath.toString()).use { reader ->
+                    assertEquals(1, reader.tensors.size)
+                    assertEquals("test.weight", reader.tensors[0].name)
+
+                    // 1) Look-up by name overload.
+                    val byName = reader.loadTensorStorageMapped("test.weight")
+                    assertEquals(listOf(2, 2), byName.shape.dimensions.toList())
+                    assertTrue(byName.isFileBacked)
+                    val handle = byName.buffer
+                    assertTrue(handle is BufferHandle.FileBacked, "Expected FileBacked, got ${handle::class}")
+                    handle as BufferHandle.FileBacked
+                    assertEquals(shardPath.toString(), handle.path)
+                    assertEquals(16L, handle.sizeInBytes) // 4 floats × 4 bytes
+
+                    // 2) Look-up by ShardedTensorInfo overload — must produce
+                    // the same TensorStorage shape.
+                    val byInfo = reader.loadTensorStorageMapped(reader.tensors[0])
+                    assertEquals(byName.shape.dimensions.toList(), byInfo.shape.dimensions.toList())
+                    assertTrue(byInfo.isFileBacked)
+                }
+            }
+        } finally {
+            tempDir.toFile().deleteRecursively()
+        }
+    }
+
+    @Test
+    fun `loadTensorStorageMapped throws on missing tensor name`() {
+        val tempDir = Files.createTempDirectory("sharded-mapped-missing-")
+        try {
+            val shardPath = tempDir.resolve("model.safetensors")
+            val indexPath = tempDir.resolve("model.safetensors.index.json")
+            writeSingleTensor(shardPath, "real.weight", listOf(1L), floatArrayOf(1f))
+            val total = Files.size(shardPath)
+            Files.writeString(
+                indexPath,
+                """{"metadata":{"total_size": $total},"weight_map":{"real.weight":"model.safetensors"}}""",
+            )
+
+            runBlocking {
+                StreamingShardedSafeTensorsReader.openFromIndex(indexPath.toString()).use { reader ->
+                    val thrown = runCatching { reader.loadTensorStorageMapped("missing.weight") }
+                        .exceptionOrNull()
+                    assertNotNull(thrown, "Expected an exception for missing tensor name")
+                    assertTrue(thrown is IllegalArgumentException)
+                }
+            }
+        } finally {
+            tempDir.toFile().deleteRecursively()
+        }
+    }
+
+    private fun writeSingleTensor(
+        outputPath: Path,
+        name: String,
+        shape: List<Long>,
+        data: FloatArray,
+    ) {
+        outputPath.deleteIfExists()
+        val ioPath = IoPath(outputPath.toString())
+        SystemFileSystem.sink(ioPath).buffered().use { sink ->
+            SafeTensorsWriter.write(sink) {
+                tensorF32(name, shape, data)
+            }
+        }
+        check(outputPath.exists()) { "Failed to write fixture at $outputPath" }
+    }
+}