brainstate 0.1.3__py2.py3-none-any.whl → 0.1.5__py2.py3-none-any.whl

brainstate/nn/_synaptic_projection.py

@@ -15,119 +15,409 @@
 # -*- coding: utf-8 -*-
 
 
-from typing import Callable, Union
+from typing import Callable, Union, Tuple
 
 import brainunit as u
 
-from brainstate._compatible_import import brainevent
-from brainstate.mixin import ParamDescriber, AlignPost, UpdateReturn
+from brainstate import init
+from brainstate._state import ParamState
+from brainstate.typing import ArrayLike
 from ._dynamics import Dynamics, Projection
-from ._projection import AlignPostProj, RawProj
-from ._stp import ShortTermPlasticity
-from ._synapse import Synapse
-from ._synouts import SynOut
 
 __all__ = [
-    'align_pre_projection',
-    'align_post_projection',
+    'SymmetryGapJunction',
+    'AsymmetryGapJunction',
 ]
 
 
-class align_pre_projection(Projection):
+class align_pre_ltp(Projection):
+    pass
+
+
+class align_post_ltp(Projection):
+    pass
+
+
+def get_gap_junction_post_key(i: int):
+    return f'gap_junction_post_{i}'
+
+
+def get_gap_junction_pre_key(i: int):
+    return f'gap_junction_pre_{i}'
+
+
+class SymmetryGapJunction(Projection):
     """
-    Represents a pre-synaptic alignment projection mechanism.
-
-    This class inherits from the `Projection` base class and is designed to
-    manage the pre-synaptic alignment process in neural network simulations.
-    It takes into account pre-synaptic dynamics, synaptic properties, delays,
-    communication functions, synaptic outputs, post-synaptic dynamics, and
-    short-term plasticity.
-
-    Attributes:
-        pre (Dynamics): The pre-synaptic dynamics object.
-        syn (Synapse): The synaptic object after pre-synaptic alignment.
-        delay (u.Quantity[u.second]): The output delay from the synapse.
-        projection (RawProj): The raw projection object handling communication,
-            output, and post-synaptic dynamics.
-        stp (ShortTermPlasticity, optional): The short-term plasticity object,
-            defaults to None.
+    Implements a symmetric electrical coupling (gap junction) between neuron populations.
+
+    This class represents electrical synapses where the conductance is identical in both
+    directions. Gap junctions allow bidirectional flow of electrical current directly between
+    neurons, with the current magnitude proportional to the voltage difference between
+    connected neurons.
+
+    Parameters
+    ----------
+    couples : Union[Tuple[Dynamics, Dynamics], Dynamics]
+        Either a single Dynamics object (when pre and post populations are the same)
+        or a tuple of two Dynamics objects (pre, post) representing the coupled neuron populations.
+    states : Union[str, Tuple[str, str]]
+        Either a single string (when pre and post states are the same)
+        or a tuple of two strings (pre_state, post_state) representing the state variables
+        to use for calculating voltage differences (typically membrane potentials).
+    conn : Callable
+        Connection function that returns pre_ids and post_ids arrays defining connections.
+    weight : Union[Callable, ArrayLike]
+        Conductance weights for the gap junctions. The same weight applies in both directions
+        of the connection.
+    param_type : type, optional
+        The parameter state type to use for weights, defaults to ParamState.
+
+    Notes
+    -----
+    The symmetric gap junction applies identical conductance in both directions between
+    connected neurons, ensuring balanced electrical coupling in the network.
+
+    See Also
+    --------
+    AsymmetryGapJunction : For gap junctions with different conductances in each direction.
     """
 
     def __init__(
         self,
-        pre: Dynamics,
-        syn: Union[Synapse, ParamDescriber[Synapse]],
-        delay: u.Quantity[u.second] | None,
-        comm: Callable,
-        out: SynOut,
-        post: Dynamics,
-        stp: ShortTermPlasticity = None,
+        couples: Union[Tuple[Dynamics, Dynamics], Dynamics],
+        states: Union[str, Tuple[str, str]],
+        conn: Callable,
+        weight: Union[Callable, ArrayLike],
+        param_type: type = ParamState
     ):
         super().__init__()
