brainstate 0.1.0.post20250422__py2.py3-none-any.whl → 0.1.0.post20250501__py2.py3-none-any.whl

brainstate/_compatible_import.py

@@ -18,7 +18,8 @@
 
 import importlib.util
 from contextlib import contextmanager
-from typing import Iterable, Hashable
+from functools import partial
+from typing import Iterable, Hashable, TypeVar, Callable
 
 import jax
 
@@ -31,8 +32,18 @@ __all__ = [
     'Tracer',
     'to_concrete_aval',
     'brainevent',
+    'safe_map',
+    'safe_zip',
+    'unzip2',
+    'unzip3',
+    'wraps',
 ]
 
+T = TypeVar("T")
+T1 = TypeVar("T1")
+T2 = TypeVar("T2")
+T3 = TypeVar("T3")
+
 brainevent_installed = importlib.util.find_spec('brainevent') is not None
 
 from jax.core import get_aval, Tracer
@@ -53,6 +64,93 @@ else:
         finally:
             trace_ctx.set_axis_env(prev)
 
+if jax.__version_info__ < (0, 6, 0):
+    from jax.util import safe_map, safe_zip, unzip2, unzip3, wraps
+
+else:
+    def safe_map(f, *args):
+        args = list(map(list, args))
+        n = len(args[0])
+        for arg in args[1:]:
+            assert len(arg) == n, f'length mismatch: {list(map(len, args))}'
+        return list(map(f, *args))
+
+
+    def safe_zip(*args):
+        args = list(map(list, args))
+        n = len(args[0])
+        for arg in args[1:]:
+            assert len(arg) == n, f'length mismatch: {list(map(len, args))}'
+        return list(zip(*args))
+
+
+    def unzip2(xys: Iterable[tuple[T1, T2]]
+               ) -> tuple[tuple[T1, ...], tuple[T2, ...]]:
+        """Unzip sequence of length-2 tuples into two tuples."""
+        # Note: we deliberately don't use zip(*xys) because it is lazily evaluated,
+        # is too permissive about inputs, and does not guarantee a length-2 output.
+        xs: list[T1] = []
+        ys: list[T2] = []
+        for x, y in xys:
+            xs.append(x)
+            ys.append(y)
+        return tuple(xs), tuple(ys)
+
+
+    def unzip3(xyzs: Iterable[tuple[T1, T2, T3]]
+               ) -> tuple[tuple[T1, ...], tuple[T2, ...], tuple[T3, ...]]:
+        """Unzip sequence of length-3 tuples into three tuples."""
+        # Note: we deliberately don't use zip(*xyzs) because it is lazily evaluated,
+        # is too permissive about inputs, and does not guarantee a length-3 output.
+        xs: list[T1] = []
+        ys: list[T2] = []
+        zs: list[T3] = []
+        for x, y, z in xyzs:
+            xs.append(x)
+            ys.append(y)
+            zs.append(z)
+        return tuple(xs), tuple(ys), tuple(zs)
+
+
+    def fun_name(fun: Callable):
+        name = getattr(fun, "__name__", None)
+        if name is not None:
+            return name
+        if isinstance(fun, partial):
+            return fun_name(fun.func)
+        else:
+            return "<unnamed function>"
+
+
+    def wraps(
+        wrapped: Callable,
+        namestr: str | None = None,
+        docstr: str | None = None,
+        **kwargs,
+    ) -> Callable[[T], T]:
+        """
+        Like functools.wraps, but with finer-grained control over the name and docstring
+        of the resulting function.
+        """
+
+        def wrapper(fun: T) -> T:
+            try:
+                name = fun_name(wrapped)
+                doc = getattr(wrapped, "__doc__", "") or ""
+                fun.__dict__.update(getattr(wrapped, "__dict__", {}))
+                fun.__annotations__ = getattr(wrapped, "__annotations__", {})
+                fun.__name__ = name if namestr is None else namestr.format(fun=name)
+                fun.__module__ = getattr(wrapped, "__module__", "<unknown module>")
+                fun.__doc__ = (doc if docstr is None
+                               else docstr.format(fun=name, doc=doc, **kwargs))
+                fun.__qualname__ = getattr(wrapped, "__qualname__", fun.__name__)
+                fun.__wrapped__ = wrapped
+            except Exception:
+                pass
+            return fun
+
+        return wrapper
+
 
 def to_concrete_aval(aval):
     aval = get_aval(aval)

brainstate/compile/_make_jaxpr.py

@@ -67,9 +67,15 @@ from jax._src.traceback_util import api_boundary
 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 brainstate._compatible_import import ClosedJaxpr, extend_axis_env_nd
+from brainstate._compatible_import import (
+    ClosedJaxpr,
+    extend_axis_env_nd,
+    safe_map,
+    safe_zip,
+    unzip2,
+    wraps,
+)
 from brainstate._state import State, StateTraceStack
 from brainstate._utils import set_module_as
 from brainstate.typing import PyTree
@@ -89,7 +95,7 @@ def _ensure_index_tuple(x: Any) -> tuple[int, ...]:
     try:
         return (operator.index(x),)
     except TypeError:
-        return tuple(jax.util.safe_map(operator.index, x))
+        return tuple(safe_map(operator.index, x))
 
 
 def _new_arg_fn(frame, trace, aval):
