Add SafeTensorsParametersLoader.withPolicy(DTypePolicy) (W0b of #615)

Companion factory + adapter that lets callers drive the loader by
the generalised DTypePolicy from W1 instead of the BF16-specific
Bf16LoadPolicy enum. Existing callers keep working unchanged — the
factory is additive.

Policy → behaviour mapping covers all four DTypePolicy arms and
follows the RFC semantics:
- Any         -> DEQUANT_TO_FP32 (adaptive default)
- Require(BF16) -> KEEP_NATIVE
- Require(FP32) -> DEQUANT_TO_FP32
- Require(FP16) -> throws (no Fp16DenseTensorData backing yet)
- Require(other) -> throws (loader can't fabricate dtypes)
- Prefer(BF16) -> KEEP_NATIVE; Prefer(other) -> DEQUANT (soft fall-through)
- OneOf containing BF16 -> KEEP_NATIVE; otherwise DEQUANT

The Require error messages name the failing target explicitly,
suggest the working alternative (Require(FP32) / Any), and point
at the missing backing class for the FP16 case — fail-fast with
diagnostic detail, per the RFC.

Round-trip test guards the parity between Bf16LoadPolicy.toDTypePolicy()
(W2) and this mapper: every enum arm survives the round trip
through DTypePolicy intact.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: michalharakal

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

@@ -8,7 +8,10 @@ import sk.ainet.lang.tensor.Shape
 import sk.ainet.lang.tensor.Tensor
 import sk.ainet.lang.tensor.data.Bf16DenseTensorData
 import sk.ainet.lang.tensor.data.TensorData
+import sk.ainet.lang.types.BF16
 import sk.ainet.lang.types.DType
+import sk.ainet.lang.types.DTypePolicy
+import sk.ainet.lang.types.FP16
 import sk.ainet.lang.types.FP32
 import sk.ainet.lang.types.Int32
 import sk.ainet.lang.types.Int8
@@ -286,4 +289,64 @@ class SafeTensorsParametersLoader(
             }
         }
     }
+
+    companion object {
+
+        /**
+         * Constructs a SafeTensorsParametersLoader from a generalised
+         * [DTypePolicy] instead of the BF16-specific [Bf16LoadPolicy].
+         * Bridge for the policy-driven loader path described in the
+         * dtype-policy RFC (#615).
+         *
+         * Policy → behaviour mapping (BF16 source tensors only —
+         * other dtypes are handled per the per-arm `require` checks
+         * in [load]):
+         * - [DTypePolicy.Any]: BF16 dequants to FP32 (the existing
+         *   adaptive default).
+         * - [DTypePolicy.Require] target = `BF16`: KEEP_NATIVE.
+         * - [DTypePolicy.Require] target = `FP32`: DEQUANT_TO_FP32.
+         * - [DTypePolicy.Require] target = `FP16`: throws — F16
+         *   KEEP_NATIVE is a follow-up (no `Fp16DenseTensorData`
+         *   yet); use `Require(FP32)` if you want F16 dequanted, or
+         *   `Any` to inherit the adaptive default.
+         * - [DTypePolicy.Require] target = anything else: throws —
+         *   SafeTensors can't fabricate dtypes the file doesn't carry.
+         * - [DTypePolicy.Prefer] target = `BF16`: KEEP_NATIVE.
+         * - [DTypePolicy.Prefer] target = anything else: DEQUANT_TO_FP32
+         *   (the soft path falls through).
+         * - [DTypePolicy.OneOf] containing `BF16`: KEEP_NATIVE.
+         * - [DTypePolicy.OneOf] without `BF16`: DEQUANT_TO_FP32.
+         */
+        fun withPolicy(
+            sourceProvider: () -> RandomAccessSource,
+            policy: DTypePolicy,
+            onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> },
+        ): SafeTensorsParametersLoader = SafeTensorsParametersLoader(
+            sourceProvider = sourceProvider,
+            onProgress = onProgress,
+            bf16Policy = mapPolicyToBf16(policy),
+        )
+
+        internal fun mapPolicyToBf16(policy: DTypePolicy): Bf16LoadPolicy = when (policy) {
+            DTypePolicy.Any -> Bf16LoadPolicy.DEQUANT_TO_FP32
+            is DTypePolicy.Require -> when (policy.target) {
+                BF16 -> Bf16LoadPolicy.KEEP_NATIVE
+                FP32 -> Bf16LoadPolicy.DEQUANT_TO_FP32
+                FP16 -> throw IllegalArgumentException(
+                    "SafeTensorsParametersLoader: Require(FP16) is not supported — " +
+                        "F16 KEEP_NATIVE has no Fp16DenseTensorData backing yet. " +
+                        "Use Require(FP32) to dequant F16 sources, or Any to inherit the adaptive default.",
+                )
+                else -> throw IllegalArgumentException(
+                    "SafeTensorsParametersLoader: Require(${policy.target.name}) is not satisfiable — " +
+                        "the loader produces FP32 / BF16 / Int32 / Int8 tensors depending on source dtype; " +
+                        "it cannot fabricate ${policy.target.name} from arbitrary sources.",
+                )
+            }
+            is DTypePolicy.Prefer -> if (policy.target == BF16) Bf16LoadPolicy.KEEP_NATIVE
+                                    else Bf16LoadPolicy.DEQUANT_TO_FP32
+            is DTypePolicy.OneOf -> if (BF16 in policy.allowed) Bf16LoadPolicy.KEEP_NATIVE
+                                    else Bf16LoadPolicy.DEQUANT_TO_FP32
+        }
+    }
 }

skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderPolicyTest.kt

@@ -0,0 +1,114 @@
+package sk.ainet.io.safetensors
+
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import sk.ainet.lang.types.BF16
+import sk.ainet.lang.types.DTypePolicy
+import sk.ainet.lang.types.FP16
+import sk.ainet.lang.types.FP32
+import sk.ainet.lang.types.Int8
+
+/**
+ * Unit tests for the `DTypePolicy` → `Bf16LoadPolicy` adapter in
+ * [SafeTensorsParametersLoader.mapPolicyToBf16]. The `withPolicy`
+ * factory is a thin wrapper over this mapper plus the existing
+ * constructor; testing the mapper covers the routing logic without
+ * needing a real SafeTensors fixture.
+ */
+class SafeTensorsParametersLoaderPolicyTest {
+
+    @Test
+    fun any_maps_to_dequant_to_fp32() {
+        assertEquals(
+            Bf16LoadPolicy.DEQUANT_TO_FP32,
+            SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Any),
+        )
+    }
+
+    @Test
+    fun require_bf16_maps_to_keep_native() {
+        assertEquals(
+            Bf16LoadPolicy.KEEP_NATIVE,
+            SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Require(BF16)),
+        )
+    }
+
+    @Test
+    fun require_fp32_maps_to_dequant() {
+        assertEquals(
+            Bf16LoadPolicy.DEQUANT_TO_FP32,
+            SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Require(FP32)),
+        )
+    }
+
+    @Test
+    fun require_fp16_fails_with_explicit_message() {
+        val ex = assertFailsWith<IllegalArgumentException> {
+            SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Require(FP16))
+        }
+        // The error message must point the operator at the alternative —
+        // RFC says "fail-fast with clear diagnostics," not just throw.
+        val msg = ex.message ?: ""
+        assertEquals(true, msg.contains("Require(FP16)"), "msg: $msg")
+        assertEquals(true, msg.contains("Fp16DenseTensorData"), "msg: $msg")
+    }
+
+    @Test
+    fun require_unsupported_target_fails_with_explicit_message() {
+        val ex = assertFailsWith<IllegalArgumentException> {
+            SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Require(Int8))
+        }
+        val msg = ex.message ?: ""
+        assertEquals(true, msg.contains("Require(Int8)"), "msg: $msg")
+        assertEquals(true, msg.contains("cannot fabricate"), "msg: $msg")
+    }
+
+    @Test
+    fun prefer_bf16_maps_to_keep_native() {
+        assertEquals(
+            Bf16LoadPolicy.KEEP_NATIVE,
+            SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Prefer(BF16)),
+        )
+    }
+
+    @Test
+    fun prefer_fp32_or_anything_else_maps_to_dequant() {
+        assertEquals(
+            Bf16LoadPolicy.DEQUANT_TO_FP32,
+            SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Prefer(FP32)),
+        )
+        assertEquals(
+            Bf16LoadPolicy.DEQUANT_TO_FP32,
+            SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Prefer(FP16)),
+            "Prefer is soft — unsatisfiable preferences fall through silently, no throw",
+        )
+    }
+
+    @Test
+    fun oneOf_with_bf16_maps_to_keep_native() {
+        assertEquals(
+            Bf16LoadPolicy.KEEP_NATIVE,
+            SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.OneOf(setOf(BF16, FP32))),
+        )
+    }
+
+    @Test
+    fun oneOf_without_bf16_maps_to_dequant() {
+        assertEquals(
+            Bf16LoadPolicy.DEQUANT_TO_FP32,
+            SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.OneOf(setOf(FP32, FP16))),
+        )
+    }
+
+    @Test
+    fun parity_with_bf16LoadPolicy_toDTypePolicy() {
+        // Round-trip property: the BF16 enum's adapter should land on a
+        // policy that the inverse mapper sends back to the original enum.
+        for (arm in Bf16LoadPolicy.entries) {
+            val asDTypePolicy = arm.toDTypePolicy()
+            val back = SafeTensorsParametersLoader.mapPolicyToBf16(asDTypePolicy)
+            assertEquals(arm, back, "round-trip failed for $arm via $asDTypePolicy")
+        }
+    }
+}