brainstate 0.0.2.post20240825__py2.py3-none-any.whl → 0.0.2.post20240826__py2.py3-none-any.whl

brainstate/transform/_loop_no_collection.py

@@ -0,0 +1,170 @@
+# Copyright 2024 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.
+# ==============================================================================
+
+from __future__ import annotations
+
+import math
+from typing import Any, Callable, TypeVar
+
+import jax
+
+from brainstate._utils import set_module_as
+from ._conditions import _wrapped_fun
+from ._loop_collect_return import _bounded_while_loop
+from ._make_jaxpr import StatefulFunction, _assign_state_values
+
+X = TypeVar('X')
+Y = TypeVar('Y')
+T = TypeVar('T')
+Carry = TypeVar('Carry')
+BooleanNumeric = Any  # A bool, or a Boolean array.
+
+__all__ = [
+  # while loop
+  'while_loop', 'bounded_while_loop',
+]
+
+
+@set_module_as('brainstate.transform')
+def while_loop(
+    cond_fun: Callable[[T], BooleanNumeric],
+    body_fun: Callable[[T], T],
+    init_val: T
+) -> T:
+  """
+  Call ``body_fun`` repeatedly in a loop while ``cond_fun`` is True.
+
+  The `Haskell-like type signature`_ in brief is
+
+  .. code-block:: haskell
+
+    while_loop :: (a -> Bool) -> (a -> a) -> a -> a
+
+  The semantics of ``while_loop`` are given by this Python implementation::
+
+    def while_loop(cond_fun, body_fun, init_val):
+      val = init_val
+      while cond_fun(val):
+        val = body_fun(val)
+      return val
+
+  Unlike that Python version, ``while_loop`` is a JAX primitive and is lowered
+  to a single WhileOp. That makes it useful for reducing compilation times
+  for jit-compiled functions, since native Python loop constructs in an ``@jit``
+  function are unrolled, leading to large XLA computations.
+
+  Also unlike the Python analogue, the loop-carried value ``val`` must hold a
+  fixed shape and dtype across all iterations (and not just be consistent up to
+  NumPy rank/shape broadcasting and dtype promotion rules, for example). In
+  other words, the type ``a`` in the type signature above represents an array
+  with a fixed shape and dtype (or a nested tuple/list/dict container data
+  structure with a fixed structure and arrays with fixed shape and dtype at the
+  leaves).
+
+  Another difference from using Python-native loop constructs is that
+  ``while_loop`` is not reverse-mode differentiable because XLA computations
+  require static bounds on memory requirements.
+
+  Args:
+    cond_fun: function of type ``a -> Bool``.
+    body_fun: function of type ``a -> a``.
+    init_val: value of type ``a``, a type that can be a scalar, array, or any
+      pytree (nested Python tuple/list/dict) thereof, representing the initial
+      loop carry value.
+
+  Returns:
+    The output from the final iteration of body_fun, of type ``a``.
+
+  .. _Haskell-like type signature: https://wiki.haskell.org/Type_signature
+  """
+  if not (callable(body_fun) and callable(cond_fun)):
+    raise TypeError("while_loop: body_fun and cond_fun arguments should be callable.")
+  if jax.config.jax_disable_jit:
+    try:
+      val = init_val
+      while cond_fun(val):
+        val = body_fun(val)
+      return val
+    except jax.core.ConcretizationTypeError:
+      # Can't run this while_loop in Python (e.g. because there's a vmap
+      # transformation on it), so we fall back to the primitive version.
+      pass
+
+  # evaluate jaxpr
+  stateful_cond = StatefulFunction(cond_fun).make_jaxpr(init_val)
+  stateful_body = StatefulFunction(body_fun).make_jaxpr(init_val)
+  all_states = tuple(set(stateful_cond.get_states() + stateful_body.get_states()))
+  new_cond_fun = _wrapped_fun(stateful_cond, all_states, return_states=False)
+  new_body_fun = _wrapped_fun(stateful_body, all_states, return_states=True)
+
+  # while_loop
+  state_vals, final_val = jax.lax.while_loop(new_cond_fun,
+                                             new_body_fun,
+                                             (tuple(st.value for st in all_states), init_val))
+  _assign_state_values(all_states, state_vals)
+  return final_val
+
+
+def bounded_while_loop(
+    cond_fun: Callable[[T], BooleanNumeric],
+    body_fun: Callable[[T], T],
+    init_val: T,
+    *,
+    max_steps: int,
+    base: int = 16,
+):
+  """
+  While loop with a bound on the maximum number of steps.
+
+  This function is useful when you want to ensure that a while loop terminates
+  even if the condition function is never false. The function is implemented
+  using a scan operation, so it is reverse-mode differentiable.
+
+  Args:
+    cond_fun: A function of type ``a -> Bool``.
+    body_fun: A function of type ``a -> a``.
+    init_val: The initial value of type ``a``.
+    max_steps: A bound on the maximum number of steps, after which the loop
+      terminates unconditionally.
+    base: Run time will increase slightly as `base` increases. Compilation time will
+      decrease substantially as `math.ceil(math.log(max_steps, base))` decreases.
+      (Which happens as `base` increases.)
+
+  Returns:
+    The final value, as if computed by a `lax.while_loop`.
+  """
+
+  # checking
+  if not isinstance(max_steps, int) or max_steps < 0:
+    raise ValueError("max_steps must be a non-negative integer")
+  init_val = jax.tree.map(jax.numpy.asarray, init_val)
+  if max_steps == 0:
+    return init_val
+
+  # evaluate jaxpr
+  stateful_cond = StatefulFunction(cond_fun).make_jaxpr(init_val)
+  stateful_body = StatefulFunction(body_fun).make_jaxpr(init_val)
+  all_states = tuple(set(stateful_cond.get_states() + stateful_body.get_states()))
+  new_cond_fun = _wrapped_fun(stateful_cond, all_states, return_states=False)
+  new_body_fun = _wrapped_fun(stateful_body, all_states, return_states=True)
+
+  # initial value
+  init_val = (tuple(st.value for st in all_states), init_val)
+
+  # while_loop
+  rounded_max_steps = base ** int(math.ceil(math.log(max_steps, base)))
+  state_vals, val = _bounded_while_loop(new_cond_fun, new_body_fun, init_val, rounded_max_steps, base, None)
+  _assign_state_values(all_states, state_vals)
+  return val