@@ -733,7 +739,7 @@ def _make_jaxpr(
         else:
             axes_specs = _flat_axes_specs(abstracted_axes, *args, **kwargs)
             in_type = pe.infer_lambda_input_type(axes_specs, flat_args)
-            in_avals, keep_inputs = jax.util.unzip2(in_type)
+            in_avals, keep_inputs = unzip2(in_type)
             return in_avals, in_tree, keep_inputs
 
     @wraps(fun)
@@ -744,7 +750,7 @@ def _make_jaxpr(
             dyn_argnums = [i for i in range(len(args)) if i not in static_argnums]
             f, args = jax.api_util.argnums_partial(f, dyn_argnums, args)
         in_avals, in_tree, keep_inputs = _abstractify(args, kwargs)
-        in_type = tuple(jax.util.safe_zip(in_avals, keep_inputs))
+        in_type = tuple(safe_zip(in_avals, keep_inputs))
         f, out_tree = _flatten_fun(f, in_tree)
         f = annotate(f, in_type)
         if jax.__version_info__ < (0, 5, 0):
@@ -758,7 +764,7 @@ def _make_jaxpr(
                 jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f)
         closed_jaxpr = ClosedJaxpr(jaxpr, consts)
         if return_shape:
-            out_avals, _ = jax.util.unzip2(out_type)
+            out_avals, _ = unzip2(out_type)
             out_shapes_flat = [jax.ShapeDtypeStruct(a.shape, a.dtype) for a in out_avals]
             return closed_jaxpr, jax.tree.unflatten(out_tree(), out_shapes_flat)
         return closed_jaxpr

{brainstate-0.1.0.post20250422.dist-info → brainstate-0.1.0.post20250501.dist-info}/METADATA

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

@@ -1,5 +1,5 @@
 brainstate/__init__.py,sha256=AkZyyFkn4fB8g2aT6Rc2MO1xICPpUZuDtdze-eUQNc0,1496
-brainstate/_compatible_import.py,sha256=ibxO3Zklu3fxdVsgP4RUe7FgQqs3cZkKbdRGo9GaZKA,2091
+brainstate/_compatible_import.py,sha256=nEmjwX7aCH_kiXsVJXMjOd_SMj2etMiHxCF3z3H8GF0,5444
 brainstate/_state.py,sha256=Qvb6O3LFUcq_V8wG6GLAgkckImEwfzDi79JFrZ-lVWc,60753
 brainstate/_state_test.py,sha256=b6uvZdVRyC4n6-fYzmHNry1b-gJ6zE_kRSxGinqiHaw,1638
 brainstate/_utils.py,sha256=uJ6WWKq3yb05ZdktCQGLWOXsOJveL1H9pR7eev70Jes,1693
@@ -31,7 +31,7 @@ brainstate/compile/_loop_collect_return.py,sha256=lMw66pARbKmWQnNISkNjuCAzSrKF92
 brainstate/compile/_loop_collect_return_test.py,sha256=r-q-D2fYuO9oIGKCu7ZUZDlZj-RlPthdn7fmP04TxbQ,1804
 brainstate/compile/_loop_no_collection.py,sha256=ij1On2Rj1_Wf7ZixbePNCWJXOi2XqjVHyH31NejV5D8,7587
 brainstate/compile/_loop_no_collection_test.py,sha256=oStB1CSG_iLp9sHdXd1hJNFvlxbzjck9Iy4sABoJDj4,1419
-brainstate/compile/_make_jaxpr.py,sha256=gZdjXq9kd0G_R0Wz_ttiCXGGaKG3MpdJKMZ48xuIJzg,33081
+brainstate/compile/_make_jaxpr.py,sha256=cYK6ReeKJANdoMv4TCuCKYdC6m9J_r-h3AGCmvbUNW0,33088
 brainstate/compile/_make_jaxpr_test.py,sha256=3rfb1oz0u1wHPmnppnmo2BlUBbEE02FlbvD6Rehub9I,4290
 brainstate/compile/_progress_bar.py,sha256=5pCMCEmbTO5XmKtzRUJGA178tuBznWKuh9Kw00wAL1I,7524
 brainstate/compile/_unvmap.py,sha256=oS_Pd2JgnMIUWfdsx_gkTLWnFAcbj8r5QM0D1JjgBT8,4156
@@ -126,8 +126,8 @@ brainstate/util/_pretty_table.py,sha256=c3c2UH8hIZ-lCas_KL462kLgsnBc6pjJiMuapPi6
 brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7562
 brainstate/util/_struct.py,sha256=RNQewdpyNcfntOJYefsDVWQjZW3RLsISKmTHpBhzyc8,17517
 brainstate/util/filter.py,sha256=blTktYNaNgCsuwv7xABjvbWsoi4Fozov6C2cXX-ta2g,14124
-brainstate-0.1.0.post20250422.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250422.dist-info/METADATA,sha256=oKYcNUrBWjw-Hp4-77lv6T2G18f883loEao6NEgHOyU,3655
-brainstate-0.1.0.post20250422.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250422.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250422.dist-info/RECORD,,
+brainstate-0.1.0.post20250501.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250501.dist-info/METADATA,sha256=l6z_DHPZegmbLm2JoQT101SmHbULrVJ5yLn-Dd_NvhM,3655
+brainstate-0.1.0.post20250501.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250501.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250501.dist-info/RECORD,,