brainstate 0.1.0.post20250102__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

@@ -31,13 +31,12 @@ from jax import config, devices, numpy as jnp
 from jax.typing import DTypeLike
 
 from .mixin import Mode
-from .util import MemScaling
 
 __all__ = [
     # functions for environment settings
     'set', 'context', 'get', 'all', 'set_host_device_count', 'set_platform',
     # functions for getting default behaviors
-    'get_host_device_count', 'get_platform', 'get_dt', 'get_mode', 'get_mem_scaling', 'get_precision',
+    'get_host_device_count', 'get_platform', 'get_dt', 'get_mode', 'get_precision',
     # functions for default data types
     'dftype', 'ditype', 'dutype', 'dctype',
     # others
@@ -94,7 +93,7 @@ def context(**kwargs):
                          'Please use set_host_device_count() or set() for the global setting.')
 
     if 'precision' in kwargs:
-        last_precision = get_precision()
+        last_precision = _get_precision()
         _set_jax_precision(kwargs['precision'])
 
     try:
@@ -203,17 +202,6 @@ def get_mode() -> Mode:
     return get('mode')
 
 
-def get_mem_scaling() -> MemScaling:
-    """Get the default computing membrane_scaling.
-
-    Returns
-    -------
-    membrane_scaling: MemScaling
-      The default computing membrane_scaling.
-    """
-    return get('mem_scaling')
-
-
 def get_platform() -> str:
     """Get the computing platform.
 
@@ -239,7 +227,7 @@ def get_host_device_count():
     return int(match.group(1)) if match else 1
 
 
-def get_precision() -> int:
+def _get_precision() -> int | str:
     """
     Get the default precision.
 
@@ -251,11 +239,29 @@ def get_precision() -> int:
     return get('precision')
 
 
+def get_precision() -> int:
+    """
+    Get the default precision.
+
+    Returns
+    -------
+    precision: int
+      The default precision.
+    """
+    precision = get('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(
     platform: str = None,
     host_device_count: int = None,
-    mem_scaling: MemScaling = None,
-    precision: int = None,
+    precision: int | str = None,
     mode: Mode = None,
     **kwargs
 ):
@@ -267,8 +273,7 @@ def set(
     Args:
       platform: str. The computing platform. Either 'cpu', 'gpu' or 'tpu'.
       host_device_count: int. The number of host devices.
-      mem_scaling: MemScaling. The membrane scaling.
-      precision: int. The default precision.
+      precision: int, str. The default precision.
       mode: Mode. The computing mode.
       **kwargs: dict. Other environment settings.
     """
@@ -276,9 +281,6 @@ def set(
         set_platform(platform)
     if host_device_count is not None:
         set_host_device_count(host_device_count)
-    if mem_scaling is not None:
-        assert isinstance(mem_scaling, MemScaling), 'mem_scaling must be a MemScaling instance.'
-        kwargs['mem_scaling'] = mem_scaling
     if precision is not None:
         _set_jax_precision(precision)
         kwargs['precision'] = precision
@@ -342,15 +344,15 @@ def set_platform(platform: str):
         DFAULT.functions['platform'](platform)
 
 
-def _set_jax_precision(precision: int):
+def _set_jax_precision(precision: int | str):
     """
     Set the default precision.
 
     Args:
       precision: int. The default precision.
     """
-    assert precision in [64, 32, 16, 8], f'Precision must be in [64, 32, 16, 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)
@@ -358,13 +360,13 @@ def _set_jax_precision(precision: int):
 
 @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 == 16:
+    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}')
@@ -372,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 == 16:
+    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}')
@@ -386,25 +388,30 @@ 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 in ['bf16']:
         return jnp.bfloat16
-        # return np.float16
+    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, '16', 'bf16']:
+        return np.complex64
+    elif precision == [8, '8']:
         return np.complex64
-    elif precision == 16:
-        return np.complex32
     else:
         raise ValueError(f'Unsupported precision: {precision}')
 
@@ -430,7 +437,7 @@ def dftype() -> DTypeLike:
     float_dtype: DTypeLike
       The default floating data type.
     """
-    return _get_float(get_precision())
+    return _get_float(_get_precision())
 
 
 def ditype() -> DTypeLike:
@@ -455,7 +462,7 @@ def ditype() -> DTypeLike:
     int_dtype: DTypeLike
       The default integer data type.
     """
-    return _get_int(get_precision())
+    return _get_int(_get_precision())
 
 
 def dutype() -> DTypeLike:
@@ -481,7 +488,7 @@ def dutype() -> DTypeLike:
     uint_dtype: DTypeLike
       The default unsigned integer data type.
     """
-    return _get_uint(get_precision())
+    return _get_uint(_get_precision())
 
 
 def dctype() -> DTypeLike:
@@ -506,7 +513,7 @@ def dctype() -> DTypeLike:
     complex_dtype: DTypeLike
       The default complex data type.
     """
-    return _get_complex(get_precision())
+    return _get_complex(_get_precision())
 
 
 def tolerance():

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.post20250102.dist-info → brainstate-0.1.0.post20250105.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250102
+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.post20250102.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=G6r_rqfbofRbjFFalRu_DHaL7ruFTeLRXBQDXM6P-tQ,17477
-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.post20250102.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250102.dist-info/METADATA,sha256=KtQbKvFh7Z_WjFQ4e0s3IkEvtsag6MpOZEvCJQ5Mj5k,3533
-brainstate-0.1.0.post20250102.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250102.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250102.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,,