brainstate/transform/_mapping.py

@@ -0,0 +1,109 @@
+# Copyright 2024 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.
+# ==============================================================================
+
+from __future__ import annotations
+
+import jax
+
+from ._loop_collect_return import scan
+
+__all__ = [
+  'map',
+]
+
+
+def _batch_and_remainder(x, batch_size: int):
+  leaves, treedef = jax.tree.flatten(x)
+
+  scan_leaves = []
+  remainder_leaves = []
+
+  for leaf in leaves:
+    num_batches, _ = divmod(leaf.shape[0], batch_size)
+    total_batch_elems = num_batches * batch_size
+    scan_leaves.append(leaf[:total_batch_elems].reshape(num_batches, batch_size, *leaf.shape[1:]))
+    remainder_leaves.append(leaf[total_batch_elems:])
+
+  scan_tree = treedef.unflatten(scan_leaves)
+  remainder_tree = treedef.unflatten(remainder_leaves)
+  return scan_tree, remainder_tree
+
+
+def map(
+    f,
+    xs,
+    *,
+    batch_size: int | None = None,
+):
+  """Map a function over leading array axes.
+
+  Like Python's builtin map, except inputs and outputs are in the form of
+  stacked arrays. Consider using the :func:`~jax.vmap` transform instead, unless you
+  need to apply a function element by element for reduced memory usage or
+  heterogeneous computation with other control flow primitives.
+
+  When ``xs`` is an array type, the semantics of :func:`~map` are given by this
+  Python implementation::
+
+    def map(f, xs):
+      return np.stack([f(x) for x in xs])
+
+  Like :func:`~scan`, :func:`~map` is implemented in terms of JAX primitives so
+  many of the same advantages over a Python loop apply: ``xs`` may be an
+  arbitrary nested pytree type, and the mapped computation is compiled only
+  once.
+
+  If ``batch_size`` is provided, the computation is executed in batches of that size
+  and parallelized using :func:`~jax.vmap`. This can be used as either a more performant
+  version of ``map`` or as a memory-efficient version of ``vmap``. If the axis is not
+  divisible by the batch size, the remainder is processed in a separate ``vmap`` and
+  concatenated to the result.
+
+    >>> x = jax.numpy.ones((10, 3, 4))
+    >>> def f(x):
+    ...   print('inner shape:', x.shape)
+    ...   return x + 1
+    >>> y = map(f, x, batch_size=3)
+    inner shape: (3, 4)
+    inner shape: (3, 4)
+    >>> y.shape
+    (10, 3, 4)
+
+  In the example above, "inner shape" is printed twice, once while tracing the batched
+  computation and once while tracing the remainder computation.
+
+  Args:
+    f: a Python function to apply element-wise over the first axis or axes of
+      ``xs``.
+    xs: values over which to map along the leading axis.
+    batch_size: (optional) integer specifying the size of the batch for each step to execute
+      in parallel.
+
+  Returns:
+    Mapped values.
+  """
+  if batch_size is not None:
+    scan_xs, remainder_xs = _batch_and_remainder(xs, batch_size)
+    g = lambda _, x: ((), jax.vmap(f)(x))
+    _, scan_ys = scan(g, (), scan_xs)
+    remainder_ys = jax.vmap(f)(remainder_xs)
+    flatten = lambda x: x.reshape(-1, *x.shape[2:])
+    ys = jax.tree.map(
+      lambda x, y: jax.numpy.concatenate([flatten(x), y], axis=0), scan_ys, remainder_ys,
+    )
+  else:
+    g = lambda _, x: ((), f(x))
+    _, ys = scan(g, (), xs)
+  return ys