-        pre = pre
-        syn: Synapse = pre.align_pre(syn)
-        assert isinstance(syn, UpdateReturn), "Synapse must implement UpdateReturn interface"
-        # require "syn" implement the "update_return()" function
-        self.delay = syn.output_delay(delay)
-        self.projection = RawProj(comm=comm, out=out, post=post)
-        self.stp = stp
-
-    def update(self):
-        x = self.delay()
-        if self.stp is not None:
-            x = self.stp(x)
-        return self.projection(x)
-
-
-class align_post_projection(Projection):
+
+        if isinstance(states, str):
+            pre_state = states
+            post_state = states
+        else:
+            pre_state, post_state = states
+        if isinstance(couples, Dynamics):
+            pre = couples
+            post = couples
+        else:
+            pre, post = couples
+        assert isinstance(pre_state, str), "pre_state must be a string representing the pre-synaptic state"
+        assert isinstance(post_state, str), "post_state must be a string representing the post-synaptic state"
+        assert isinstance(pre, Dynamics), "pre must be a Dynamics object"
+        assert isinstance(post, Dynamics), "post must be a Dynamics object"
+        self.pre_state = pre_state
+        self.post_state = post_state
+        self.pre = pre
+        self.post = post
+        self.pre_ids, self.post_ids = conn(pre.out_size, post.out_size)
+        self.weight = param_type(init.param(weight, (len(self.pre_ids),)))
+
+    def update(self, *args, **kwargs):
+        if not hasattr(self.pre, self.pre_state):
+            raise ValueError(f"pre_state {self.pre_state} not found in pre-synaptic neuron group")
+        if not hasattr(self.post, self.post_state):
+            raise ValueError(f"post_state {self.post_state} not found in post-synaptic neuron group")
+        pre = getattr(self.pre, self.pre_state).value
+        post = getattr(self.post, self.post_state).value
+
+        return symmetry_gap_junction_projection(
+            pre=self.pre,
+            pre_value=pre,
+            post=self.post,
+            post_value=post,
+            pre_ids=self.pre_ids,
+            post_ids=self.post_ids,
+            weight=self.weight.value,
+        )
+
+
+def symmetry_gap_junction_projection(
+    pre: Dynamics,
+    pre_value: ArrayLike,
+    post: Dynamics,
+    post_value: ArrayLike,
+    pre_ids: ArrayLike,
+    post_ids: ArrayLike,
+    weight: ArrayLike,
+):
+    """
+    Calculate symmetrical electrical coupling through gap junctions between neurons.
+
+    This function implements bidirectional gap junction coupling where the same weight is
+    applied in both directions. It computes the electrical current flowing between
+    connected neurons based on their potential differences and updates both pre-synaptic
+    and post-synaptic neuron groups with appropriate input currents.
+
+    Parameters
+    ----------
+    pre : Dynamics
+        The pre-synaptic neuron group dynamics object.
+    pre_value : ArrayLike
+        State values (typically membrane potentials) of the pre-synaptic neurons.
+    post : Dynamics
+        The post-synaptic neuron group dynamics object.
+    post_value : ArrayLike
+        State values (typically membrane potentials) of the post-synaptic neurons.
+    pre_ids : ArrayLike
+        Indices of pre-synaptic neurons that form gap junctions.
+    post_ids : ArrayLike
+        Indices of post-synaptic neurons that form gap junctions,
+        where each pre_ids[i] is connected to post_ids[i].
+    weight : ArrayLike
+        Conductance weights for the gap junctions. Can be a scalar (same weight for all connections)
+        or an array with length matching pre_ids.
+
+    Returns
+    -------
+    ArrayLike
+        The input currents that were added to the pre-synaptic neuron group.
+
+    Notes
+    -----
+    The electrical coupling is implemented as I = g(V_pre - V_post), where:
+    - I is the current flowing from pre to post neuron
+    - g is the gap junction conductance (weight)
+    - V_pre and V_post are the membrane potentials of connected neurons
+
+    Equal but opposite currents are applied to both connected neurons, ensuring
+    conservation of current in the network.
+
+    Raises
+    ------
+    AssertionError
+        If weight dimensionality is incorrect or pre_ids and post_ids have different lengths.
+    """
+    assert u.math.ndim(weight) == 0 or weight.shape[0] == len(pre_ids), \
+        "weight must be a scalar or have the same length as pre_ids"
+    assert len(pre_ids) == len(post_ids), "pre_ids and post_ids must have the same length"
+    # Calculate the voltage difference between connected pre-synaptic and post-synaptic neurons
+    # and multiply by the connection weights
+    diff = (pre_value[pre_ids] - post_value[post_ids]) * weight
+
+    # add to post-synaptic neuron group
+    # Initialize the input currents for the post-synaptic neuron group
+    inputs = u.math.zeros(post.out_size, unit=u.get_unit(diff))
+    # Add the calculated current to the corresponding post-synaptic neurons
+    inputs = inputs.at[post_ids].add(diff)
+    # Generate a unique key for the post-synaptic input currents
+    key = get_gap_junction_post_key(0 if post.current_inputs is None else len(post.current_inputs))
+    # Add the input currents to the post-synaptic neuron group
+    post.add_current_input(key, inputs)
+
+    # add to pre-synaptic neuron group
+    # Initialize the input currents for the pre-synaptic neuron group
+    inputs = u.math.zeros(pre.out_size, unit=u.get_unit(diff))
+    # Add the calculated current to the corresponding pre-synaptic neurons
+    inputs = inputs.at[pre_ids].add(diff)
+    # Generate a unique key for the pre-synaptic input currents
+    key = get_gap_junction_pre_key(0 if pre.current_inputs is None else len(pre.current_inputs))
+    # Add the input currents to the pre-synaptic neuron group with opposite polarity
+    pre.add_current_input(key, -inputs)
+    return inputs
+
+
+class AsymmetryGapJunction(Projection):
     """
-    Represents a post-synaptic alignment projection mechanism.
+    Implements an asymmetric electrical coupling (gap junction) between neuron populations.
+
+    This class represents electrical synapses where the conductance in one direction can differ
+    from the conductance in the opposite direction. Unlike chemical synapses, gap junctions
+    allow bidirectional flow of electrical current directly between neurons, with the current
+    magnitude proportional to the voltage difference between connected neurons.
 
-    This class inherits from the `Projection` base class and is designed to
-    manage the post-synaptic alignment process in neural network simulations.
-    It takes into account spike generators, communication functions, synaptic
-    properties, synaptic outputs, post-synaptic dynamics, and short-term plasticity.
+    Parameters
+    ----------
+    pre : Dynamics
+        The pre-synaptic neuron group dynamics object.
+    pre_state : str
+        Name of the state variable in pre-synaptic neurons (typically membrane potential).
+    post : Dynamics
+        The post-synaptic neuron group dynamics object.
+    post_state : str
+        Name of the state variable in post-synaptic neurons (typically membrane potential).
+    conn : Callable
+        Connection function that returns pre_ids and post_ids arrays defining connections.
+    weight : Union[Callable, ArrayLike]
+        Conductance weights for the gap junctions. Must have shape [..., 2] where the last
+        dimension contains [pre_weight, post_weight] for each connection, defining
+        potentially different conductances in each direction.
+    param_type : type, optional
+        The parameter state type to use for weights, defaults to ParamState.
 
-    Args:
-        *spike_generator: Callable(s) that generate spike events or transform input spikes.
-        comm (Callable): Communication function for the projection.
-        syn (Union[AlignPost, ParamDescriber[AlignPost]]): The post-synaptic alignment object or its parameter describer.
-        out (Union[SynOut, ParamDescriber[SynOut]]): The synaptic output object or its parameter describer.
-        post (Dynamics): The post-synaptic dynamics object.
-        stp (ShortTermPlasticity, optional): The short-term plasticity object, defaults to None.
+    Examples
+    --------
+    >>> import brainstate
+    >>> import brainunit as u
+    >>> import numpy as np
+    >>>
+    >>> # Create two neuron populations
+    >>> n_neurons = 100
+    >>> pre_pop = brainstate.nn.LIF(n_neurons, V_rest=-70*u.mV, V_threshold=-50*u.mV)
+    >>> post_pop = brainstate.nn.LIF(n_neurons, V_rest=-70*u.mV, V_threshold=-50*u.mV)
+    >>> pre_pop.init_state()
+    >>> post_pop.init_state()
+    >>>
+    >>> # Create asymmetric gap junction with different weights in each direction
+    >>> weights = np.ones((n_neurons, 2)) * u.nS
+    >>> weights[:, 0] *= 2.0  # Double weight in pre->post direction
+    >>>
+    >>> gap_junction = brainstate.nn.AsymmetryGapJunction(
+    ...     pre=pre_pop,
+    ...     pre_state='V',
+    ...     post=post_pop,
+    ...     post_state='V',
+    ...     conn=one_to_one,
+    ...     weight=weights
+    ... )
 
+    Notes
+    -----
+    The asymmetric gap junction allows for different conductances in each direction between
+    the same pair of neurons. This can model rectifying electrical synapses that preferentially
+    allow current to flow in one direction.
+
+    See Also
+    --------
+    SymmetryGapJunction : For gap junctions with identical conductance in both directions.
     """
