Record MTPLX upgrade defect reproductions

tools/issue103-reference.py

@@ -0,0 +1,92 @@
+"""Model-free #103 oracle receipts. Set PYTHONPATH to the immutable MTPLX checkout.
+
+Run in the existing MLX 0.32.2 reference environment, passing the unchanged
+Bare-Speed artifact folder. Only tokenizer assets are read; no model is loaded.
+"""
+import hashlib
+import importlib.metadata
+import json
+import math
+import sys
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+from mtplx.generation import _mx_lazy_sample, _sample_from_logits
+from mtplx.sampling import SamplerConfig, distribution_from_logits
+from mtplx.runtime import _load_tokenizer_resilient
+
+
+def emit(**receipt):
+    print(json.dumps(receipt, ensure_ascii=False), flush=True)
+
+
+root = Path(sys.argv[1])
+emit(case="environment", mlx=importlib.metadata.version("mlx"),
+     device=mx.device_info(),
+     hashes={n: hashlib.sha256((root/n).read_bytes()).hexdigest()
+             for n in ["config.json", "tokenizer.json", "merges.txt", "chat_template.jinja"]})
+config = SamplerConfig(temperature=1., top_p=0.5, top_k=2)
+values = np.array([math.log(.4), math.log(.3)] + [math.log(.3/62)]*62, dtype=np.float32)
+p = distribution_from_logits(values.astype(np.float64), config)
+emit(case="sampling_law", probabilities=p[:2].tolist())
+counts = [0]*64
+for seed in range(128):
+    token = _mx_lazy_sample(mx.array(values), config, mx.random.key(seed))
+    mx.eval(token)
+    counts[int(token.item())] += 1
+emit(case="sampling_law_lazy", counts=counts[:2])
+for temperature, top_k, top_p in [(.6,20,.95),(.2,8,.8),(1.,2,1.)]:
+    row = values * np.float32(temperature)
+    config = SamplerConfig(temperature=temperature, top_k=top_k, top_p=top_p)
+    probabilities = distribution_from_logits(row.astype(np.float64), config)
+    counts = [0]*64
+    for seed in range(512):
+        token = _mx_lazy_sample(mx.array(row), config, mx.random.key(seed))
+        mx.eval(token)
+        counts[int(token.item())] += 1
+    emit(case="sampling_matrix", temperature=temperature, top_k=top_k, top_p=top_p,
+         ordinary_probabilities=probabilities.tolist(), lazy_counts=counts)
+for name, values in [
+    ("nan", [float("nan")]*64),
+    ("positive_inf", [float("inf")]+[0.]*63),
+    ("all_negative_inf", [float("-inf")]*64),
+    ("valid_mask", [float("-inf"),2.,float("-inf"),1.]+[float("-inf")]*60),
+]:
+    for temperature in [0.,.6,1.]:
+        config = SamplerConfig(temperature=temperature, top_p=.95, top_k=20)
+        try:
+            token, _ = _sample_from_logits(mx.array(values), config, np.random.default_rng(0))
+            emit(case="invalid_logits", row=name, temperature=temperature, token=int(token))
+        except (ValueError, RuntimeError) as error:
+            emit(case="invalid_logits", row=name, temperature=temperature, error=type(error).__name__)
+    token = _mx_lazy_sample(mx.array(values), SamplerConfig(temperature=1., top_k=20, top_p=.95), mx.random.key(7))
+    mx.eval(token)
+    emit(case="invalid_logits_lazy", row=name, token=int(token.item()))
+
+tokenizer = _load_tokenizer_resilient(root, json.loads((root/"config.json").read_text()))
+for text in [
+    "भारत एक विशाल देश है जिसकी संस्कृति बहुत पुरानी है।",
+    "ประเทศไทยมีวัฒนธรรมที่เก่าแก่และสวยงามมาก",
+    "إِنَّ اللُّغَةَ الْعَرَبِيَّةَ جَمِيلَةٌ",
+    "The quick brown fox jumps over the lazy dog near the river bank.",
+    "def foo(x):\n    return x + 1\n",
+    "中国的文化历史悠久。",
+    "a🙂b ©c ∑d —e …f €g ™h\n",
+    "x\u0301y १२३ ١٢٣ ①Ⅷ²\u200dZ",
+    " ①", " Ⅷ", " ²", " 𝟘",
+]:
+    emit(case="tokenizer", text=text, ids=tokenizer.encode(text, add_special_tokens=False))
+
+try:
+    from mtplx.attention_context import attention_phase
+    from mtplx.attention_math import attention_gate
+except ImportError:
+    emit(case="bf16_verify_gate", available=False)
+    sys.exit(0)
+gate = mx.array([-6.84375]*4, dtype=mx.bfloat16)
+value = mx.ones((4,), dtype=mx.bfloat16)
+with attention_phase("decode_verify"):
+    output = attention_gate(value, gate)
+    mx.eval(output)
+    emit(case="bf16_verify_gate", values=output.astype(mx.float32).tolist())