wassname/ml_debug
1
0
Fork 0
mirror of https://github.com/wassname/ml_debug.git synced 2026-08-12 12:10:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(ml_debug): add JAX grep patterns and diagnostic equivalents

Browse Source 
refs/static_analysis.md: JAX-specific grep patterns (in-place mutation,
print side effects, key reuse, numpy escape, cast behavior).
refs/diagnostics.md: JAX equivalents table (NaN detection, gradcheck,
disable_jit, debug.print, debug.breakpoint, checkify).
This commit is contained in:
wassname
2026-03-06 14:10:39 +08:00
parent 7ac7aacac7
commit bbe3fe0985
2 changed files with 28 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
refs/diagnostics.md
+13
View File
@@ -169,3 +169,16 @@ for name, p in model.named_parameters():
# Bad signs: all zeros, huge values (>100), std ~0 (collapsed), NaN.
# After training: weights diverging to +/-inf = exploding. All same value = dead.
```
---
## JAX diagnostic equivalents
| Diagnostic | PyTorch | JAX |
|------------|---------|-----|
| NaN detection | `torch.autograd.detect_anomaly()` | `jax.config.update("jax_debug_nans", True)` |
| Gradient check | `torch.autograd.gradcheck(fn, inputs)` | `jax.test_util.check_grads(fn, args, order=2)` |
| Eager debug (no compile) | N/A (already eager) | `jax.config.update("jax_disable_jit", True)` |
| Print inside compiled | N/A | `jax.debug.print("{x}", x=x)` |
| Breakpoint inside compiled | `pdb.set_trace()` | `jax.debug.breakpoint()` |
| Runtime assertions inside compiled | `assert` | `jax.experimental.checkify` |