+
     def __init__(
         self,
-        *spike_generator,
-        comm: Callable,
-        syn: Union[AlignPost, ParamDescriber[AlignPost]],
-        out: Union[SynOut, ParamDescriber[SynOut]],
+        pre: Dynamics,
+        pre_state: str,
         post: Dynamics,
-        stp: ShortTermPlasticity = None,
+        post_state: str,
+        conn: Callable,
+        weight: Union[Callable, ArrayLike],
+        param_type: type = ParamState
     ):
         super().__init__()
-        self.spike_generator = spike_generator
-        self.projection = AlignPostProj(comm=comm, syn=syn, out=out, post=post)
-        self.stp = stp
-
-    def update(self, *x):
-        for fun in self.spike_generator:
-            x = fun(*x)
-            if isinstance(x, (tuple, list)):
-                x = tuple(x)
-            else:
-                x = (x,)
-        assert len(x) == 1, "Spike generator must return a single value or a tuple/list of values"
-        x = brainevent.BinaryArray(x[0])  # Ensure input is a BinaryFloat for spike generation
-        if self.stp is not None:
-            x = brainevent.MaskedFloat(self.stp(x))  # Ensure STP output is a MaskedFloat
-        return self.projection(x)
 
+        assert isinstance(pre_state, str), "pre_state must be a string representing the pre-synaptic state"
+        assert isinstance(post_state, str), "post_state must be a string representing the post-synaptic state"
+        self.pre_state = pre_state
+        self.post_state = post_state
+        self.pre = pre
+        self.post = post
+        self.pre_ids, self.post_ids = conn(pre.out_size, post.out_size)
+        self.weight = param_type(init.param(weight, (len(self.pre_ids), 2)))
 
