brainstate 0.0.2.post20241010__py2.py3-none-any.whl → 0.1.0__py2.py3-none-any.whl

brainstate/functional/_normalization.py

@@ -24,7 +24,7 @@ from brainstate._utils import set_module_as
 from brainstate.typing import ArrayLike
 
 __all__ = [
-  'weight_standardization',
+    'weight_standardization',
 ]
 
 
@@ -35,49 +35,49 @@ def weight_standardization(
     gain: Optional[jax.Array] = None,
     out_axis: int = -1,
 ) -> Union[jax.Array, u.Quantity]:
-  """
-  Scaled Weight Standardization,
-  see `Micro-Batch Training with Batch-Channel Normalization and Weight Standardization <https://paperswithcode.com/paper/weight-standardization>`_.
+    """
+    Scaled Weight Standardization,
+    see `Micro-Batch Training with Batch-Channel Normalization and Weight Standardization <https://paperswithcode.com/paper/weight-standardization>`_.
 
-  Parameters
-  ----------
-  w : ArrayLike
-      The weight tensor.
-  eps : float
-      A small value to avoid division by zero.
-  gain : Array
-      The gain function, by default None.
-  out_axis : int
-      The output axis, by default -1.
+    Parameters
+    ----------
+    w : ArrayLike
+        The weight tensor.
+    eps : float
+        A small value to avoid division by zero.
+    gain : Array
+        The gain function, by default None.
+    out_axis : int
+        The output axis, by default -1.
 
-  Returns
-  -------
-  ArrayLike
-      The scaled weight tensor.
-  """
-  if out_axis < 0:
-    out_axis = w.ndim + out_axis
-  fan_in = 1  # get the fan-in of the weight tensor
-  axes = []  # get the axes of the weight tensor
-  for i in range(w.ndim):
-    if i != out_axis:
-      fan_in *= w.shape[i]
-      axes.append(i)
-  # normalize the weight
-  mean = u.math.mean(w, axis=axes, keepdims=True)
-  var = u.math.var(w, axis=axes, keepdims=True)
+    Returns
+    -------
+    ArrayLike
+        The scaled weight tensor.
+    """
+    if out_axis < 0:
+        out_axis = w.ndim + out_axis
+    fan_in = 1  # get the fan-in of the weight tensor
+    axes = []  # get the axes of the weight tensor
+    for i in range(w.ndim):
+        if i != out_axis:
+            fan_in *= w.shape[i]
+            axes.append(i)
+    # normalize the weight
+    mean = u.math.mean(w, axis=axes, keepdims=True)
+    var = u.math.var(w, axis=axes, keepdims=True)
 
-  temp = u.math.maximum(var * fan_in, eps)
-  if isinstance(temp, u.Quantity):
-    unit = temp.unit
-    temp = temp.mantissa
-    if unit.is_unitless:
-      scale = jax.lax.rsqrt(temp)
+    temp = u.math.maximum(var * fan_in, eps)
+    if isinstance(temp, u.Quantity):
+        unit = temp.unit
+        temp = temp.mantissa
+        if unit.is_unitless:
+            scale = jax.lax.rsqrt(temp)
+        else:
+            scale = u.Quantity(jax.lax.rsqrt(temp), unit=1 / unit ** 0.5)
     else:
-      scale = u.Quantity(jax.lax.rsqrt(temp), unit=1 / unit ** 0.5)
-  else:
-    scale = jax.lax.rsqrt(temp)
-  if gain is not None:
-    scale = gain * scale
-  shift = mean * scale
-  return w * scale - shift
+        scale = jax.lax.rsqrt(temp)
+    if gain is not None:
+        scale = gain * scale
+    shift = mean * scale
+    return w * scale - shift

brainstate/functional/_others.py

@@ -23,7 +23,7 @@ import jax.numpy as jnp
 from brainstate.typing import PyTree
 
 __all__ = [
-  'clip_grad_norm',
+    'clip_grad_norm',
 ]
 
 
@@ -32,17 +32,17 @@ def clip_grad_norm(
     max_norm: float | jax.Array,
     norm_type: int | str | None = None
 ):
-  """
-  Clips gradient norm of an iterable of parameters.
-
-  The norm is computed over all gradients together, as if they were
-  concatenated into a single vector. Gradients are modified in-place.
-
-  Args:
-      grad (PyTree): an iterable of Tensors or a single Tensor that will have gradients normalized
-      max_norm (float): max norm of the gradients.
-      norm_type (int, str, None): type of the used p-norm. Can be ``'inf'`` for infinity norm.
-  """
-  norm_fn = partial(jnp.linalg.norm, ord=norm_type)
-  norm = norm_fn(jnp.asarray(jax.tree.leaves(jax.tree.map(norm_fn, grad))))
-  return jax.tree.map(lambda x: jnp.where(norm < max_norm, x, x * max_norm / (norm + 1e-6)), grad)
+    """
+    Clips gradient norm of an iterable of parameters.
+
+    The norm is computed over all gradients together, as if they were
+    concatenated into a single vector. Gradients are modified in-place.
+
+    Args:
+        grad (PyTree): an iterable of Tensors or a single Tensor that will have gradients normalized
+        max_norm (float): max norm of the gradients.
+        norm_type (int, str, None): type of the used p-norm. Can be ``'inf'`` for infinity norm.
+    """
+    norm_fn = partial(jnp.linalg.norm, ord=norm_type)
+    norm = norm_fn(jnp.asarray(jax.tree.leaves(jax.tree.map(norm_fn, grad))))
+    return jax.tree.map(lambda x: jnp.where(norm < max_norm, x, x * max_norm / (norm + 1e-6)), grad)

