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

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,23 @@ def get_precision() -> int:
     return get('precision')
 
 
+def get_precision() -> int:
+    """
+    Get the default precision.
+
+    Returns
+    -------
+    precision: int
+      The default precision.
+    """
+    precision = get('precision')
+    return 16 if precision == 'bf16' else 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 +267,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 +275,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,14 +338,14 @@ 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}.'
+    assert precision in [64, 32, 16, 'bf16', 8], f'Precision must be in [64, 32, 16, "bf16", 8]. But got {precision}.'
     if precision == 64:
         config.update("jax_enable_x64", True)
     else:
@@ -362,7 +358,7 @@ def _get_uint(precision: int):
         return np.uint64
     elif precision == 32:
         return np.uint32
-    elif precision == 16:
+    elif precision in [16, 'bf16']:
         return np.uint16
     elif precision == 8:
         return np.uint8
@@ -376,7 +372,7 @@ def _get_int(precision: int):
         return np.int64
     elif precision == 32:
         return np.int32
-    elif precision == 16:
+    elif precision in [16, 'bf16']:
         return np.int16
     elif precision == 8:
         return np.int8
@@ -391,8 +387,9 @@ def _get_float(precision: int):
     elif precision == 32:
         return np.float32
     elif precision == 16:
+        return np.float16
+    elif precision == 'bf16':
         return jnp.bfloat16
-        # return np.float16
     else:
         raise ValueError(f'Unsupported precision: {precision}')
 
@@ -403,8 +400,8 @@ def _get_complex(precision: int):
         return np.complex128
     elif precision == 32:
         return np.complex64
-    elif precision == 16:
-        return np.complex32
+    elif precision in [16, 'bf16']:
+        return np.complex64
     else:
         raise ValueError(f'Unsupported precision: {precision}')
 
@@ -430,7 +427,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 +452,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 +478,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 +503,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-0.1.0.post20250102.dist-info → brainstate-0.1.0.post20250104.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250102
+Version: 0.1.0.post20250104
 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.post20250104.dist-info}/RECORD

@@ -2,7 +2,7 @@ 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.py,sha256=wCVlSjavJ9OXc1STOucg-VfXeK9KE443a1SDhkK9lA8,17270
 brainstate/environ_test.py,sha256=jXX3nR1CO74aow5YqfqSd73isj9MWgHQxrwSsEjTDY8,1901
 brainstate/mixin.py,sha256=g7uVUwZphZWsNs9pb48ozG2cDGaj0hs0g3lq8tDk-Sg,11310
 brainstate/mixin_test.py,sha256=Oq_0fwC9vpXDN4t4dTBhWzLdFDNlcYsrcip14F1yECI,3079
@@ -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.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,,