brainstate 0.1.0.post20250406__py2.py3-none-any.whl → 0.1.0.post20250420__py2.py3-none-any.whl

brainstate/_compatible_import.py

@@ -0,0 +1,58 @@
+# Copyright 2025 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+
+from contextlib import contextmanager
+from typing import Iterable, Hashable
+
+import jax
+
+__all__ = [
+    'ClosedJaxpr',
+    'Primitive',
+    'extend_axis_env_nd',
+    'jaxpr_as_fun',
+    'get_aval',
+    'Tracer',
+    'to_concrete_aval',
+]
+
+from jax.core import get_aval, Tracer
+
+if jax.__version_info__ < (0, 4, 38):
+    from jax.core import ClosedJaxpr, extend_axis_env_nd, Primitive, jaxpr_as_fun
+else:
+    from jax.extend.core import ClosedJaxpr, Primitive, jaxpr_as_fun
+    from jax.core import trace_ctx
+
+
+    @contextmanager
+    def extend_axis_env_nd(name_size_pairs: Iterable[tuple[Hashable, int]]):
+        prev = trace_ctx.axis_env
+        try:
+            trace_ctx.set_axis_env(prev.extend_pure(name_size_pairs))
+            yield
+        finally:
+            trace_ctx.set_axis_env(prev)
+
+
+def to_concrete_aval(aval):
+    aval = get_aval(aval)
+    if isinstance(aval, Tracer):
+        return aval.to_concrete_value()
+    return aval
+

brainstate/compile/_conditions.py

@@ -24,6 +24,8 @@ from brainstate._utils import set_module_as
 from ._error_if import jit_error_if
 from ._make_jaxpr import StatefulFunction
 from ._util import wrap_single_fun_in_multi_branches, write_back_state_values
+from brainstate._compatible_import import to_concrete_aval, Tracer
+
 
 __all__ = [
     'cond', 'switch', 'ifelse',
@@ -86,7 +88,7 @@ def cond(pred, true_fun: Callable, false_fun: Callable, *operands):
             raise TypeError("Pred type must be either boolean or number, got {}.".format(pred_dtype))
 
     # not jit
-    if jax.config.jax_disable_jit and isinstance(jax.core.get_aval(pred), jax.core.ConcreteArray):
+    if jax.config.jax_disable_jit and not isinstance(to_concrete_aval(pred), Tracer):
         if pred:
             return true_fun(*operands)
         else:

brainstate/compile/_make_jaxpr.py

@@ -55,10 +55,12 @@ from __future__ import annotations
 
 import functools
 import inspect
-import jax
 import operator
 from collections.abc import Hashable, Iterable, Sequence
-from contextlib import ExitStack
+from contextlib import ExitStack, contextmanager
+from typing import Any, Callable, Tuple, Union, Dict, Optional
+
+import jax
 from jax._src import source_info_util
 from jax._src.linear_util import annotate
 from jax._src.traceback_util import api_boundary
@@ -66,17 +68,12 @@ from jax.api_util import shaped_abstractify
 from jax.extend.linear_util import transformation_with_aux, wrap_init
 from jax.interpreters import partial_eval as pe
 from jax.util import wraps
-from typing import Any, Callable, Tuple, Union, Dict, Optional
 
 from brainstate._state import State, StateTraceStack
 from brainstate._utils import set_module_as
 from brainstate.typing import PyTree
 from brainstate.util import PrettyObject
-
-if jax.__version_info__ < (0, 4, 38):
-    from jax.core import ClosedJaxpr
-else:
-    from jax.extend.core import ClosedJaxpr
+from brainstate._compatible_import import ClosedJaxpr, extend_axis_env_nd
 
 AxisName = Hashable
 
@@ -200,7 +197,7 @@ class StatefulFunction(PrettyObject):
 
         # implicit parameters
         self.cache_type = cache_type
-        self._cached_jaxpr: Dict[Any, jax.core.ClosedJaxpr] = dict()
+        self._cached_jaxpr: Dict[Any, ClosedJaxpr] = dict()
         self._cached_out_shapes: Dict[Any, PyTree] = dict()
         self._cached_jaxpr_out_tree: Dict[Any, PyTree] = dict()
         self._cached_state_trace: Dict[Any, StateTraceStack] = dict()
@@ -210,7 +207,7 @@ class StatefulFunction(PrettyObject):
             return None
         return k, v
 
-    def get_jaxpr(self, cache_key: Hashable = ()) -> jax.core.ClosedJaxpr:
+    def get_jaxpr(self, cache_key: Hashable = ()) -> ClosedJaxpr:
         """
         Read the JAX Jaxpr representation of the function.
 
@@ -507,8 +504,8 @@ def make_jaxpr(
     return_shape: bool = False,
     abstracted_axes: Optional[Any] = None,
     state_returns: Union[str, Tuple[str, ...]] = ('read', 'write')
-) -> Callable[..., (Tuple[jax.core.ClosedJaxpr, Tuple[State, ...]] |
-                    Tuple[jax.core.ClosedJaxpr, Tuple[State, ...], PyTree])]:
+) -> Callable[..., (Tuple[ClosedJaxpr, Tuple[State, ...]] |
+                    Tuple[ClosedJaxpr, Tuple[State, ...], PyTree])]:
     """
     Creates a function that produces its jaxpr given example args.
 
@@ -754,12 +751,12 @@ def _make_jaxpr(
             debug_info = pe.debug_info(fun, in_tree, out_tree, True, 'make_jaxpr')
         with ExitStack() as stack:
             if axis_env is not None:
-                stack.enter_context(jax.core.extend_axis_env_nd(axis_env))
+                stack.enter_context(extend_axis_env_nd(axis_env))
             if jax.__version_info__ < (0, 5, 0):
                 jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f, debug_info=debug_info)
             else:
                 jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f)
