nx_shapley.nx source
↩ module page · 474 lines · 16341 B
1// nx_shapley.nx -- game-theoretic feature attribution.
2//
3// Shapley 1953 ("A Value for n-Person Games"): the UNIQUE fair
4// allocation of a coalition's value back to its members. The
5// substrate's third attribution primitive after nx_attribution
6// (top-K of W @ x) and nx_patch_attribution (Cohen's d per axis).
7//
8// phi_i = sum over S subset N\{i} of
9// |S|! * (N-|S|-1)! / N! * [ f(S union {i}) - f(S) ]
10//
11// Closes the WORLD-CLASS attribution gap named in
12// nxc2/docs/TOPDOWN_BOTTOMUP_GAP_ANALYSIS.md. No existing
13// image-gen / LLM stack ships Shapley in typed bits-up form;
14// they all defer to SHAP / KernelSHAP Python libs that hide
15// per-coalition evaluation under floating-point opacity.
16//
17// Two modes:
18// EXACT -- N <= 10 features; full 2^N coalition enumeration.
19// SAMPLED -- N > 10; Castro/Gomez/Tejada 2009 permutation
20// sampling estimator. Unbiased.
21//
22// Caller supplies f(active_mask, n_features, ctx) returning a
23// scalar value in Q10. Substrate computes the contribution
24// vector phi[] (Q10 each) such that sum(phi) + baseline = f(all_on).
25//
26// genealogy_id: shapley_1953_castro_2009
27// lineage_id: substrate_attribution_v3
28//
29// 5W+H+GLP linkage: each phi[i] is the WHY answer for "feature i
30// contributed how much to output f(all_on)". Caller emits one
31// Nx5whRecord with what=ATTRIBUTION per non-trivial phi[i].
32
33// nx_safety_envelope:
34// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
35// sil_target: SIL1
36// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
37// verdict: NOT_YET_EVALUATED
38
39import "nx_syscalls.nx"
40import "nx_runtime.nx"
41import "nx_tier.nx"
42import "nx_random.nx"
43const NX_MAGIC_5040: i64 = 5040
44const NX_MAGIC_40320: i64 = 40320
45const NX_MAGIC_362880: i64 = 362880
46const NX_MAGIC_3628800: i64 = 3628800
47const NX_MAGIC_4096: i64 = 4096
48
49// ===== mode + cap constants =======================================
50
51const NX_SHAPLEY_MODE_EXACT: nx_int = 0
52const NX_SHAPLEY_MODE_SAMPLED: nx_int = 1
53
54// 2^10 = 1024 coalitions; safe for exact enumeration.
55const NX_SHAPLEY_MAX_EXACT_N: nx_int = 10
56
57// JPL Rule 2: bounded loops. Max features the result can carry.
58const NX_SHAPLEY_MAX_FEATURES: nx_int = 256
59
60// Default sample count; caller may override.
61const NX_SHAPLEY_DEFAULT_SAMPLES: nx_int = 2048
62const NX_SHAPLEY_MAX_SAMPLES: nx_int = 65536
63
64// Q10 unit so phi values stay in fixed-point.
65const NX_SHAPLEY_Q10: nx_int = 1024
66
67// ===== result struct ==============================================
68
69struct NxShapleyResult {
70 n_features: nx_int,
71 mode: nx_int,
72 n_evaluations: nx_int,
73 baseline_q10: nx_int,
74 full_q10: nx_int,
75 phi: *nx_int,
76}
77
78const NX_SHAPLEY_RESULT_BYTES: nx_size = 48
79
80// ===== factorial table (precomputed up to 10!) ====================
81//
82// Used by the exact mode for the Shapley weight |S|!*(N-|S|-1)!/N!.
83// Closed-form table avoids overflow at N=10 (10! = 3628800 fits i64
84// trivially; intermediate products in weight stay small).
85
86func _shapley_factorial(n: nx_int) -> nx_int {
87 if n < 0 { return 0 }
88 if n == 0 { return 1 }
89 if n == 1 { return 1 }
90 if n == 2 { return 2 }
91 if n == 3 { return 6 }
92 if n == 4 { return 24 }
93 if n == 5 { return 120 }
94 if n == 6 { return 720 }
95 if n == 7 { return NX_MAGIC_5040 }
96 if n == 8 { return NX_MAGIC_40320 }
97 if n == 9 { return NX_MAGIC_362880 }
98 if n == 10 { return NX_MAGIC_3628800 }
99 // n > 10 not used by exact mode; sampled mode never calls this.
100 return 0
101}
102
103// Returns Shapley weight for coalition size s, total players n,
104// scaled by NX_SHAPLEY_Q10 so we stay in integer arithmetic.
105// weight_q10 = round( Q10 * s! * (n-s-1)! / n! )
106func _shapley_weight_q10(s: nx_int, n: nx_int) -> nx_int {
107 let num1: nx_int = _shapley_factorial(s)
108 let num2: nx_int = _shapley_factorial(n - s - 1)
109 let den: nx_int = _shapley_factorial(n)
110 if den == 0 { return 0 }
111 let num_q10: nx_int = NX_SHAPLEY_Q10 * num1
112 let prod: nx_int = num_q10 * num2
113 return prod / den
114}
115
116// ===== popcount + bit test helpers ================================
117//
118// Coalition encoding: integer 0..2^N-1, bit i = "feature i in S".
119
120func _shapley_popcount(mask: nx_int) -> nx_int {
121 var count: nx_int = 0
122 var x: nx_int = mask
123 var safety: nx_int = 0
124 while safety < 64 {
125 if x == 0 { return count }
126 let lsb: nx_int = x & 1
127 if lsb == 1 { count = count + 1 }
128 x = x >> 1
129 safety = safety + 1
130 }
131 return count
132}
133
134func _shapley_bit_set(mask: nx_int, bit: nx_int) -> nx_int {
135 if bit < 0 { return mask }
136 if bit >= 64 { return mask }
137 let one_shifted: nx_int = 1 << bit
138 let result: nx_int = mask | one_shifted
139 return result
140}
141
142func _shapley_bit_test(mask: nx_int, bit: nx_int) -> nx_int {
143 if bit < 0 { return 0 }
144 if bit >= 64 { return 0 }
145 let shifted: nx_int = mask >> bit
146 let lsb: nx_int = shifted & 1
147 return lsb
148}
149
150// ===== exact Shapley over 2^N coalitions ==========================
151//
152// fn_ptr signature: takes the active-mask + n_features + caller-ctx
153// and returns a Q10 scalar value. Substrate sums marginal
154// contributions weighted by the Shapley kernel.
155
156func nx_shapley_compute_exact(
157 fn_ptr: func(nx_int, nx_int, *u8) -> nx_int,
158 n_features: nx_int,
159 ctx: *u8,
160 baseline_q10: nx_int) -> *NxShapleyResult {
161
162 if n_features <= 0 { return 0 as *NxShapleyResult }
163 if n_features > NX_SHAPLEY_MAX_EXACT_N { return 0 as *NxShapleyResult }
164
165 let result_ptr: *u8 = sys_mmap(NX_SHAPLEY_RESULT_BYTES)
166 let result: *NxShapleyResult = result_ptr as *NxShapleyResult
167 let phi_bytes: nx_size = (n_features as nx_size) * 8
168 let phi: *nx_int = (sys_mmap(phi_bytes)) as *nx_int
169
170 var init_i: nx_int = 0
171 while init_i < n_features {
172 phi[init_i] = 0
173 init_i = init_i + 1
174 }
175
176 let one_shifted: nx_int = 1 << n_features
177 let total_coalitions: nx_int = one_shifted
178 var subset: nx_int = 0
179 var n_calls: nx_int = 0
180
181 // For each feature i, sum its marginal contribution across all
182 // coalitions that DON'T include i.
183 while subset < total_coalitions {
184 let s_size: nx_int = _shapley_popcount(subset)
185 if s_size < n_features {
186 // f(S) and f(S union {i}) only computed when needed below.
187 // To avoid recomputing f(S) for each i, cache it once here.
188 let f_s: nx_int = fn_ptr(subset, n_features, ctx)
189 n_calls = n_calls + 1
190 var i: nx_int = 0
191 while i < n_features {
192 if _shapley_bit_test(subset, i) == 0 {
193 let s_with_i: nx_int = _shapley_bit_set(subset, i)
194 let f_si: nx_int = fn_ptr(s_with_i, n_features, ctx)
195 n_calls = n_calls + 1
196 let marginal: nx_int = f_si - f_s
197 let w_q10: nx_int = _shapley_weight_q10(s_size, n_features)
198 let weighted_q10: nx_int = w_q10 * marginal
199 let weighted: nx_int = weighted_q10 / NX_SHAPLEY_Q10
200 phi[i] = phi[i] + weighted
201 }
202 i = i + 1
203 }
204 }
205 subset = subset + 1
206 }
207
208 // Compute f(all_on) for the sum-check field.
209 let all_mask: nx_int = total_coalitions - 1
210 let f_all: nx_int = fn_ptr(all_mask, n_features, ctx)
211 n_calls = n_calls + 1
212
213 result.n_features = n_features
214 result.mode = NX_SHAPLEY_MODE_EXACT
215 result.n_evaluations = n_calls
216 result.baseline_q10 = baseline_q10
217 result.full_q10 = f_all
218 result.phi = phi
219 return result
220}
221
222// ===== sampled Shapley (permutation estimator) ====================
223//
224// Castro/Gomez/Tejada 2009 unbiased estimator: for each sample,
225// pick a random permutation pi of features; for each feature i in
226// position k of pi, marginal contribution =
227// f({pi[0]..pi[k]}) - f({pi[0]..pi[k-1]})
228// Average across M samples gives unbiased phi_i.
229//
230// 2*M*N function evaluations; tractable for N up to 256.
231
232func _shapley_permute(perm: *nx_int, n: nx_int, rng: *NxRng) -> nx_int {
233 var k: nx_int = 0
234 while k < n {
235 perm[k] = k
236 k = k + 1
237 }
238 // Fisher-Yates 1938 shuffle.
239 var i: nx_int = n - 1
240 while i > 0 {
241 let bound: nx_int = i + 1
242 let j: nx_int = nx_rng_below(rng, bound)
243 let tmp: nx_int = perm[i]
244 perm[i] = perm[j]
245 perm[j] = tmp
246 i = i - 1
247 }
248 return 0
249}
250
251func nx_shapley_compute_sampled(
252 fn_ptr: func(nx_int, nx_int, *u8) -> nx_int,
253 n_features: nx_int,
254 n_samples: nx_int,
255 seed: nx_int,
256 ctx: *u8,
257 baseline_q10: nx_int) -> *NxShapleyResult {
258
259 if n_features <= 0 { return 0 as *NxShapleyResult }
260 if n_features > NX_SHAPLEY_MAX_FEATURES { return 0 as *NxShapleyResult }
261 var m: nx_int = n_samples
262 if m <= 0 { m = NX_SHAPLEY_DEFAULT_SAMPLES }
263 if m > NX_SHAPLEY_MAX_SAMPLES { m = NX_SHAPLEY_MAX_SAMPLES }
264
265 let result_ptr: *u8 = sys_mmap(NX_SHAPLEY_RESULT_BYTES)
266 let result: *NxShapleyResult = result_ptr as *NxShapleyResult
267 let phi_bytes: nx_size = (n_features as nx_size) * 8
268 let phi: *nx_int = (sys_mmap(phi_bytes)) as *nx_int
269
270 var init_i: nx_int = 0
271 while init_i < n_features {
272 phi[init_i] = 0
273 init_i = init_i + 1
274 }
275
276 let perm_bytes: nx_size = (n_features as nx_size) * 8
277 let perm: *nx_int = (sys_mmap(perm_bytes)) as *nx_int
278 let rng: *NxRng = nx_rng_new(seed)
279
280 var n_calls: nx_int = 0
281 var s: nx_int = 0
282 while s < m {
283 _shapley_permute(perm, n_features, rng)
284 var built_mask: nx_int = 0
285 let f_empty: nx_int = fn_ptr(0, n_features, ctx)
286 n_calls = n_calls + 1
287 var prev_f: nx_int = f_empty
288 var k: nx_int = 0
289 while k < n_features {
290 let feat: nx_int = perm[k]
291 built_mask = _shapley_bit_set(built_mask, feat)
292 let f_cur: nx_int = fn_ptr(built_mask, n_features, ctx)
293 n_calls = n_calls + 1
294 let marginal: nx_int = f_cur - prev_f
295 phi[feat] = phi[feat] + marginal
296 prev_f = f_cur
297 k = k + 1
298 }
299 s = s + 1
300 }
301
302 // Average marginals.
303 var avg_i: nx_int = 0
304 while avg_i < n_features {
305 phi[avg_i] = phi[avg_i] / m
306 avg_i = avg_i + 1
307 }
308
309 let all_mask: nx_int = (1 << n_features) - 1
310 let f_all: nx_int = fn_ptr(all_mask, n_features, ctx)
311 n_calls = n_calls + 1
312
313 result.n_features = n_features
314 result.mode = NX_SHAPLEY_MODE_SAMPLED
315 result.n_evaluations = n_calls
316 result.baseline_q10 = baseline_q10
317 result.full_q10 = f_all
318 result.phi = phi
319 return result
320}
321
322// ===== consistency check: sum(phi) ?= f(all) - f(empty) ===========
323//
324// Efficiency axiom: sum_i phi_i = f(N) - f(empty). Returns absolute
325// difference; caller decides tolerance threshold.
326
327func nx_shapley_efficiency_residual(
328 result: *NxShapleyResult,
329 f_empty: nx_int) -> nx_int {
330
331 var sum_phi: nx_int = 0
332 var i: nx_int = 0
333 while i < result.n_features {
334 sum_phi = sum_phi + result.phi[i]
335 i = i + 1
336 }
337 let expected: nx_int = result.full_q10 - f_empty
338 let diff: nx_int = sum_phi - expected
339 if diff < 0 { return 0 - diff }
340 return diff
341}
342
343// ===== top-K phi by magnitude ======================================
344//
345// Returns the K features with largest |phi| -- the "most-attributed"
346// features for the output. Output indices in out_idx[]; out_phi[]
347// holds signed phi values. K must be <= n_features.
348
349func nx_shapley_topk(
350 result: *NxShapleyResult,
351 k: nx_int,
352 out_idx: *nx_int,
353 out_phi: *nx_int) -> nx_int {
354
355 if k <= 0 { return 0 }
356 if k > result.n_features { return 0 - 1 }
357
358 // Linear-scan partial sort (k typically small). For larger k a
359 // heap would win; left for future optimisation.
360 var picked: *nx_int = (sys_mmap((result.n_features as nx_size) * 8)) as *nx_int
361 var init_i: nx_int = 0
362 while init_i < result.n_features {
363 picked[init_i] = 0
364 init_i = init_i + 1
365 }
366
367 var out_count: nx_int = 0
368 while out_count < k {
369 var best_idx: nx_int = -1
370 var best_abs: nx_int = -1
371 var i: nx_int = 0
372 while i < result.n_features {
373 if picked[i] == 0 {
374 var abs_phi: nx_int = result.phi[i]
375 if abs_phi < 0 { abs_phi = 0 - abs_phi }
376 if abs_phi > best_abs {
377 best_abs = abs_phi
378 best_idx = i
379 }
380 }
381 i = i + 1
382 }
383 if best_idx < 0 { return out_count }
384 out_idx[out_count] = best_idx
385 out_phi[out_count] = result.phi[best_idx]
386 picked[best_idx] = 1
387 out_count = out_count + 1
388 }
389 return out_count
390}
391
392// ===== test fixture: 3-feature additive function ==================
393//
394// f(S) = 10*[1 in S] + 20*[2 in S] + 40*[3 in S]
395// Expected phi = [10, 20, 40] exactly (additive game).
396
397func _test_fn_additive(active_mask: nx_int, n_features: nx_int, ctx: *u8) -> nx_int {
398 var total: nx_int = 0
399 if _shapley_bit_test(active_mask, 0) == 1 { total = total + 10 }
400 if _shapley_bit_test(active_mask, 1) == 1 { total = total + 20 }
401 if _shapley_bit_test(active_mask, 2) == 1 { total = total + 40 }
402 return total
403}
404
405// f(S) = 100 if both features 0 and 1 in S, else 0.
406// Symmetric superadditive game: phi[0] = phi[1] = 50.
407func _test_fn_and_gate(active_mask: nx_int, n_features: nx_int, ctx: *u8) -> nx_int {
408 let b0: nx_int = _shapley_bit_test(active_mask, 0)
409 let b1: nx_int = _shapley_bit_test(active_mask, 1)
410 if b0 == 1 {
411 if b1 == 1 { return 100 }
412 }
413 return 0
414}
415
416// ===== self-test ===================================================
417
418func main() -> nx_int {
419 // ---- exact: 3-feature additive game ----
420 let null_ctx: *u8 = 0 as *u8
421 let r_add: *NxShapleyResult = nx_shapley_compute_exact(
422 _test_fn_additive, 3, null_ctx, 0)
423 if r_add == (0 as *NxShapleyResult) { return 1 }
424 if r_add.n_features != 3 { return 2 }
425 if r_add.mode != NX_SHAPLEY_MODE_EXACT { return 3 }
426 if r_add.phi[0] != 10 { return 4 }
427 if r_add.phi[1] != 20 { return 5 }
428 if r_add.phi[2] != 40 { return 6 }
429
430 // Efficiency: sum(phi) = f(all) - f(empty) = 70 - 0 = 70.
431 let resid_add: nx_int = nx_shapley_efficiency_residual(r_add, 0)
432 if resid_add != 0 { return 7 }
433
434 // ---- exact: 2-feature AND-gate symmetric superadditive ----
435 let r_and: *NxShapleyResult = nx_shapley_compute_exact(
436 _test_fn_and_gate, 2, null_ctx, 0)
437 if r_and == (0 as *NxShapleyResult) { return 10 }
438 if r_and.phi[0] != 50 { return 11 }
439 if r_and.phi[1] != 50 { return 12 }
440 let resid_and: nx_int = nx_shapley_efficiency_residual(r_and, 0)
441 if resid_and != 0 { return 13 }
442
443 // ---- sampled: 3-feature additive, 4096 samples ----
444 let r_samp: *NxShapleyResult = nx_shapley_compute_sampled(
445 _test_fn_additive, 3, NX_MAGIC_4096, 42, null_ctx, 0)
446 if r_samp == (0 as *NxShapleyResult) { return 20 }
447 if r_samp.mode != NX_SHAPLEY_MODE_SAMPLED { return 21 }
448 // Additive games: every permutation gives exact marginal. The
449 // sampled estimator must converge to exact values.
450 if r_samp.phi[0] != 10 { return 22 }
451 if r_samp.phi[1] != 20 { return 23 }
452 if r_samp.phi[2] != 40 { return 24 }
453
454 // ---- top-K: largest-magnitude phi in additive game = feat 2 ----
455 let top_idx: *nx_int = (sys_mmap(16)) as *nx_int
456 let top_phi: *nx_int = (sys_mmap(16)) as *nx_int
457 let n_top: nx_int = nx_shapley_topk(r_add, 2, top_idx, top_phi)
458 if n_top != 2 { return 30 }
459 if top_idx[0] != 2 { return 31 }
460 if top_phi[0] != 40 { return 32 }
461 if top_idx[1] != 1 { return 33 }
462 if top_phi[1] != 20 { return 34 }
463
464 // ---- input validation: refuse N too large for exact ----
465 let r_bad: *NxShapleyResult = nx_shapley_compute_exact(
466 _test_fn_additive, 11, null_ctx, 0)
467 if r_bad != (0 as *NxShapleyResult) { return 40 }
468
469 let r_bad_n: *NxShapleyResult = nx_shapley_compute_exact(
470 _test_fn_additive, 0, null_ctx, 0)
471 if r_bad_n != (0 as *NxShapleyResult) { return 41 }
472
473 return 0
474}