-class align_pre_ltp(Projection):
-    pass
+    def update(self, *args, **kwargs):
+        if not hasattr(self.pre, self.pre_state):
+            raise ValueError(f"pre_state {self.pre_state} not found in pre-synaptic neuron group")
+        if not hasattr(self.post, self.post_state):
+            raise ValueError(f"post_state {self.post_state} not found in post-synaptic neuron group")
+        pre = getattr(self.pre, self.pre_state).value
+        post = getattr(self.post, self.post_state).value
 
+        return asymmetry_gap_junction_projection(
+            pre=self.pre,
+            pre_value=pre,
+            post=self.post,
+            post_value=post,
+            pre_ids=self.pre_ids,
+            post_ids=self.post_ids,
+            weight=self.weight.value,
+        )
 
-class align_post_ltp(Projection):
-    pass
+
+def asymmetry_gap_junction_projection(
+    pre: Dynamics,
+    pre_value: ArrayLike,
+    post: Dynamics,
+    post_value: ArrayLike,
+    pre_ids: ArrayLike,
+    post_ids: ArrayLike,
+    weight: ArrayLike,
+):
+    """
+    Calculate asymmetrical electrical coupling through gap junctions between neurons.
+
+    This function implements bidirectional gap junction coupling where different weights
+    can be applied in each direction. It computes the electrical current flowing between
+    connected neurons based on their potential differences and updates both pre-synaptic
+    and post-synaptic neuron groups with appropriate input currents.
+
+    Parameters
+    ----------
+    pre : Dynamics
+        The pre-synaptic neuron group dynamics object.
+    pre_value : ArrayLike
+        State values (typically membrane potentials) of the pre-synaptic neurons.
+    post : Dynamics
+        The post-synaptic neuron group dynamics object.
+    post_value : ArrayLike
+        State values (typically membrane potentials) of the post-synaptic neurons.
+    pre_ids : ArrayLike
+        Indices of pre-synaptic neurons that form gap junctions.
+    post_ids : ArrayLike
+        Indices of post-synaptic neurons that form gap junctions,
+        where each pre_ids[i] is connected to post_ids[i].
+    weight : ArrayLike
+        Conductance weights for the gap junctions. Must have shape [..., 2], where
+        the last dimension contains [pre_weight, post_weight] for each connection.
+        Can be a 1D array [pre_weight, post_weight] (same weights for all connections)
+        or a 2D array with shape [len(pre_ids), 2] for connection-specific weights.
+
+    Returns
+    -------
+    ArrayLike
+        The input currents that were added to the pre-synaptic neuron group.
+
+    Notes
+    -----
+    The electrical coupling is implemented with direction-specific conductances:
+    - I_pre2post = g_pre * (V_pre - V_post) flowing from pre to post neuron
+    - I_post2pre = g_post * (V_pre - V_post) flowing from post to pre neuron
+    where g_pre and g_post can be different, allowing for asymmetrical coupling.
+
+    Raises
+    ------
+    AssertionError
+        If weight dimensionality is incorrect or pre_ids and post_ids have different lengths.
+    ValueError
+        If weight shape is incompatible with asymmetrical gap junction requirements.
+    """
+    assert weight.shape[-1] == 2, 'weight must be a 2-element array for asymmetry gap junctions'
+    assert len(pre_ids) == len(post_ids), "pre_ids and post_ids must have the same length"
+    if u.math.ndim(weight) == 1:
+        # If weight is a 1D array, it should have two elements for pre and post weights
+        assert weight.shape[0] == 2, "weight must be a 2-element array for asymmetry gap junctions"
+        pre_weight = weight[0]
+        post_weight = weight[1]
+    elif u.math.ndim(weight) == 2:
+        # If weight is a 2D array, it should have two rows for pre and post weights
+        pre_weight = weight[:, 0]
+        post_weight = weight[:, 1]
+        assert pre_weight.shape[0] == len(pre_ids), "pre_weight must have the same length as pre_ids"
+        assert post_weight.shape[0] == len(post_ids), "post_weight must have the same length as post_ids"
+    else:
+        raise ValueError("weight must be a 1D or 2D array for asymmetry gap junctions")
+
+    # Calculate the voltage difference between connected pre-synaptic and post-synaptic neurons
+    # and multiply by the connection weights
+    diff = pre_value[pre_ids] - post_value[post_ids]
+    pre2post_current = diff * pre_weight
+    post2pre_current = diff * post_weight
+
+    # add to post-synaptic neuron group
+    # Initialize the input currents for the post-synaptic neuron group
+    inputs = u.math.zeros(post.out_size, unit=u.get_unit(pre2post_current))
+    # Add the calculated current to the corresponding post-synaptic neurons
+    inputs = inputs.at[post_ids].add(pre2post_current)
+    # Generate a unique key for the post-synaptic input currents
+    key = get_gap_junction_post_key(0 if post.current_inputs is None else len(post.current_inputs))
+    # Add the input currents to the post-synaptic neuron group
+    post.add_current_input(key, inputs)
+
+    # add to pre-synaptic neuron group
+    # Initialize the input currents for the pre-synaptic neuron group
+    inputs = u.math.zeros(pre.out_size, unit=u.get_unit(post2pre_current))
+    # Add the calculated current to the corresponding pre-synaptic neurons
+    inputs = inputs.at[pre_ids].add(post2pre_current)
+    # Generate a unique key for the pre-synaptic input currents
+    key = get_gap_junction_pre_key(0 if pre.current_inputs is None else len(pre.current_inputs))
+    # Add the input currents to the pre-synaptic neuron group with opposite polarity
+    pre.add_current_input(key, -inputs)
+    return inputs

