Update precision policy docs and config

Document the new reference/float32 runtime policies and make the default explicit in config.json.

README.md

@@ -2,10 +2,10 @@
 license: apache-2.0
 library_name: mlx
 tags:
-  - mlx
-  - protein-structure-prediction
-  - chai-1
-  - biology
+ - mlx
+ - protein-structure-prediction
+ - chai-1
+ - biology
 pipeline_tag: other
 ---
 
@@ -28,9 +28,18 @@ converted from the upstream TorchScript distribution for use with the
 | `model-feature_embedding.safetensors` | Feature embedding stack | 4.8 MB |
 | `model-bond_loss_input_proj.safetensors` | Bond feature projection | 2 KB |
 
-Total: **~1.2 GB**. All parameters are stored in `float32`;
-`chai-mlx` casts to `bfloat16` at inference time by default (matching the
-mixed-precision policy baked into the upstream TorchScript graphs).
+Total: **~1.2 GB**. All parameters are stored in `float32`; runtime
+precision is selected by `compute_dtype`:
+
+- `reference` (default): matches the reference bundle's mixed-precision
+  boundary. Trunk / confidence run in bf16 where the TorchScript graphs do,
+  while diffusion and other preserved numerically sensitive parameters stay
+  in fp32.
+- `float32`: keeps the MLX port in fp32 throughout.
+
+The shipped `config.json` sets `config_version` to `"1"` and
+`compute_dtype` to `"reference"`, so `ChaiMLX.from_pretrained(...)`
+picks up the reference precision policy by default.
 
 ## Usage
 
@@ -39,17 +48,23 @@ from huggingface_hub import snapshot_download
 from chai_mlx import ChaiMLX
 
 path = snapshot_download("josephjojoe/chai-mlx")
-model = ChaiMLX.from_pretrained(path)
+model = ChaiMLX.from_pretrained(path)  # default: compute_dtype="reference"
+model_fp32 = ChaiMLX.from_pretrained(path, compute_dtype="float32")
 ```
 
 Or directly:
 
 ```python
 model = ChaiMLX.from_pretrained("josephjojoe/chai-mlx")
+model_fp32 = ChaiMLX.from_pretrained(
+    "josephjojoe/chai-mlx",
+    compute_dtype="float32",
+)
 ```
 
 See the [chai-mlx README](https://github.com/josephjojoe/chai-mlx) for
-featurization, inference, and CUDA-comparison workflows.
+featurization, inference, and CUDA-comparison workflows, including the
+CLI `--dtype reference|float32` switch.
 
 ## Provenance
 
@@ -66,11 +81,11 @@ for the pinned `chai-lab` SHA that the CUDA comparison harness runs against.
 
 ## Validation
 
-On 1L2Y (Trp-cage, 20 residues), MLX-vs-CUDA (H100, bf16) Cα RMSD is
-**0.75 Å mean** across 15 sample pairs (3 seeds × 5 diffusion samples);
-CUDA vs NMR ground truth is 0.57 Å mean, so MLX sits ~0.26 Å further from
-experimental truth than CUDA. GDT-TS between MLX and CUDA is 95.1%,
-Cα lDDT 89.8%. Monomers up to ~76 residues (1L2Y, 1VII, 1CRN, 1UBQ) have
+On 1L2Y (Trp-cage, 20 residues), MLX-vs-CUDA (H100, `compute_dtype="reference"`)
+C-alpha RMSD is **0.75 A mean** across 15 sample pairs (3 seeds x 5 diffusion
+samples); CUDA vs NMR ground truth is 0.57 A mean, so MLX sits ~0.26 A further
+from experimental truth than CUDA. GDT-TS between MLX and CUDA is 95.1%,
+C-alpha lDDT 89.8%. Monomers up to ~76 residues (1L2Y, 1VII, 1CRN, 1UBQ) have
 been validated; multimer and ligand targets are untested.
 
 The remaining MLX-vs-CUDA gap is dominated by bf16 fused-kernel rounding

config.json

@@ -1,4 +1,5 @@
 {
+  "config_version": "1",
   "feature_dims": {
     "token": 2638,
     "token_pair": 163,
@@ -92,6 +93,7 @@
       21.375
     ]
   },
+  "compute_dtype": "reference",
   "supported_token_sizes": [
     256,
     384,
@@ -101,6 +103,11 @@
     1536,
     2048
   ],
+  "template_restype_vocab": 33,
+  "template_restype_embed_dim": 32,
+  "num_rbf_radii": 6,
+  "distance_rbf_scale": 4.8,
+  "pocket_rbf_scale": 2.8,
   "attention_mask_value": -10000.0,
   "layer_norm_eps": 1e-05,
   "centroid_eps": 0.0001,