Add on-device NaN debugging and unify StatefulFunction cache (#139)

* refactor(debug_nan_gradient): remove unused exp_exprel function and streamline main examples

* refactor(debug_nan): remove unused parameters and streamline NaN detection tests

* refactor(make_jaxpr): replace mutable hashabledict with immutable CacheKey and unify compilation cache

* refactor(traceback): register brainstate with JAX's traceback filter to exclude internal frames

* refactor(debug): enhance traceback filtering and improve dynamic argument hashing

* refactor(debug_nan_gradient): improve NaN detection output and streamline gradient examples

* refactor(debug_test): update NaN detection test to reference the primitive that introduced NaN

Author: chaoming0625

brainstate/__init__.py

@@ -20,6 +20,15 @@
 __version__ = "0.2.10"
 __versio_info__ = tuple(map(int, __version__.split('.')))
 
+# Register brainstate with JAX's traceback filter so that brainstate internal
+# frames are hidden in user-facing error tracebacks.  This is the same pattern
+# used by Flax, Equinox, and other JAX ecosystem libraries.  Users can still
+# see the full traceback by setting JAX_TRACEBACK_FILTERING=off.
+import os as _os
+from jax._src import traceback_util as _traceback_util
+_traceback_util.register_exclusion(_os.path.dirname(_os.path.abspath(__file__)))
+del _os, _traceback_util
+
 from . import environ
 from . import graph
 from . import mixin