brainstate/transform/_unvmap.py

@@ -0,0 +1,143 @@
+import jax
+import jax.core
+import jax.interpreters.batching as batching
+import jax.interpreters.mlir as mlir
+import jax.numpy as jnp
+
+from brainstate._utils import set_module_as
+
+__all__ = [
+  "unvmap",
+]
+
+
+@set_module_as('brainstate.transform')
+def unvmap(x, op: str = 'any'):
+  if op == 'all':
+    return unvmap_all(x)
+  elif op == 'any':
+    return unvmap_any(x)
+  elif op == 'none':
+    return _without_vmap(x)
+  elif op == 'max':
+    return unvmap_max(x)
+  else:
+    raise ValueError(f'Do not support type: {op}')
+
+
+# unvmap_all
+
+unvmap_all_p = jax.core.Primitive("unvmap_all")
+
+
+def unvmap_all(x):
+  """As `jnp.all`, but ignores batch dimensions."""
+  return unvmap_all_p.bind(x)
+
+
+def _unvmap_all_impl(x):
+  return jnp.all(x)
+
+
+def _unvmap_all_abstract_eval(x):
+  return jax.core.ShapedArray(shape=(), dtype=jax.numpy.bool_.dtype)  # pyright: ignore
+
+
+def _unvmap_all_batch(x, batch_axes):
+  (x,) = x
+  return unvmap_all(x), batching.not_mapped
+
+
+unvmap_all_p.def_impl(_unvmap_all_impl)
+unvmap_all_p.def_abstract_eval(_unvmap_all_abstract_eval)
+batching.primitive_batchers[unvmap_all_p] = _unvmap_all_batch  # pyright: ignore
+mlir.register_lowering(
+  unvmap_all_p,
+  mlir.lower_fun(_unvmap_all_impl, multiple_results=False),
+)
+
+# unvmap_any
+
+unvmap_any_p = jax.core.Primitive("unvmap_any")
+
+
+def unvmap_any(x):
+  """As `jnp.any`, but ignores batch dimensions."""
+  return unvmap_any_p.bind(x)
+
+
+def _unvmap_any_impl(x):
+  return jnp.any(x)
+
+
+def _unvmap_any_abstract_eval(x):
+  return jax.core.ShapedArray(shape=(), dtype=jax.numpy.bool_.dtype)  # pyright: ignore
+
+
+def _unvmap_any_batch(x, batch_axes):
+  (x,) = x
+  return unvmap_any(x), batching.not_mapped
+
+
+unvmap_any_p.def_impl(_unvmap_any_impl)
+unvmap_any_p.def_abstract_eval(_unvmap_any_abstract_eval)
+batching.primitive_batchers[unvmap_any_p] = _unvmap_any_batch  # pyright: ignore
+mlir.register_lowering(
+  unvmap_any_p,
+  mlir.lower_fun(_unvmap_any_impl, multiple_results=False),
+)
+
+# unvmap_max
+
+unvmap_max_p = jax.core.Primitive("unvmap_max")
+
+
+def unvmap_max(x):
+  """As `jnp.max`, but ignores batch dimensions."""
+  return unvmap_max_p.bind(x)
+
+
+def _unvmap_max_impl(x):
+  return jnp.max(x)
+
+
+def _unvmap_max_abstract_eval(x):
+  return jax.core.ShapedArray(shape=(), dtype=x.dtype)
+
+
+def _unvmap_max_batch(x, batch_axes):
+  (x,) = x
+  return unvmap_max(x), batching.not_mapped
+
+
+unvmap_max_p.def_impl(_unvmap_max_impl)
+unvmap_max_p.def_abstract_eval(_unvmap_max_abstract_eval)
+batching.primitive_batchers[unvmap_max_p] = _unvmap_max_batch  # pyright: ignore
+mlir.register_lowering(
+  unvmap_max_p,
+  mlir.lower_fun(_unvmap_max_impl, multiple_results=False),
+)
+
+
+def _without_vmap(x):
+  return _no_vmap_prim.bind(x)
+
+
+def _without_vmap_imp(x):
+  return x
+
+
+def _without_vmap_abs(x):
+  return x
+
+
+def _without_vmap_batch(x, batch_axes):
+  (x,) = x
+  return _without_vmap(x), batching.not_mapped
+
+
+_no_vmap_prim = jax.core.Primitive('no_vmap')
+_no_vmap_prim.def_impl(_without_vmap_imp)
+_no_vmap_prim.def_abstract_eval(_without_vmap_abs)
+batching.primitive_batchers[_no_vmap_prim] = _without_vmap_batch
+mlir.register_lowering(_no_vmap_prim, mlir.lower_fun(_without_vmap_imp, multiple_results=False))

