braintools 0.0.2__py2.py3-none-any.whl

braintools/functional/_normalization.py

@@ -0,0 +1,62 @@
+# Copyright 2024- BrainPy 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
+from typing import Sequence, Optional
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+__all__ = [
+  'weight_standardization',
+]
+
+
+def weight_standardization(
+    w: jax.typing.ArrayLike,
+    axes: Sequence[int],
+    eps: float,
+    gain: Optional[jax.Array] = None
+):
+  """
+  Scaled Weight Standardization.
+
+  Parameters
+  ----------
+  w : jax.typing.ArrayLike
+      The weight tensor.
+  axes : Sequence[int]
+      The axes to calculate the mean and variance.
+  eps : float
+      A small value to avoid division by zero.
+  gain : Array
+      The gain function, by default None.
+
+  Returns
+  -------
+  jax.typing.ArrayLike
+      The scaled weight tensor.
+
+  """
+  # Get Scaled WS weight HWIO;
+  fan_in = np.prod(w.shape[:-1])
+  mean = jnp.mean(w, axis=axes[:-1], keepdims=True)
+  var = jnp.var(w, axis=axes[:-1], keepdims=True)
+  weight = (w - mean) / (var * fan_in + eps) ** 0.5
+  if gain is not None:
+    weight = gain * weight
+  return weight

braintools/functional/_spikes.py

@@ -0,0 +1,90 @@
+# Copyright 2024- BrainPy 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
+
+__all__ = [
+  '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
+
+
+def spike_bitwise_and(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
+
+
+def spike_bitwise_not(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
+
+
+def spike_bitwise_ixor(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}.")
+

braintools/init/__init__.py

@@ -0,0 +1,26 @@
+# Copyright 2024- BrainPy 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 ._base import *
+from ._base import __all__ as _base_all
+from ._generic import *
+from ._generic import __all__ as _generic_all
+from ._random_inits import *
+from ._random_inits import __all__ as _random_inits_all
+from ._regular_inits import *
+from ._regular_inits import __all__ as _regular_inits_all
+
+__all__ = _generic_all + _base_all + _regular_inits_all + _random_inits_all

braintools/init/_base.py

@@ -0,0 +1,36 @@
+# Copyright 2024- BrainPy 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 typing import Optional, Tuple
+
+import numpy as np
+
+__all__ = ['Initializer', 'to_size']
+
+
+class Initializer(object):
+  def __call__(self, *args, **kwargs):
+    raise NotImplementedError
+
+
+def to_size(x) -> Optional[Tuple[int]]:
+  if isinstance(x, (tuple, list)):
+    return tuple(x)
+  if isinstance(x, (int, np.integer)):
+    return (x,)
+  if x is None:
+    return x
+  raise ValueError(f'Cannot make a size for {x}')

braintools/init/_generic.py

@@ -0,0 +1,205 @@
+# Copyright 2024- BrainPy 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.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+import numbers
+from typing import Union, Callable, Optional, Sequence
+
+import braincore as bc
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from ._base import to_size
+
+__all__ = [
+  'parameter',
+  'state',
+  'noise',
+]
+
+
+def _is_scalar(x):
+  return isinstance(x, numbers.Number)
+
+
+def parameter(
+    param: Union[Callable, np.ndarray, jax.Array, float, int, bool],
+    sizes: Union[int, Sequence[int]],
+    batch_or_mode: Optional[Union[int, bool, bc.mixin.Mode]] = None,
+    allow_none: bool = True,
+    allow_scalar: bool = True,
+):
+  """Initialize parameters.
+
+  Parameters
+  ----------
+  param: callable, Initializer, bm.ndarray, jnp.ndarray, onp.ndarray, float, int, bool
+    The initialization of the parameter.
+    - If it is None, the created parameter will be None.
+    - If it is a callable function :math:`f`, the ``f(size)`` will be returned.
+    - If it is an instance of :py:class:`init.Initializer``, the ``f(size)`` will be returned.
+    - If it is a tensor, then this function check whether ``tensor.shape`` is equal to the given ``size``.
+  sizes: int, sequence of int
+    The shape of the parameter.
+  batch_or_mode: int, bool, braincore.mixin.Mode
+    The batch size or the mode.
+  allow_none: bool
+    Whether allow the parameter is None.
+  allow_scalar: bool
+    Whether allow the parameter is a scalar value.
+
+  Returns
+  -------
+  param: ArrayType, float, int, bool, None
+    The initialized parameter.
+
+  See Also
+  --------
+  variable_, noise, delay
+  """
+  if param is None:
+    if allow_none:
+      return None
+    else:
+      raise ValueError(f'Expect a parameter with type of float, ArrayType, Initializer, or '
+                       f'Callable function, but we got None. ')
+  sizes = list(to_size(sizes))
+  if allow_scalar and _is_scalar(param):
+    return param
+
+  batch_axis = None
+  batch_size = None
+  if isinstance(batch_or_mode, bc.mixin.Mode):
+    if batch_or_mode.has(bc.mixin.Batching):
+      batch_axis = batch_or_mode.batch_axis
+      batch_size = batch_or_mode.batch_size
+  elif batch_or_mode in (None, False):
+    pass
+  elif isinstance(batch_or_mode, int):
+    batch_axis = 0
+    batch_size = batch_or_mode
+  else:
+    raise ValueError('Unknown batch_or_mode.')
+  if batch_size is not None:
+    sizes.insert(batch_axis, batch_size)
+
+  if callable(param):
+    return param(sizes)
+  elif isinstance(param, (np.ndarray, jax.Array)):
+    param = jnp.asarray(param)
+    if batch_size is not None:
+      param = jnp.repeat(jnp.expand_dims(param, axis=batch_axis), batch_size, axis=batch_axis)
+  elif isinstance(param, bc.State):
+    param = param
+    if batch_size is not None:
+      param = bc.State(jnp.repeat(jnp.expand_dims(param.value, axis=batch_axis), batch_size, axis=batch_axis))
+  else:
+    raise ValueError(f'Unknown parameter type: {type(param)}')
+
+  if allow_scalar:
+    if param.shape == () or param.shape == (1,):
+      return param
+  if param.shape != tuple(sizes):
+    raise ValueError(f'The shape of the parameters should be {sizes}, but we got {param.shape}')
+  return param
+
+
+def state(
+    init: Union[Callable, np.ndarray, jax.Array],
+    sizes: Union[int, Sequence[int]] = None,
+    batch_or_mode: Optional[Union[int, bool, bc.mixin.Mode]] = None,
+):
+  """
+  Initialize a :math:`~.State` from a callable function or a data.
+  """
+  sizes = to_size(sizes)
+
+  if callable(init):
+    if sizes is None:
+      raise ValueError('"varshape" cannot be None when data is a callable function.')
+    sizes = list(sizes)
+    if isinstance(batch_or_mode, bc.mixin.Mode):
+      if batch_or_mode.has(bc.mixin.Batching):
+        sizes.insert(batch_or_mode.batch_axis, batch_or_mode.batch_size)
+    elif batch_or_mode in (None, False):
+      pass
+    elif isinstance(batch_or_mode, int):
+      sizes.insert(0, batch_or_mode)
+    else:
+      raise ValueError(f'Unknown batch_size_or_mode: {batch_or_mode}')
+    return bc.State(init(sizes))
+
+  else:
+    if sizes is not None:
+      if jnp.shape(init) != sizes:
+        raise ValueError(f'The shape of "data" {jnp.shape(init)} does not match with "var_shape" {sizes}')
+    batch_axis = None
+    batch_size = None
+    if isinstance(batch_or_mode, bc.mixin.Mode):
+      if batch_or_mode.has(bc.mixin.Batching):
+        batch_axis = batch_or_mode.batch_axis
+        batch_size = batch_or_mode.batch_size
+    elif batch_or_mode in (None, False):
+      pass
+    elif isinstance(batch_or_mode, int):
+      batch_axis = 0
+      batch_size = batch_or_mode
+    else:
+      raise ValueError('Unknown batch_size_or_mode.')
+    data = bc.State(jnp.repeat(jnp.expand_dims(init, axis=batch_axis), batch_size, axis=batch_axis))
+  return data
+
+
+def noise(
+    noises: Optional[Union[int, float, np.ndarray, jax.Array, Callable]],
+    size: Union[int, Sequence[int]],
+    num_vars: int = 1,
+    noise_idx: int = 0,
+) -> Optional[Callable]:
+  """Initialize a noise function.
+
+  Parameters
+  ----------
+  noises: Any
+  size: Shape
+    The size of the noise.
+  num_vars: int
+    The number of variables.
+  noise_idx: int
+    The index of the current noise among all noise variables.
+
+  Returns
+  -------
+  noise_func: function, None
+    The noise function.
+
+  See Also
+  --------
+  variable_, parameter, delay
+
+  """
+  if callable(noises):
+    return noises
+  elif noises is None:
+    return None
+  else:
+    noises = parameter(noises, size, allow_none=False)
+    if num_vars > 1:
+      noises_ = [None] * num_vars
+      noises_[noise_idx] = noises
+      noises = tuple(noises_)
+    return lambda *args, **kwargs: noises