-        closed_jaxpr = jax.core.ClosedJaxpr(jaxpr, consts)
+        closed_jaxpr = ClosedJaxpr(jaxpr, consts)
         if return_shape:
             out_avals, _ = jax.util.unzip2(out_type)
             out_shapes_flat = [jax.ShapeDtypeStruct(a.shape, a.dtype) for a in out_avals]

brainstate/compile/_make_jaxpr_test.py

@@ -21,11 +21,7 @@ import pytest
 import unittest
 
 import brainstate as bst
-
-if jax.__version_info__ < (0, 4, 38):
-    from jax.core import jaxpr_as_fun
-else:
-    from jax.extend.core import jaxpr_as_fun
+from brainstate._compatible_import import jaxpr_as_fun
 
 
 class TestMakeJaxpr(unittest.TestCase):

brainstate/compile/_unvmap.py

@@ -21,11 +21,8 @@ import jax.interpreters.mlir as mlir
 import jax.numpy as jnp
 
 from brainstate._utils import set_module_as
+from brainstate._compatible_import import Primitive
 
-if jax.__version_info__ < (0, 4, 38):
-    from jax.core import Primitive
-else:
-    from jax.extend.core import Primitive
 
 __all__ = [
     "unvmap",

brainstate/surrogate.py

@@ -22,11 +22,8 @@ import jax.scipy as sci
 from jax.interpreters import batching, ad, mlir
 
 from brainstate.util._pretty_pytree import PrettyObject
+from brainstate._compatible_import import Primitive
 
-if jax.__version_info__ < (0, 4, 38):
-    from jax.core import Primitive
-else:
-    from jax.extend.core import Primitive
 
 __all__ = [
     'Surrogate',

{brainstate-0.1.0.post20250406.dist-info → brainstate-0.1.0.post20250420.dist-info}/METADATA

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

@@ -1,4 +1,5 @@
 brainstate/__init__.py,sha256=AkZyyFkn4fB8g2aT6Rc2MO1xICPpUZuDtdze-eUQNc0,1496
+brainstate/_compatible_import.py,sha256=rFpe8602qIzFTSkGJ5aDm-TaWAsmXGb1p2pphU1jzE8,1720
 brainstate/_state.py,sha256=KJclcHKGrIt8K_rDW3E2dO8g_f_UMcZwGID4UUb9MBE,60751
 brainstate/_state_test.py,sha256=UBbbGJ8cb9dJ3NeySf-TNs_nNP47Ax8CP7QL_b32MAA,1636
 brainstate/_utils.py,sha256=uJ6WWKq3yb05ZdktCQGLWOXsOJveL1H9pR7eev70Jes,1693
@@ -6,7 +7,7 @@ brainstate/environ.py,sha256=PllYYZKqany3G7NzIwoUPplLAePbyza6kJGXTPgJK-c,17698
 brainstate/environ_test.py,sha256=khZ_-SUJL6rQCgndeYV98ruUIHGTwFDtITrOs_olmuo,2043
 brainstate/mixin.py,sha256=g7uVUwZphZWsNs9pb48ozG2cDGaj0hs0g3lq8tDk-Sg,11310
 brainstate/mixin_test.py,sha256=Oq_0fwC9vpXDN4t4dTBhWzLdFDNlcYsrcip14F1yECI,3079
-brainstate/surrogate.py,sha256=wWYw-TxaFxHVneXuHjWD1UtTcOTk3XRSnhRtUkt_Hb8,53580
+brainstate/surrogate.py,sha256=nSv2oaO0HBfbsKECr8ctpQSKbZveQAzidGdUrMcZQYU,53508
 brainstate/transform.py,sha256=OvshYpPnp3YXPG6riy15Ve7jcX8t0aaATl1xZnsFeic,858
 brainstate/typing.py,sha256=988gX1tvwtyYnYjmej90OaRxoMoBIPO0-DSrXXGxojM,10523
 brainstate/augment/__init__.py,sha256=Q9-JIwQ1FNn8VLS1MA9MrSylbvUjWSw98whrI3NIuKo,1229
@@ -20,7 +21,7 @@ brainstate/augment/_random.py,sha256=ikRzNoDDE2BkARajDsBhNlngCUrghzGSZUDmEGvVors
 brainstate/compile/__init__.py,sha256=fQtG316MLkeeu1Ssp54Kghw1PwbGK5gNq9yRVJu0wjA,1474
 brainstate/compile/_ad_checkpoint.py,sha256=3wv-f89oo94XeWwRV5LcRot0Nz7xTk5_PdjEDyUMsoo,9394
 brainstate/compile/_ad_checkpoint_test.py,sha256=R1I76nG4zIqb6g3M_VxWts7rUC1OHJCjtQhPkcbXodk,1746
-brainstate/compile/_conditions.py,sha256=POPdqRvFbd0dETkShPbsXUvhRQI8MhgU71ijnvtcqP0,10164
+brainstate/compile/_conditions.py,sha256=G-MioUGu2gjYV7wqulxTx8TfdVPx4AWbBi9V3sFw0uY,10221
 brainstate/compile/_conditions_test.py,sha256=3Y1drxdNvTAm5CZu_ui00v9RTIaLDWoiRbzbmJjZrjQ,8387
 brainstate/compile/_error_if.py,sha256=JIzvnAAW2Ze8aniw_CLZYktb1_WVtKLhP3eGCUhQNgg,2687
 brainstate/compile/_error_if_test.py,sha256=j4x2bzWIWstwLzzt3R9hmencinSSveSrLt5AB7PJC1A,2042
@@ -30,10 +31,10 @@ brainstate/compile/_loop_collect_return.py,sha256=-LsP7fkHmAyGnDOKa3BxxYOEWe8M2J
 brainstate/compile/_loop_collect_return_test.py,sha256=D9RQ5RyQHkqBr4nmSK-yM_uge3EC6uVm_Dzy42g3vtg,1802
 brainstate/compile/_loop_no_collection.py,sha256=2OEVtv5XztOx-e0focZ1UnWkXmFzmDskjHJXuVXmuhA,7587
 brainstate/compile/_loop_no_collection_test.py,sha256=oStB1CSG_iLp9sHdXd1hJNFvlxbzjck9Iy4sABoJDj4,1419
-brainstate/compile/_make_jaxpr.py,sha256=8iV8XyvkMH3n3wbEWZAgZtbrUxryljwQJD6o5DMW9Lc,33189
-brainstate/compile/_make_jaxpr_test.py,sha256=fZe3K4RHFLmMAeXZoFZ5RyxgXvncTcuMQdjmOROJtKU,4365
+brainstate/compile/_make_jaxpr.py,sha256=kh2bX1kQL9VLBic8E6MNa-gBO_G-4tJuFKTP517_PYE,33097
+brainstate/compile/_make_jaxpr_test.py,sha256=EYUDoVKm2lDQeVCC3rBElRzH61LO76C61XgmEFtKcbU,4288
 brainstate/compile/_progress_bar.py,sha256=3Z3OVcc5sl9FK9Fkt813l20MNzEfa6UZ9lJrvSgXTCU,7522
-brainstate/compile/_unvmap.py,sha256=uCvQjvb8J7kT0kalX576mrAPvQuCh_W76EPdgZ53kTM,4230
+brainstate/compile/_unvmap.py,sha256=dHLA6jkZdZ9_hPZ-4ovlPmZpi7Fl6Z4P6bta_Zgo1e0,4158
 brainstate/compile/_util.py,sha256=a_tunKZ1OzVowCI2JmcniQz5P6bqZ4BJkDKmA_h7s6Y,6313
 brainstate/functional/__init__.py,sha256=j6-3Er4fgqWpvntzYCZVB3e5hoz-Z3aqvapITCuDri0,1107
 brainstate/functional/_activations.py,sha256=VmCU9HOKWbysxuJFBN-JsShS4loNMG_E6IXfky1tX-s,21724
@@ -121,8 +122,8 @@ brainstate/util/_pretty_table.py,sha256=NM_6VAW6oL9jojsK0-RkQGHnDzLy_fn_hgzl5R8o
 brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7562
 brainstate/util/_struct.py,sha256=F5GfFURITAIYTwf17_xypkZU1wvoL4dUCviPnr_eCtw,17515
 brainstate/util/filter.py,sha256=Zw0H42NwAi2P7dBr3ISv2VpkB5jqoWnV4Kpd61gq66o,14126
-brainstate-0.1.0.post20250406.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250406.dist-info/METADATA,sha256=VK2aJl8p_7ET5VL8_UcAQTTOg2ynbbs5msPCZiFEjpE,3689
-brainstate-0.1.0.post20250406.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250406.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250406.dist-info/RECORD,,
+brainstate-0.1.0.post20250420.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250420.dist-info/METADATA,sha256=gHk2WgxJm1CirBO0QaiO85mcexK3eziLD9ahYST_RZ0,3689
+brainstate-0.1.0.post20250420.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250420.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250420.dist-info/RECORD,,