brainstate/random/_rand_state.py

@@ -384,7 +384,10 @@ class RandomState(State):
         loc = _check_py_seq(loc)
         scale = _check_py_seq(scale)
         if size is None:
-            size = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
+            size = lax.broadcast_shapes(
+                jnp.shape(loc) if loc is not None else (),
+                jnp.shape(scale) if scale is not None else ()
+            )
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
         r = _loc_scale(loc, scale, jr.logistic(key, shape=_size2shape(size), dtype=dtype))
@@ -399,7 +402,10 @@ class RandomState(State):
         loc = _check_py_seq(loc)
         scale = _check_py_seq(scale)
         if size is None:
-            size = lax.broadcast_shapes(jnp.shape(scale), jnp.shape(loc))
+            size = lax.broadcast_shapes(
+                jnp.shape(scale) if scale is not None else (),
+                jnp.shape(loc) if loc is not None else ()
+            )
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
         r = _loc_scale(loc, scale, jr.normal(key, shape=_size2shape(size), dtype=dtype))
@@ -456,7 +462,7 @@ class RandomState(State):
                        dtype: DTypeLike = None):
         shape = _check_py_seq(shape)
         if size is None:
-            size = jnp.shape(shape)
+            size = jnp.shape(shape) if shape is not None else ()
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
         r = jr.gamma(key, a=shape, shape=_size2shape(size), dtype=dtype)
