brainstate 0.0.1__py2.py3-none-any.whl → 0.0.1.post20240612__py2.py3-none-any.whl

brainstate/transform/_make_jaxpr.py

@@ -54,20 +54,27 @@ function.
 from __future__ import annotations
 
 import functools
+import inspect
 import operator
 from collections.abc import Hashable, Iterable, Sequence
+from contextlib import ExitStack
 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
+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 jax.interpreters.xla import abstractify
+from jax.util import wraps
 
 from brainstate._state import State, StateTrace
 from brainstate._utils import set_module_as
 
 PyTree = Any
+AxisName = Hashable
+
 
 __all__ = [
   "StatefulFunction",
@@ -393,7 +400,8 @@ class StatefulFunction(object):
     if cache_key not in self._state_trace:
       try:
         # jaxpr
-        jaxpr, (out_shapes, state_shapes) = jax.make_jaxpr(
+        # jaxpr, (out_shapes, state_shapes) = jax.make_jaxpr(
+        jaxpr, (out_shapes, state_shapes) = _make_jaxpr(
           functools.partial(self._wrapped_fun_to_eval, cache_key),
           static_argnums=self.static_argnums,
           axis_env=self.axis_env,
@@ -474,7 +482,8 @@ def make_jaxpr(
     state_returns: Union[str, Tuple[str, ...]] = ('read', 'write')
 ) -> Callable[..., (Tuple[jax.core.ClosedJaxpr, Tuple[State, ...]] |
                     Tuple[jax.core.ClosedJaxpr, Tuple[State, ...], PyTree])]:
-  """Creates a function that produces its jaxpr given example args.
+  """
+  Creates a function that produces its jaxpr given example args.
 
   Args:
     fun: The function whose ``jaxpr`` is to be computed. Its positional
@@ -571,3 +580,156 @@ def make_jaxpr(
   if hasattr(fun, "__name__"):
     make_jaxpr_f.__name__ = f"make_jaxpr({fun.__name__})"
   return make_jaxpr_f
+
+
+def _check_callable(fun):
+  # In Python 3.10+, the only thing stopping us from supporting staticmethods
+  # is that we can't take weak references to them, which the C++ JIT requires.
+  if isinstance(fun, staticmethod):
+    raise TypeError(f"staticmethod arguments are not supported, got {fun}")
+  if not callable(fun):
+    raise TypeError(f"Expected a callable value, got {fun}")
+  if inspect.isgeneratorfunction(fun):
+    raise TypeError(f"Expected a function, got a generator function: {fun}")
+
+
+def _broadcast_prefix(
+    prefix_tree: Any,
+    full_tree: Any,
+    is_leaf: Callable[[Any], bool] | None = None
+) -> list[Any]:
+  # If prefix_tree is not a tree prefix of full_tree, this code can raise a
+  # ValueError; use prefix_errors to find disagreements and raise more precise
+  # error messages.
+  result = []
+  num_leaves = lambda t: jax.tree.structure(t).num_leaves
+  add_leaves = lambda x, subtree: result.extend([x] * num_leaves(subtree))
+  jax.tree.map(add_leaves, prefix_tree, full_tree, is_leaf=is_leaf)
+  return result
+
+
+def _flat_axes_specs(
+    abstracted_axes, *args, **kwargs
+) -> list[pe.AbstractedAxesSpec]:
+  if kwargs:
+    raise NotImplementedError
+
+  def ax_leaf(l):
+    return (isinstance(l, dict) and jax.tree_util.all_leaves(l.values()) or
+            isinstance(l, tuple) and jax.tree_util.all_leaves(l, lambda x: x is None))
+
+  return _broadcast_prefix(abstracted_axes, args, ax_leaf)
+
+
+@transformation_with_aux
+def _flatten_fun(in_tree, *args_flat):
+  py_args, py_kwargs = jax.tree.unflatten(in_tree, args_flat)
+  ans = yield py_args, py_kwargs
+  yield jax.tree.flatten(ans)
+
+
+def _make_jaxpr(
+    fun: Callable,
+    static_argnums: int | Iterable[int] = (),
+    axis_env: Sequence[tuple[AxisName, int]] | None = None,
+    return_shape: bool = False,
+    abstracted_axes: Any | None = None,
+) -> Callable[..., (jax.core.ClosedJaxpr | tuple[jax.core.ClosedJaxpr, Any])]:
+  """Creates a function that produces its jaxpr given example args.
+
+  Args:
+    fun: The function whose ``jaxpr`` is to be computed. Its positional
+      arguments and return value should be arrays, scalars, or standard Python
+      containers (tuple/list/dict) thereof.
+    static_argnums: See the :py:func:`jax.jit` docstring.
+    axis_env: Optional, a sequence of pairs where the first element is an axis
+      name and the second element is a positive integer representing the size of
+      the mapped axis with that name. This parameter is useful when lowering
+      functions that involve parallel communication collectives, and it
+      specifies the axis name/size environment that would be set up by
+      applications of :py:func:`jax.pmap`.
+    return_shape: Optional boolean, defaults to ``False``. If ``True``, the
+      wrapped function returns a pair where the first element is the
+      ``ClosedJaxpr`` representation of ``fun`` and the second element is a
+      pytree with the same structure as the output of ``fun`` and where the
+      leaves are objects with ``shape``, ``dtype``, and ``named_shape``
+      attributes representing the corresponding types of the output leaves.
+
+  Returns:
+    A wrapped version of ``fun`` that when applied to example arguments returns
+    a ``ClosedJaxpr`` representation of ``fun`` on those arguments. If the
+    argument ``return_shape`` is ``True``, then the returned function instead
+    returns a pair where the first element is the ``ClosedJaxpr``
+    representation of ``fun`` and the second element is a pytree representing
+    the structure, shape, dtypes, and named shapes of the output of ``fun``.
+
+  A ``jaxpr`` is JAX's intermediate representation for program traces. The
+  ``jaxpr`` language is based on the simply-typed first-order lambda calculus
+  with let-bindings. :py:func:`make_jaxpr` adapts a function to return its
+  ``jaxpr``, which we can inspect to understand what JAX is doing internally.
+  The ``jaxpr`` returned is a trace of ``fun`` abstracted to
+  :py:class:`ShapedArray` level. Other levels of abstraction exist internally.
+
+  We do not describe the semantics of the ``jaxpr`` language in detail here, but
+  instead give a few examples.
+
+  >>> import jax
+  >>>
+  >>> def f(x): return jax.numpy.sin(jax.numpy.cos(x))
+  >>> print(f(3.0))
+  -0.83602
+  >>> _make_jaxpr(f)(3.0)
+  { lambda ; a:f32[]. let b:f32[] = cos a; c:f32[] = sin b in (c,) }
+  >>> _make_jaxpr(jax.grad(f))(3.0)
+  { lambda ; a:f32[]. let
+      b:f32[] = cos a
+      c:f32[] = sin a
+      _:f32[] = sin b
+      d:f32[] = cos b
+      e:f32[] = mul 1.0 d
+      f:f32[] = neg e
+      g:f32[] = mul f c
+    in (g,) }
+  """
+  _check_callable(fun)
+  static_argnums = _ensure_index_tuple(static_argnums)
+
+  def _abstractify(args, kwargs):
+    flat_args, in_tree = jax.tree.flatten((args, kwargs))
+    if abstracted_axes is None:
+      return map(jax.api_util.shaped_abstractify, flat_args), in_tree, [True] * len(flat_args)
+    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)
+      return in_avals, in_tree, keep_inputs
+
+  @wraps(fun)
+  @api_boundary
+  def make_jaxpr_f(*args, **kwargs):
+    f = wrap_init(fun)
+    if static_argnums:
+      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))
+    f, out_tree = _flatten_fun(f, in_tree)
+    f = annotate(f, in_type)
+    debug_info = pe.debug_info(fun, in_tree, out_tree, True, 'make_jaxpr')
+    with ExitStack() as stack:
+      for axis_name, size in axis_env or []:
+        stack.enter_context(jax.core.extend_axis_env(axis_name, size, None))
+      jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f, debug_info=debug_info)
+    closed_jaxpr = jax.core.ClosedJaxpr(jaxpr, consts)
+    if return_shape:
+      out_avals, _ = jax.util.unzip2(out_type)
+      out_shapes_flat = [jax.ShapeDtypeStruct(a.shape, a.dtype, a.named_shape) for a in out_avals]
+      return closed_jaxpr, jax.tree.unflatten(out_tree(), out_shapes_flat)
+    return closed_jaxpr
+
+  make_jaxpr_f.__module__ = "brainstate.transform"
+  if hasattr(fun, "__qualname__"):
+    make_jaxpr_f.__qualname__ = f"make_jaxpr({fun.__qualname__})"
+  if hasattr(fun, "__name__"):
+    make_jaxpr_f.__name__ = f"make_jaxpr({fun.__name__})"
+  return make_jaxpr_f