brainstate/typing.py

@@ -15,8 +15,9 @@
 
 
 import functools as ft
+import inspect
 import typing
-from typing import Sequence, Protocol, Union, Any, Generic, TypeVar
+from typing import Sequence, Protocol, Union, Any, Generic, TypeVar, Tuple
 
 import brainunit as bu
 import jax
@@ -34,6 +35,59 @@ __all__ = [
 
 _T = TypeVar("_T")
 
+_Annotation = TypeVar("_Annotation")
+
+
+class _Array(Generic[_Annotation]):
+  pass
+
+
+_Array.__module__ = "builtins"
+
+
+def _item_to_str(item: Union[str, type, slice]) -> str:
+  if isinstance(item, slice):
+    if item.step is not None:
+      raise NotImplementedError
+    return _item_to_str(item.start) + ": " + _item_to_str(item.stop)
+  elif item is ...:
+    return "..."
+  elif inspect.isclass(item):
+    return item.__name__
+  else:
+    return repr(item)
+
+
+def _maybe_tuple_to_str(
+    item: Union[str, type, slice, Tuple[Union[str, type, slice], ...]]
+) -> str:
+  if isinstance(item, tuple):
+    if len(item) == 0:
+      # Explicit brackets
+      return "()"
+    else:
+      # No brackets
+      return ", ".join([_item_to_str(i) for i in item])
+  else:
+    return _item_to_str(item)
+
+
+class Array:
+  def __class_getitem__(cls, item):
+    class X:
+      pass
+
+    X.__module__ = "builtins"
+    X.__qualname__ = _maybe_tuple_to_str(item)
+    return _Array[X]
+
+
+# Same __module__ trick here again. (So that we get the correct display when
+# doing `def f(x: Array)` as well as `def f(x: Array["dim"])`.
+#
+# Don't need to set __qualname__ as that's already correct.
+Array.__module__ = "builtins"
+
 
 class _FakePyTree(Generic[_T]):
   pass

{brainstate-0.0.2.post20240825.dist-info → brainstate-0.0.2.post20240826.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.0.2.post20240825
+Version: 0.0.2.post20240826
 Summary: A State-based Transformation System for Brain Dynamics Programming.
 Home-page: https://github.com/brainpy/brainstate
 Author: BDP
@@ -31,7 +31,7 @@ License-File: LICENSE
 Requires-Dist: jax
 Requires-Dist: jaxlib
 Requires-Dist: numpy
-Requires-Dist: brainunit >=0.0.2
+Requires-Dist: brainunit (>=0.0.2)
 Provides-Extra: cpu
 Requires-Dist: jaxlib ; extra == 'cpu'
 Provides-Extra: cuda12

{brainstate-0.0.2.post20240825.dist-info → brainstate-0.0.2.post20240826.dist-info}/RECORD

@@ -1,5 +1,5 @@
 brainstate/__init__.py,sha256=zipNSih9Tyvi4-5cXqNPGsDF7VeestkLp-lcjJ4-dA0,1408
-brainstate/_module.py,sha256=YJDp9aD38wBa_lY6BojWjWV9LJ2aFMAMYh-KZe5a4eM,52443
+brainstate/_module.py,sha256=L0cwF_6p9cSZlWGi33Mb5s4vUnbfSeQi2TPUbvsGCzo,52461
 brainstate/_module_test.py,sha256=oQaoaZBTo1o3wHrMEJTInQCc7RdcVs1gcfQGvdSb1SI,7843
 brainstate/_random_for_unit.py,sha256=1rHr7gfH_bYrJfpxbDhQUk_j00Yosx-GzyZCXrLxsd0,2007
 brainstate/_state.py,sha256=C0widCOj_ca6zfqh95jzFXf_G5vi0hJyuQ5GIqEqOUs,12102
@@ -8,10 +8,10 @@ brainstate/_utils.py,sha256=RLorgGJkt2BhbX4C-ygd-PPG0wfcGCghjSP93sRvzqM,833
 brainstate/environ.py,sha256=k0p1oyi9jbsPfuvqrPL-_zgSd7VW3LRs0LboxlaaIfc,11806
 brainstate/mixin.py,sha256=OumTTSVyYSbtudjfS_MRThsBaeVJ_0JggeMClY7xtBA,10758
 brainstate/mixin_test.py,sha256=-Ej9oUOu8O1M4oy37SVMj7xNRYhHHyAHwrjS_aISayo,2923
-brainstate/random.py,sha256=BqEBYVD9TGe8dSzp8U0suK0O4r6Ox59GCq0mwfUndVQ,188073
+brainstate/random.py,sha256=pTZvTH06hv08_TpwzAWCqAjy-8oNGmB6-Jp6MKfkLaY,188087
 brainstate/random_test.py,sha256=cCeuYvlZkCS2_RgG0vipZFNSHG8b-uJ7SXM9SZDCYQM,17866
 brainstate/surrogate.py,sha256=1kgbn82GSlpReIytIVl29yozk75gkdZv0gTBlixQ4C4,43798
-brainstate/typing.py,sha256=6BlkLSN5TiaNO49q8b0OYyzcuSxmdoG3noIJTbyhE3s,7895
+brainstate/typing.py,sha256=szCYee9R15YQfsEAQOx95_LqfrD9AYuE5dfTBTPd8sg,9165
 brainstate/util.py,sha256=y-6eX1z3EMyg6pfZt4YdDalOnJ3HDAT1IPBCJDp-gQI,19876
 brainstate/functional/__init__.py,sha256=j6-3Er4fgqWpvntzYCZVB3e5hoz-Z3aqvapITCuDri0,1107
 brainstate/functional/_activations.py,sha256=gvZ9E1-TsEUlyO7Om0eYzlM9DF-14_A32-gta1mjGo4,17798
@@ -47,20 +47,24 @@ brainstate/optim/__init__.py,sha256=1L6x_qZprw3PJYddB1nX-uTFGUl6_Qt3PM0OdY6g968,
 brainstate/optim/_lr_scheduler.py,sha256=emKnA52UVqOfUcX7LJqwP-FVDVlGGzTQi2djYmbCWUo,15627
 brainstate/optim/_lr_scheduler_test.py,sha256=OwF8Iz-PorEbO0gO--A7IIgQEytqEfYWbPucAgzqL90,1598
 brainstate/optim/_sgd_optimizer.py,sha256=JiK_AVGregL0wn8uHhRQvK9Qq7Qja7dEyLW6Aan7b70,45826
-brainstate/transform/__init__.py,sha256=my2X4ZW0uKZRfN82zyGEPizWNJ0fsSP2akvmkjn43ck,1458
+brainstate/transform/__init__.py,sha256=hqef3a4sLQ_Oihuqs8E5IghSLr9o2bS7CWmwRL8jX6E,1887
 brainstate/transform/_autograd.py,sha256=Pj_YxpU52guaxQs1NcB6qDtXgkvaPcoJbuvIF8T-Wmk,23964
 brainstate/transform/_autograd_test.py,sha256=RWriMemIF9FVFUjQh4IHzLhT9LGyd1JXpjXfFZKHn10,38654
-brainstate/transform/_control.py,sha256=0NFUGLIenqKuBhBiTmY0YgCrl2GI1ZbuWMW0DSOolpE,26874
-brainstate/transform/_controls_test.py,sha256=mPUa_qmXXVxDziAJrPWRBwsGnc3cHR9co08eJB_fJwA,7648
+brainstate/transform/_conditions.py,sha256=bdqdHCPCJIpRJxNr0enO2u81924YoIuA8kS8GUGY98g,12970
+brainstate/transform/_conditions_test.py,sha256=hg3gyOk4jn88F_ZYYqqwf6m87N3GlOUE9dC2V3BnMTA,7691
+brainstate/transform/_error_if.py,sha256=0JThfFqt9B3K3H6mS84qecBS22yTi3-FPzviaYacaMY,2597
+brainstate/transform/_error_if_test.py,sha256=kQZujlgr9bYnL-Vf7x4Zfc7jJk7rCLNVu-bsiry40dQ,1874
 brainstate/transform/_jit.py,sha256=sjQHFV8Tt75fpdl12jjPRDPT92_IZxBBJAG4gapdbNQ,11471
-brainstate/transform/_jit_error.py,sha256=8rGRx8dtvmPWmHVOsfz30EUMXSix-m2PKM3Ni_9-_7I,4829
-brainstate/transform/_jit_error_test.py,sha256=GAVGL0eNJ5Fu0lHABCGc-nLfa_0x0tw_VPfURB-nhLc,1862
 brainstate/transform/_jit_test.py,sha256=5ltT7izh_OS9dcHnRymmVhq01QomjwZGdA8XzwJRLb4,2868
+brainstate/transform/_loop_collect_return.py,sha256=aUZSK5MX4stVP8Te4R8glm2SdP18rUUfHjcV4TXOPC8,20768
+brainstate/transform/_loop_no_collection.py,sha256=p2vHoNNesDH2cM7b5LgLzSv90M8iNQPkRZEl0jhf7yA,6476
 brainstate/transform/_make_jaxpr.py,sha256=ZkrOZu4_0xcILuPUA3RFEkorJ-xbDuDtXorJI_qVThE,30450
 brainstate/transform/_make_jaxpr_test.py,sha256=K3vRUBroDTCCx0lnmhgHtgrlWvWglJO2f1K2phTvU70,3819
+brainstate/transform/_mapping.py,sha256=G9XUsD1xKLCprwwE0wv0gSXS0NYZ-ZIsv-PKKRlOoTA,3821
 brainstate/transform/_progress_bar.py,sha256=VGoRZPRBmB8ELNwLc6c7S8QhUUTvn0FY46IbBm9cuYM,3502
-brainstate-0.0.2.post20240825.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.0.2.post20240825.dist-info/METADATA,sha256=COfpxoCL7w1xGa1OFYFeANFLhAKmSioWVtmF_i2st34,3849
-brainstate-0.0.2.post20240825.dist-info/WHEEL,sha256=GUeE9LxUgRABPG7YM0jCNs9cBsAIx0YAkzCB88PMLgc,109
-brainstate-0.0.2.post20240825.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.0.2.post20240825.dist-info/RECORD,,
+brainstate/transform/_unvmap.py,sha256=8Se_23QrwDdcJpFcUnnMgD6EP-4XylbhP9K5TDhW358,3311
+brainstate-0.0.2.post20240826.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.0.2.post20240826.dist-info/METADATA,sha256=unkbYHiPHAtNHGENoSt47mA0N3cXWsCafRC8Fo2NPyk,3851
+brainstate-0.0.2.post20240826.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.0.2.post20240826.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.0.2.post20240826.dist-info/RECORD,,

{brainstate-0.0.2.post20240825.dist-info → brainstate-0.0.2.post20240826.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (73.0.1)
+Generator: bdist_wheel (0.38.4)
 Root-Is-Purelib: true
 Tag: py2-none-any
 Tag: py3-none-any