brainstate 0.1.0.post20250104__py2.py3-none-any.whl → 0.1.0.post20250105__py2.py3-none-any.whl

brainstate/compile/_error_if_test.py

@@ -32,6 +32,7 @@ class TestJitError(unittest.TestCase):
         def err_f(x):
             raise ValueError(f'error: {x}')
 
+        bst.compile.jit_error_if(False, err_f, 1.)
         with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
             bst.compile.jit_error_if(True, err_f, 1.)
 

brainstate/compile/_loop_collect_return.py

@@ -206,7 +206,7 @@ def scan(
 
     # evaluate jaxpr, get all states #
     # ------------------------------ #
-    xs_avals = [jax.core.raise_to_shaped(jax.core.get_aval(x)) for x in xs_flat]
+    xs_avals = [jax.core.get_aval(x) for x in xs_flat]
     x_avals = [jax.core.mapped_aval(length, 0, aval) for aval in xs_avals]
     stateful_fun = StatefulFunction(f).make_jaxpr(init, xs_tree.unflatten(x_avals))
     state_trace = stateful_fun.get_state_trace()
@@ -302,7 +302,7 @@ def checkpointed_scan(
         pbar_runner = None
 
     # evaluate jaxpr
-    xs_avals = [jax.core.raise_to_shaped(jax.core.get_aval(x)) for x in xs_flat]
+    xs_avals = [jax.core.get_aval(x) for x in xs_flat]
     x_avals = [jax.core.mapped_aval(length, 0, aval) for aval in xs_avals]
     stateful_fun = StatefulFunction(f).make_jaxpr(init, xs_tree.unflatten(x_avals))
     state_trace = stateful_fun.get_state_trace()

brainstate/compile/_make_jaxpr_test.py

@@ -18,6 +18,7 @@ from __future__ import annotations
 import unittest
 
 import jax
+import jax.extend as je
 import jax.numpy as jnp
 import pytest
 
@@ -84,7 +85,7 @@ class TestMakeJaxpr(unittest.TestCase):
         print(jaxpr)
         jaxpr, _ = bst.compile.make_jaxpr(f3)(jnp.zeros(1))
         print(jaxpr)
-        self.assertTrue(jnp.allclose(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value)[0],
+        self.assertTrue(jnp.allclose(je.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value)[0],
                                      f3(jnp.zeros(1))))
 
     def test_compar_jax_make_jaxpr2(self):

brainstate/compile/_unvmap.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 
 import jax
 import jax.core
+import jax.extend as je
 import jax.interpreters.batching as batching
 import jax.interpreters.mlir as mlir
 import jax.numpy as jnp
@@ -43,7 +44,7 @@ def unvmap(x, op: str = 'any'):
 
 # unvmap_all
 
-unvmap_all_p = jax.core.Primitive("unvmap_all")
+unvmap_all_p = je.core.Primitive("unvmap_all")
 
 
 def unvmap_all(x):

brainstate/environ.py

@@ -249,7 +249,13 @@ def get_precision() -> int:
       The default precision.
     """
     precision = get('precision')
-    return 16 if precision == 'bf16' else precision
+    if precision == 'bf16':
+        return 16
+    if isinstance(precision, int):
+        return precision
+    if isinstance(precision, str):
+        return int(precision)
+    raise ValueError(f'Unsupported precision: {precision}')
 
 
 def set(
@@ -345,8 +351,8 @@ def _set_jax_precision(precision: int | str):
     Args:
       precision: int. The default precision.
     """
-    assert precision in [64, 32, 16, 'bf16', 8], f'Precision must be in [64, 32, 16, "bf16", 8]. But got {precision}.'
-    if precision == 64:
+    # assert precision in [64, 32, 16, 'bf16', 8], f'Precision must be in [64, 32, 16, "bf16", 8]. But got {precision}.'
+    if precision in [64, '64']:
         config.update("jax_enable_x64", True)
     else:
         config.update("jax_enable_x64", False)
@@ -354,13 +360,13 @@ def _set_jax_precision(precision: int | str):
 
 @functools.lru_cache()
 def _get_uint(precision: int):
-    if precision == 64:
+    if precision in [64, '64']:
         return np.uint64
-    elif precision == 32:
+    elif precision in [32, '32']:
         return np.uint32
-    elif precision in [16, 'bf16']:
+    elif precision in [16, '16', 'bf16']:
         return np.uint16
-    elif precision == 8:
+    elif precision in [8, '8']:
         return np.uint8
     else:
         raise ValueError(f'Unsupported precision: {precision}')
@@ -368,13 +374,13 @@ def _get_uint(precision: int):
 
 @functools.lru_cache()
 def _get_int(precision: int):
-    if precision == 64:
+    if precision in [64, '64']:
         return np.int64
-    elif precision == 32:
+    elif precision in [32, '32']:
         return np.int32
-    elif precision in [16, 'bf16']:
+    elif precision in [16, '16', 'bf16']:
         return np.int16
-    elif precision == 8:
+    elif precision in [8, '8']:
         return np.int8
     else:
         raise ValueError(f'Unsupported precision: {precision}')
@@ -382,25 +388,29 @@ def _get_int(precision: int):
 
 @functools.lru_cache()
 def _get_float(precision: int):
-    if precision == 64:
+    if precision in [64, '64']:
         return np.float64
-    elif precision == 32:
+    elif precision in [32, '32']:
         return np.float32
-    elif precision == 16:
+    elif precision in [16, '16']:
         return np.float16
-    elif precision == 'bf16':
+    elif precision in ['bf16']:
         return jnp.bfloat16
+    elif precision in [8, '8']:
+        return jnp.float8_e5m2
     else:
         raise ValueError(f'Unsupported precision: {precision}')
 
 
 @functools.lru_cache()
 def _get_complex(precision: int):
-    if precision == 64:
+    if precision == [64, '64']:
         return np.complex128
-    elif precision == 32:
+    elif precision == [32, '32']:
         return np.complex64
-    elif precision in [16, 'bf16']:
+    elif precision in [16, '16', 'bf16']:
+        return np.complex64
+    elif precision == [8, '8']:
         return np.complex64
     else:
         raise ValueError(f'Unsupported precision: {precision}')

brainstate/environ_test.py

@@ -32,6 +32,10 @@ class TestEnviron(unittest.TestCase):
             self.assertEqual(a.dtype, jnp.float32)
 
         with bst.environ.context(precision=16):
+            a = bst.random.randn(1)
+            self.assertEqual(a.dtype, jnp.float16)
+
+        with bst.environ.context(precision='bf16'):
             a = bst.random.randn(1)
             self.assertEqual(a.dtype, jnp.bfloat16)
 

brainstate/event/_fixedprob_mv.py

@@ -24,6 +24,7 @@ import jax.numpy as jnp
 import numpy as np
 from jax.interpreters import ad
 
+from brainstate import environ
 from brainstate._state import ParamState
 from brainstate.augment import vmap
 from brainstate.init import param
@@ -111,7 +112,7 @@ class FixedProb(Module):
 
     def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
         if self.n_conn > 1:
-            return event_fixed_prob(
+            r = event_fixed_prob(
                 spk,
                 self.weight.value,
                 self.indices,
@@ -123,7 +124,8 @@ class FixedProb(Module):
             weight = self.weight.value
             unit = u.get_unit(weight)
             r = jnp.zeros(spk.shape[:-1] + (self.out_size[-1],), dtype=weight.dtype)
-            return u.maybe_decimal(u.Quantity(r, unit=unit))
+            r = u.maybe_decimal(u.Quantity(r, unit=unit))
+        return u.math.asarray(r, dtype=environ.dftype())
 
 
 def event_fixed_prob(

brainstate/event/_fixedprob_mv_test.py

@@ -128,4 +128,5 @@ class TestFixedProbCSR(parameterized.TestCase):
 
         o2, r2 = jax.jvp(f2, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
         self.assertTrue(jnp.allclose(o1, o2))
-        self.assertTrue(jnp.allclose(r1, r2))
+        # assert jnp.allclose(r1, r2), f'r1={r1}, r2={r2}'
+        self.assertTrue(jnp.allclose(r1, r2, rtol=1e-4, atol=1e-4))

brainstate/event/_xla_custom_op.py

@@ -12,9 +12,13 @@ from jax import tree_util
 from jax.core import Primitive
 from jax.interpreters import batching, ad
 from jax.interpreters import xla, mlir
-from jax.lib import xla_client
 from jaxlib.hlo_helpers import custom_call
 
+if jax.__version_info__ < (0, 4, 35):
+    from jax.lib import xla_client
+else:
+    import jax.extend as je
+
 numba_installed = importlib.util.find_spec('numba') is not None
 
 __all__ = [
@@ -143,7 +147,10 @@ def numba_cpu_custom_call_target(output_ptrs, input_ptrs):
     xla_c_rule = cfunc(sig)(new_f)
     target_name = f'numba_custom_call_{str(xla_c_rule.address)}'
     capsule = ctypes.pythonapi.PyCapsule_New(xla_c_rule.address, b"xla._CUSTOM_CALL_TARGET", None)
-    xla_client.register_custom_call_target(target_name, capsule, "cpu")
+    if jax.__version_info__ < (0, 4, 35):
+        xla_client.register_custom_call_target(target_name, capsule, "cpu")
+    else:
+        je.ffi.register_ffi_target(target_name, capsule, "cpu", api_version=0)
 
     # call
     return custom_call(

{brainstate-0.1.0.post20250104.dist-info → brainstate-0.1.0.post20250105.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250104
+Version: 0.1.0.post20250105
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers

{brainstate-0.1.0.post20250104.dist-info → brainstate-0.1.0.post20250105.dist-info}/RECORD

@@ -2,8 +2,8 @@ brainstate/__init__.py,sha256=A-QKdOvSalsCMxgk80Iz6_xMiUin6con6JaONHfciSY,1526
 brainstate/_state.py,sha256=4aDpLyHGr1VlPXeLSfM3USQG5K4o7orF7IlaBdYrtfE,29098
 brainstate/_state_test.py,sha256=1boTp1w8DiCFLsPwNtlLrlIqGRpkasAmLid5bv2fgP4,2223
 brainstate/_utils.py,sha256=uJ6WWKq3yb05ZdktCQGLWOXsOJveL1H9pR7eev70Jes,1693
-brainstate/environ.py,sha256=wCVlSjavJ9OXc1STOucg-VfXeK9KE443a1SDhkK9lA8,17270
-brainstate/environ_test.py,sha256=jXX3nR1CO74aow5YqfqSd73isj9MWgHQxrwSsEjTDY8,1901
+brainstate/environ.py,sha256=PZnVFWPioUBuWmwCO8wwCKrHQfP3BR-5lYPRl5i5GDA,17698
+brainstate/environ_test.py,sha256=QD6sPCKNtqemVCGwkdImjMazatrvvLr6YeAVcfUnVVY,2045
 brainstate/mixin.py,sha256=g7uVUwZphZWsNs9pb48ozG2cDGaj0hs0g3lq8tDk-Sg,11310
 brainstate/mixin_test.py,sha256=Oq_0fwC9vpXDN4t4dTBhWzLdFDNlcYsrcip14F1yECI,3079
 brainstate/surrogate.py,sha256=YaY6RJ6kzpuPXWFjaWsxWt2MzJfdm5v_jeOR8V_jPoU,48369
@@ -23,17 +23,17 @@ brainstate/compile/_ad_checkpoint_test.py,sha256=R1I76nG4zIqb6g3M_VxWts7rUC1OHJC
 brainstate/compile/_conditions.py,sha256=gApsHKGQrf1QBjoKXDVL7VsoeJ2zFtSc-hFz9nbYcF0,10113
 brainstate/compile/_conditions_test.py,sha256=s9LF6h9LvigvgxUIugTqvgCHBIU8TXS1Ar1OlIxXfrw,8389
 brainstate/compile/_error_if.py,sha256=TFvhqITKkRO9m30GdlUP4eEjJvLWQUhjkujXO9zvrWs,2689
-brainstate/compile/_error_if_test.py,sha256=SJmAfosVoGd4vhfFtb1IvjeFVW914bfTccCg6DoLWYk,1992
+brainstate/compile/_error_if_test.py,sha256=OdJG483IIdOrCHxtHd49OHfOxCSnSkk7GdAUOzSt8bE,2044
 brainstate/compile/_jit.py,sha256=3WBXNTALWPYC9rQH0TPH6w4bjG0BpnZt3RAzUQF5kkc,14045
 brainstate/compile/_jit_test.py,sha256=zD7kck9SQJGmUDolh9P4luKwQ21fBGje1Z4STTEXIuA,4135
-brainstate/compile/_loop_collect_return.py,sha256=_iOVPytbctgyaIOxQZH3A2ZbsSoT7VXnFk6Q6R8-gvA,23360
+brainstate/compile/_loop_collect_return.py,sha256=XwMnKkMH0xTWB1f6GE4NQNK1R2GXTXCiVgulpkdIpc4,23308
 brainstate/compile/_loop_collect_return_test.py,sha256=bA-_11E8A_0jR5umEO3e409y7bb5QYDTgSL-SBaX7kQ,1802
 brainstate/compile/_loop_no_collection.py,sha256=0i31gdQ7sI-d6pvnh08ttUUwdAtpx4uoYhGuf_CyL9s,7343
 brainstate/compile/_loop_no_collection_test.py,sha256=oStB1CSG_iLp9sHdXd1hJNFvlxbzjck9Iy4sABoJDj4,1419
 brainstate/compile/_make_jaxpr.py,sha256=S5O9KUB3bsxoKcfptlV0MRfKA__Ija37WxkakIRL3z0,33010
-brainstate/compile/_make_jaxpr_test.py,sha256=qJUtkyj50JQ6f4UJbOLhvRdkbNn3NSKibFL9jESdQkA,4279
+brainstate/compile/_make_jaxpr_test.py,sha256=3XaX8LUuG6UjolcD83qDVo5odf8FCDppdr9Q6V0NBs4,4303
 brainstate/compile/_progress_bar.py,sha256=eInZPjiqzYE6PWxl_or_lBthDNDO0Ov60Uz0DbuBbZQ,4620
-brainstate/compile/_unvmap.py,sha256=ewbLLNXiI_dBsEBaVzSS0BEXNol22sd9gMzk606lSkM,4139
+brainstate/compile/_unvmap.py,sha256=0i-NvCLDAUe-effJIIEPVsK4WTPbCDBTgw6AqRvq7mE,4163
 brainstate/compile/_util.py,sha256=aCvkTV--g4NsqcodTdBAISt4EwgezCbKzNUV58n-Q_Y,6304
 brainstate/event/__init__.py,sha256=W0ZxbgrcFuYhWTl-GZ0UDoMGfsWmesvG4J_LbTBkex8,937
 brainstate/event/_csr.py,sha256=QDccbgXUklE2iq1w6WdyaFspXY1165uA9UlPltX16OU,30365
@@ -41,14 +41,14 @@ brainstate/event/_csr_mv.py,sha256=HStHvK3KyEMfLsIUslZjgbdU6OsD1yKGrzQOzBXG36M,1
 brainstate/event/_csr_mv_benchmark.py,sha256=xrj2DSWzw0pUHAE1jRBeSRhMW7ogXvDHEdeaZGioNE4,702
 brainstate/event/_csr_mv_test.py,sha256=WQfAvp_3UeCUGAZjr3_aqQvrB-eYZcFEN4v1PBe9fUQ,4012
 brainstate/event/_csr_test.py,sha256=v59rnwTy8jrvqjdGzN75kvLg0wLBmRbthaVRKY2f0Uw,2945
-brainstate/event/_fixedprob_mv.py,sha256=HP5uyFwue5ZNhsU71ZedMQ-Kp5-st89aLKGNhCBmBRA,25457
+brainstate/event/_fixedprob_mv.py,sha256=nR3lhd87t1Vge435QHnFuDp-UBbWoW0Qk1kbsjRHQyc,25541
 brainstate/event/_fixedprob_mv_benchmark.py,sha256=_F_8fH5MNMJZHeSqnq9DYMI9OgYr6JIxBKjbsgeWRv4,4720
-brainstate/event/_fixedprob_mv_test.py,sha256=jijrtJ5fnwcLmA7Tjd3vDlzwfbftmLoVTN4-MPuogVc,4201
+brainstate/event/_fixedprob_mv_test.py,sha256=pVEarvGbqTjnAbxgMVRTAhkyYbvDnlyCJdeOdDD927w,4283
 brainstate/event/_linear_mv.py,sha256=O5qbY31GNV1qEDrZ5kvPbA8Ae-bY5JpUgGtqDFNAeV0,11794
 brainstate/event/_linear_mv_benckmark.py,sha256=hu0WqYMIa3jMoH7Fq9dgxcBjjXGFhghPx9vztyCo1KY,2411
 brainstate/event/_linear_mv_test.py,sha256=V9w41ZP2vu95CyCdCkm-j9Eftqs2kqmeBY809N1-syY,3736
 brainstate/event/_misc.py,sha256=8IpPooXjF2m0-tuo3pGHqThq2yLSNmYziy_zdurZ3NI,1040
-brainstate/event/_xla_custom_op.py,sha256=sR06amc_3mbQDq_ONdxHwr_O8ZKj7S5SNNgo62ILR3U,10797
+brainstate/event/_xla_custom_op.py,sha256=f0OrO6CjsJOUNUCRPpHIRmsb_wgNEym0xBl1tcz8ij4,11016
 brainstate/event/_xla_custom_op_test.py,sha256=rnkGMleXzLfJj4y5QqwfBvCCLTAHe_uabwBDniY-URM,1745
 brainstate/functional/__init__.py,sha256=j6-3Er4fgqWpvntzYCZVB3e5hoz-Z3aqvapITCuDri0,1107
 brainstate/functional/_activations.py,sha256=S0Ok7sq5FTbmJWSejpOCHo1jpKX0gYOLy_TO2IUXM8s,21726
@@ -137,8 +137,8 @@ brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7
 brainstate/util/_struct.py,sha256=0exv0oOiSt1hmx20Y4J2-pCGtCTx13WcAlEYSBkyung,17640
 brainstate/util/_tracers.py,sha256=0r5T4nhxMzI79NtqroqitsdMT4YfpgV5RdYJLS5uJ0w,2285
 brainstate/util/_visualization.py,sha256=n4ZVz10z7VBqA0cKO6vyHwEMprWJgPeEqtITzDMai2Y,1519
-brainstate-0.1.0.post20250104.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250104.dist-info/METADATA,sha256=pGOGZBCF8q5La6-DIO1hdldslPnnt4bfjVnOVWgHRU4,3533
-brainstate-0.1.0.post20250104.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250104.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250104.dist-info/RECORD,,
+brainstate-0.1.0.post20250105.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250105.dist-info/METADATA,sha256=Xec1GNBlHcignyvym-EHzU-JIOUuo3T-IUU2LoCO0sk,3533
+brainstate-0.1.0.post20250105.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250105.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250105.dist-info/RECORD,,