{brainstate-0.0.1.dist-info → brainstate-0.0.1.post20240612.dist-info}/METADATA

@@ -1,8 +1,8 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.0.1
+Version: 0.0.1.post20240612
 Summary: A State-based Transformation System for Brain Dynamics Programming.
-Home-page: https://github.com/brainpy/braincore
+Home-page: https://github.com/brainpy/brainstate
 Author: BrainPy Team
 Author-email: BrainPy Team <chao.brain@qq.com>
 License: Apache-2.0 license
@@ -90,12 +90,14 @@ The official documentation is hosted on Read the Docs: [https://brainstate.readt
 
 ## See also the BDP ecosystem
 
-- [``brainpy``](https://github.com/brainpy/BrainPy): The solution for the general-purpose brain dynamics programming.
+- [``brainstate``](https://github.com/brainpy/brainstate): A ``State``-based transformation system for brain dynamics programming.
 
-- [``brainstate``](https://github.com/brainpy/brainstate): The core system for the next generation of BrainPy framework.
+- [``brainunit``](https://github.com/brainpy/brainunit): The unit system for brain dynamics programming.
 
-- [``braintools``](https://github.com/brainpy/braintools): The tools for the brain dynamics simulation and analysis.
+- [``braintaichi``](https://github.com/brainpy/braintaichi): Leveraging Taichi Lang to customize brain dynamics operators.
 
-- [``brainscale``](https://github.com/brainpy/brainscale): The scalable online learning for biological spiking neural networks.
+- [``brainscale``](https://github.com/brainpy/brainscale): The scalable online learning framework for biological neural networks.
+
+- [``braintools``](https://github.com/brainpy/braintools): The toolbox for the brain dynamics simulation, training and analysis.
 
 

{brainstate-0.0.1.dist-info → brainstate-0.0.1.post20240612.dist-info}/RECORD

@@ -1,10 +1,10 @@
 brainstate/__init__.py,sha256=3R0I9oLpIS6Z9iwmx5ODzSlZGX3MbYWvCWsFHaCiaG4,1436
-brainstate/_module.py,sha256=gwJ-G1T_Yae-YnOtY5pQ-ED06GAVco4U3jzTwJx-pTg,47013
+brainstate/_module.py,sha256=R3pBeNvqR_mEfquZU60uWj7JmopOCUciF2BgcJyA0aw,48151
 brainstate/_module_test.py,sha256=4tqtp2-j5mSoUmCITY0mVZEcXzxXCWJ_02Jdt1fxYJg,4502
-brainstate/_state.py,sha256=hnud2AZCFkRG_BOtjR3WY3pt3zpHDjAAEQVkOfc4faU,10838
+brainstate/_state.py,sha256=RWnLjMeaidxWXNAA0X-8mxj4i61j3T8w5KhugACUYhI,11422
 brainstate/_state_test.py,sha256=HDdipndRLhEHWEdTmyT1ayEBkbv6qJKykfCWKI6yJ_E,1253
 brainstate/_utils.py,sha256=RLorgGJkt2BhbX4C-ygd-PPG0wfcGCghjSP93sRvzqM,833
-brainstate/environ.py,sha256=kXPMQcBMrOJQO0Mqqjxwk6GEZMm61IEHgc5rs1q3GYY,10239
+brainstate/environ.py,sha256=RMDUACuixwk2ZTHf0UGLhcd5DCraW-l1j9T3wc2wcFc,10242
 brainstate/mixin.py,sha256=V75vjMTzYcCMlPo5wekgRZZ9o6-xN8kJocQgEliu5yI,10738
 brainstate/mixin_test.py,sha256=qDYqhHbHw3aBFW8aHQdPhov29013Eo9TJDF7RW2dapE,2919
 brainstate/random.py,sha256=Mi5i0kAsR8C-VoI8LMuIbPPr6YFzq6NBxhJ5K0w2qW4,186392
@@ -12,9 +12,10 @@ brainstate/random_test.py,sha256=cCeuYvlZkCS2_RgG0vipZFNSHG8b-uJ7SXM9SZDCYQM,178
 brainstate/surrogate.py,sha256=1kgbn82GSlpReIytIVl29yozk75gkdZv0gTBlixQ4C4,43798
 brainstate/typing.py,sha256=Ooweu7c17nYP686fyIeKNomChodSxx_OEpu8QRoB9cY,2180
 brainstate/util.py,sha256=FrBN_OZAPlWxfNK8c9Z1d-bbIa8qwMrcOsSJZJS8xOE,19878
-brainstate/functional/__init__.py,sha256=6aMwiOl2UGEd6b5SBpGaB0s2QG1iF8rRFoOfvScJ4Dg,1020
-brainstate/functional/_activations.py,sha256=8mgmctDoK_oO8D3JV-YBlLMrsg5Dry0OPcyBz_2Ur0I,18498
-brainstate/functional/_normalization.py,sha256=aLe_CsVl8gZCMTcTVmrI3mEtaLmBWqJyIHQRDgslSHw,2090
+brainstate/functional/__init__.py,sha256=Z-43coOHFAsQK0u5amlr4l0fNNPc7dVcuKXfNY4Gj_s,1107
+brainstate/functional/_activations.py,sha256=IfZ6Zy8SAwyxo166E3NmCZMUHnG_rBFAUaLTyxG5FgA,18490
+brainstate/functional/_normalization.py,sha256=IxE580waloZylZVXcpUUK4bWQdlE6oSPfafaKYfDkbg,2169
+brainstate/functional/_others.py,sha256=ifB-l82y7ZB632yLUJOEcpkRY-yOoiJ0mtDOxNilp4M,1711
 brainstate/functional/_spikes.py,sha256=uAln_Q87pr1codLxeDck3PUA9jpk7S5LifNps1kdyrU,2576
 brainstate/init/__init__.py,sha256=R1dHgub47o-WJM9QkFLc7x_Q7GsyaKKDtrRHTFPpC5g,1097
 brainstate/init/_base.py,sha256=jRTmfoUsH_315vW9YMZzyIn2KDAAsv56SplBnvOyBW0,1148
@@ -28,52 +29,43 @@ brainstate/math/_einops_parsing_test.py,sha256=JPn73yld300481J6E9cL7jHWn63Vr21VV
 brainstate/math/_einops_test.py,sha256=xj-DDTL0EsW1Obm64KCnT7eqELWjjj04Ozdwk0839Tw,13289
 brainstate/math/_misc.py,sha256=jDtREP4ojxHyj6lXcLcYLGVsLA0HFZcrs8cdlnA7aK8,7863
 brainstate/math/_misc_test.py,sha256=V41YV-RiEbukKQlzq54174cpSalOhMjaHOoVH8o82eI,2443
-brainstate/nn/__init__.py,sha256=SxzR2l8bik7hjL3UxGIrrlxZm0cnEq1pwKm1teIhSRI,2058
-brainstate/nn/_base.py,sha256=NfWWYVl4kChHZg8EwDePI_D-1eQhQRlAv7KsIeUh94A,8342
+brainstate/nn/__init__.py,sha256=YJHoI8cXKVRS8f2vUl3Zegp5wm0svMz3qo9JmQJiMQk,2162
+brainstate/nn/_base.py,sha256=lzbZpku3Q2arH6ZaAwjs6bhbV0RcFChxo2UcpnX5t84,8481
 brainstate/nn/_connections.py,sha256=GSOW2IbpJRHdPyF4nFJ2RPgO8y6SVHT1Gn-pbri9pMk,22970
-brainstate/nn/_dynamics.py,sha256=E5BmaLiaSAX-BVndlZRruOllB9xzhIU9bjeNpR1VqFQ,12211
+brainstate/nn/_dynamics.py,sha256=OeYYXv1dqjUDcCsRhZo1XS7SP2li1vlH9uhME_PE9v0,13205
 brainstate/nn/_elementwise.py,sha256=T1oCu47t11Ki7LPaL-hHk4W8bKP_Q3HLJcGngcmGK0Y,43552
+brainstate/nn/_embedding.py,sha256=WbgrIaM_14abN8zBDr0xipBOsFc8dXP2m7Z_aRLAfmU,2249
 brainstate/nn/_misc.py,sha256=Z7gdJraJ18gVNHyNOk_KmE67M3OM4z3QT4RN6al5JMc,3766
 brainstate/nn/_normalizations.py,sha256=9yVDORAEpqEkL9MYSPU4m7C4q8Qj5UNsPh9sKmIt5gQ,14329
 brainstate/nn/_others.py,sha256=8PYmCiUNzru4kmm58HY0RzCs-32dnwNFDZdTTPixaqo,4492
 brainstate/nn/_poolings.py,sha256=cNZ1PyMIaViP-_AUkEbpy3ZfHo--ib1hAhL0bEAmXIQ,45688
 brainstate/nn/_poolings_test.py,sha256=iE0NgvOIWVgwmcvP4wazhGG4RJQdU2eeagdJ1sDXIBQ,7260
-brainstate/nn/_rate_rnns.py,sha256=OxAcgrT-cD-cSf_Q2618CRaalIHfLov0eud50sryPt0,13561
-brainstate/nn/_readout.py,sha256=1sosUKBXGJ2E75crFd-SfIYBpQkiRidOWm--IZ4bLQc,3981
+brainstate/nn/_rate_rnns.py,sha256=Cebhy57UWzfwrCfq0v2qLDegmb__mXL5ht750y4aTro,14457
+brainstate/nn/_readout.py,sha256=jsQwhVnrJICKw4wFq-Du2AORPb_XXz_tZ4cURcckU-E,4240
 brainstate/nn/_synouts.py,sha256=gi3EyKlzt4UoyghwvNIr03r7YabZyl1idbq9aYG8zYM,4379
-brainstate/nn/surrogate.py,sha256=1kgbn82GSlpReIytIVl29yozk75gkdZv0gTBlixQ4C4,43798
 brainstate/nn/_projection/__init__.py,sha256=L6svNHTb8BDh2rdX2eYmcx_NdscSdKykkQbzpdCSkTA,1207
 brainstate/nn/_projection/_align_post.py,sha256=dAdNsuNf7jo8Qsh3uHXLonv9iDi3J9AnWzmFaG3b3bo,20655
 brainstate/nn/_projection/_align_pre.py,sha256=R2U6_RQ_o8y6PWXpozeWE2cx_oQ7WMhhrBR9hZtEBTs,24597
 brainstate/nn/_projection/_delta.py,sha256=KT8ySo3n_Q_7swzOH-ISDf0x9rjMkiv99H-vqeQZDR8,7122
 brainstate/nn/_projection/_utils.py,sha256=UcmELOqsINgqJr7eC5BSNNteyZ--1lyGjhUTJfxyMmA,813
 brainstate/nn/_projection/_vanilla.py,sha256=_bh_DLtF0o33SBtj6IGL8CTanFEtJwfjBrgxBEAmIlg,3397
-brainstate/nn/functional/__init__.py,sha256=6aMwiOl2UGEd6b5SBpGaB0s2QG1iF8rRFoOfvScJ4Dg,1020
-brainstate/nn/functional/_activations.py,sha256=8mgmctDoK_oO8D3JV-YBlLMrsg5Dry0OPcyBz_2Ur0I,18498
-brainstate/nn/functional/_normalization.py,sha256=aLe_CsVl8gZCMTcTVmrI3mEtaLmBWqJyIHQRDgslSHw,2090
-brainstate/nn/functional/_spikes.py,sha256=uAln_Q87pr1codLxeDck3PUA9jpk7S5LifNps1kdyrU,2576
-brainstate/nn/init/__init__.py,sha256=R1dHgub47o-WJM9QkFLc7x_Q7GsyaKKDtrRHTFPpC5g,1097
-brainstate/nn/init/_base.py,sha256=jRTmfoUsH_315vW9YMZzyIn2KDAAsv56SplBnvOyBW0,1148
-brainstate/nn/init/_generic.py,sha256=OJFS7DHYmZV0JogdsgjnUseUfvTUrAUYiXZynCQqmG4,5163
-brainstate/nn/init/_random_inits.py,sha256=STbX-mrHwNuICXkw7EldtJLdUUsWOAcGkEzx2ycV-Yc,15321
-brainstate/nn/init/_regular_inits.py,sha256=n-vF-51FM1UcUh-8h5lUk5Jhjrn04KPcGXgGhUGFAAk,3065
 brainstate/optim/__init__.py,sha256=1xH5_peSWKuZ4tOU295r9EKAv0a-cBMABx6XV3faDJI,919
 brainstate/optim/_lr_scheduler.py,sha256=emKnA52UVqOfUcX7LJqwP-FVDVlGGzTQi2djYmbCWUo,15627
-brainstate/optim/_lr_scheduler_test.py,sha256=JznLyr2BMJjsmypu8Zte8p4qTPlaWRUYK8u1QGn4ROo,1201
+brainstate/optim/_lr_scheduler_test.py,sha256=OwF8Iz-PorEbO0gO--A7IIgQEytqEfYWbPucAgzqL90,1598
 brainstate/optim/_sgd_optimizer.py,sha256=7-jMfP_Hol0XGEA6_4wVqygpLTqI1646F6eeLtwtNFY,45760
 brainstate/transform/__init__.py,sha256=9S9TLp1sF6nWRmW6jFtu6_dLmOc43V88Ruh073Z8I50,1460
 brainstate/transform/_autograd.py,sha256=sFGJ6oAhlSr54Hb1c1aNc5Q2St7eIr_X77lupc31YXg,23964
 brainstate/transform/_autograd_test.py,sha256=epQ2z97fAp_dQ_CwWGZD7sgw-p9o9fGfSeOUAJiiDY0,38658
 brainstate/transform/_control.py,sha256=NWceTIuLlj2uGTdNcqBAXgnaLuChOGgAtIXtFn5vdLU,26837
 brainstate/transform/_controls_test.py,sha256=mPUa_qmXXVxDziAJrPWRBwsGnc3cHR9co08eJB_fJwA,7648
-brainstate/transform/_jit.py,sha256=kkyPk2aBWMiFoXcDx3p_bIiaqMENMrNyyXC0dtzC2jo,10863
+brainstate/transform/_jit.py,sha256=sjQHFV8Tt75fpdl12jjPRDPT92_IZxBBJAG4gapdbNQ,11471
 brainstate/transform/_jit_error.py,sha256=lO_e5AdhkjozHjM10q0b57OaXbeZ9gQkVmZMN6VQVCw,4450
 brainstate/transform/_jit_test.py,sha256=lVXvScfXExhXwFi8jnvEY6stNVulZHCzriamajFqzrY,2891
-brainstate/transform/_make_jaxpr.py,sha256=O4auuHNEhcHK71AKA42uIgx5wtgB9Im0patR4m9qCus,23048
+brainstate/transform/_make_jaxpr.py,sha256=MTeBpPO1thu5yDytWoJijySHV7-nWmUoBMC0RCbdzcY,29972
 brainstate/transform/_make_jaxpr_test.py,sha256=4nEwZv_ebgUZgV86vOJFO_qC69mw2F3rogViF2SC1Qs,3823
 brainstate/transform/_progress_bar.py,sha256=myrAkBcUfuVGFLVwFzeSe5vdg1z49ARKqTlccG92maA,3536
-brainstate-0.0.1.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.0.1.dist-info/METADATA,sha256=LxLrWvqU_wZ5pjJZbE1XnglTOonyiMdsjw-pkZnxhxU,4103
-brainstate-0.0.1.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.0.1.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.0.1.dist-info/RECORD,,
+brainstate-0.0.1.post20240612.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.0.1.post20240612.dist-info/METADATA,sha256=VRHXnO0TBRcoo_M4iFHsywjCJbhonpStzYSmRXkR_wM,4254
+brainstate-0.0.1.post20240612.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.0.1.post20240612.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.0.1.post20240612.dist-info/RECORD,,

brainstate/nn/functional/__init__.py

@@ -1,25 +0,0 @@
-# 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 ._activations import *
-from ._activations import __all__ as __activations_all__
-from ._normalization import *
-from ._normalization import __all__ as __others_all__
-from ._spikes import *
-from ._spikes import __all__ as __spikes_all__
-
-__all__ = __spikes_all__ + __others_all__ + __activations_all__
-