brainstate/functional/_spikes.py

@@ -16,74 +16,74 @@
 from __future__ import annotations
 
 __all__ = [
-  'spike_bitwise_or',
-  'spike_bitwise_and',
-  'spike_bitwise_iand',
-  'spike_bitwise_not',
-  'spike_bitwise_xor',
-  'spike_bitwise_ixor',
-  'spike_bitwise',
+    'spike_bitwise_or',
+    'spike_bitwise_and',
+    'spike_bitwise_iand',
+    'spike_bitwise_not',
+    'spike_bitwise_xor',
+    'spike_bitwise_ixor',
+    'spike_bitwise',
 ]
 
 
 def spike_bitwise_or(x, y):
-  """Bitwise OR operation for spike tensors."""
-  return x + y - x * y
+    """Bitwise OR operation for spike tensors."""
+    return x + y - x * y
 
 
 def spike_bitwise_and(x, y):
-  """Bitwise AND operation for spike tensors."""
-  return x * y
+    """Bitwise AND operation for spike tensors."""
+    return x * y
 
 
 def spike_bitwise_iand(x, y):
-  """Bitwise IAND operation for spike tensors."""
-  return (1 - x) * y
+    """Bitwise IAND operation for spike tensors."""
+    return (1 - x) * y
 
 
 def spike_bitwise_not(x):
-  """Bitwise NOT operation for spike tensors."""
-  return 1 - x
+    """Bitwise NOT operation for spike tensors."""
+    return 1 - x
 
 
 def spike_bitwise_xor(x, y):
-  """Bitwise XOR operation for spike tensors."""
-  return x + y - 2 * x * y
+    """Bitwise XOR operation for spike tensors."""
+    return x + y - 2 * x * y
 
 
 def spike_bitwise_ixor(x, y):
-  """Bitwise IXOR operation for spike tensors."""
-  return x * (1 - y) + (1 - x) * y
+    """Bitwise IXOR operation for spike tensors."""
+    return x * (1 - y) + (1 - x) * y
 
 
 def spike_bitwise(x, y, op: str):
-  r"""Bitwise operation for spike tensors.
-
-  .. math::
-
-     \begin{array}{ccc}
-      \hline \text { Mode } & \text { Expression for } \mathrm{g}(\mathrm{x}, \mathrm{y}) & \text { Code for } \mathrm{g}(\mathrm{x}, \mathrm{y}) \\
-      \hline \text { ADD } & x+y & x+y \\
-      \text { AND } & x \cap y & x \cdot y \\
-      \text { IAND } & (\neg x) \cap y & (1-x) \cdot y \\
-      \text { OR } & x \cup y & (x+y)-(x \cdot y) \\
-      \hline
-      \end{array}
-
-  Args:
-    x: A spike tensor.
-    y: A spike tensor.
-    op: A string indicating the bitwise operation to perform.
-  """
-  if op == 'or':
-    return spike_bitwise_or(x, y)
-  elif op == 'and':
-    return spike_bitwise_and(x, y)
-  elif op == 'iand':
-    return spike_bitwise_iand(x, y)
-  elif op == 'xor':
-    return spike_bitwise_xor(x, y)
-  elif op == 'ixor':
-    return spike_bitwise_ixor(x, y)
-  else:
-    raise NotImplementedError(f"Unsupported bitwise operation: {op}.")
+    r"""Bitwise operation for spike tensors.
+
+    .. math::
+
+       \begin{array}{ccc}
+        \hline \text { Mode } & \text { Expression for } \mathrm{g}(\mathrm{x}, \mathrm{y}) & \text { Code for } \mathrm{g}(\mathrm{x}, \mathrm{y}) \\
+        \hline \text { ADD } & x+y & x+y \\
+        \text { AND } & x \cap y & x \cdot y \\
+        \text { IAND } & (\neg x) \cap y & (1-x) \cdot y \\
+        \text { OR } & x \cup y & (x+y)-(x \cdot y) \\
+        \hline
+        \end{array}
+
+    Args:
+      x: A spike tensor.
+      y: A spike tensor.
+      op: A string indicating the bitwise operation to perform.
+    """
+    if op == 'or':
+        return spike_bitwise_or(x, y)
+    elif op == 'and':
+        return spike_bitwise_and(x, y)
+    elif op == 'iand':
+        return spike_bitwise_iand(x, y)
+    elif op == 'xor':
+        return spike_bitwise_xor(x, y)
+    elif op == 'ixor':
+        return spike_bitwise_ixor(x, y)
+    else:
+        raise NotImplementedError(f"Unsupported bitwise operation: {op}.")

brainstate/graph/__init__.py

@@ -0,0 +1,33 @@
+# 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 ._graph_context import *
+from ._graph_context import __all__ as _graph_context__all__
+from ._graph_convert import *
+from ._graph_convert import __all__ as _graph_convert__all__
+from ._graph_node import *
+from ._graph_node import __all__ as _graph_node__all__
+from ._graph_operation import *
+from ._graph_operation import __all__ as _graph_operation__all__
+
+__all__ = (_graph_context__all__ +
+           _graph_convert__all__ +
+           _graph_node__all__ +
+           _graph_operation__all__)
+del (_graph_context__all__,
+     _graph_convert__all__,
+     _graph_node__all__,
+     _graph_operation__all__)