@@ -477,7 +483,7 @@ class RandomState(State):
                    dtype: DTypeLike = None):
         df = _check_py_seq(df)
         if size is None:
-            size = jnp.shape(size)
+            size = jnp.shape(size) if size is not None else ()
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
         r = jr.t(key, df=df, shape=_size2shape(size), dtype=dtype)
@@ -606,8 +612,8 @@ class RandomState(State):
 
         if size is None:
             size = jnp.broadcast_shapes(
-                jnp.shape(mean),
-                jnp.shape(sigma)
+                jnp.shape(mean) if mean is not None else (),
+                jnp.shape(sigma) if sigma is not None else ()
             )
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
@@ -822,7 +828,7 @@ class RandomState(State):
         a = _check_py_seq(a)
         scale = _check_py_seq(scale)
         if size is None:
-            size = jnp.broadcast_shapes(jnp.shape(a), jnp.shape(scale))
+            size = jnp.broadcast_shapes(jnp.shape(a), jnp.shape(scale) if scale is not None else ())
         else:
             if jnp.size(a) > 1:
                 raise ValueError(f'"a" should be a scalar when "size" is provided. But we got {a}')

brainstate/surrogate.py

@@ -21,7 +21,7 @@ import jax.scipy as sci
 from jax.interpreters import batching, ad, mlir
 
 from brainstate._compatible_import import Primitive
-from brainstate.util._pretty_pytree import PrettyObject
+from brainstate.util import PrettyObject
 
 __all__ = [
     'Surrogate',

brainstate/typing.py

@@ -257,7 +257,7 @@ f. A structure can end with a `...`, to denote that the PyTree must be a prefix
     cases, all named pieces must already have been seen and their structures bound.
 """  # noqa: E501
 
-Size = Union[int, Sequence[int]]
+Size = Union[int, Sequence[int], np.integer, Sequence[np.integer]]
 Axes = Union[int, Sequence[int]]
 SeedOrKey = Union[int, jax.Array, np.ndarray]
 Shape = Sequence[int]

brainstate/util/__init__.py

@@ -14,20 +14,20 @@
 # ==============================================================================
 
 from . import filter
-from ._error import *
-from ._error import __all__ as _error_all
-from ._others import *
-from ._others import __all__ as _others_all
-from ._pretty_pytree import *
-from ._pretty_pytree import __all__ as _mapping_all
-from ._pretty_repr import *
-from ._pretty_repr import __all__ as _pretty_repr_all
-from ._pretty_table import *
-from ._pretty_table import __all__ as _table_all
-from ._scaling import *
-from ._scaling import __all__ as _mem_scale_all
-from ._struct import *
-from ._struct import __all__ as _struct_all
+from .error import *
+from .error import __all__ as _error_all
+from .others import *
+from .others import __all__ as _others_all
+from .pretty_pytree import *
+from .pretty_pytree import __all__ as _mapping_all
+from .pretty_repr import *
+from .pretty_repr import __all__ as _pretty_repr_all
+from .pretty_table import *
+from .pretty_table import __all__ as _table_all
+from .scaling import *
+from .scaling import __all__ as _mem_scale_all
+from .struct import *
+from .struct import __all__ as _struct_all
 
 __all__ = (
     ['filter']

brainstate/util/{_pretty_pytree.py → pretty_pytree.py}

@@ -21,8 +21,8 @@ from typing import TypeVar, Hashable, Union, Iterable, Any, Optional, Tuple, Dic
 import jax
 
 from brainstate.typing import Filter, PathParts
-from ._pretty_repr import PrettyRepr, PrettyType, PrettyAttr, yield_unique_pretty_repr_items, pretty_repr_object
-from ._struct import dataclass
+from .pretty_repr import PrettyRepr, PrettyType, PrettyAttr, yield_unique_pretty_repr_items, pretty_repr_object
+from .struct import dataclass
 from .filter import to_predicate
 
 __all__ = [

{brainstate-0.1.3.dist-info → brainstate-0.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.3
+Version: 0.1.5
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers