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

brainstate/nn/functional/_normalization.py

@@ -1,69 +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 __future__ import annotations
-
-from typing import Optional
-
-import jax
-import jax.numpy as jnp
-
-__all__ = [
-  'weight_standardization',
-]
-
-
-def weight_standardization(
-    w: jax.typing.ArrayLike,
-    eps: float = 1e-4,
-    gain: Optional[jax.Array] = None,
-    out_axis: int = -1,
-):
-  """
-  Scaled Weight Standardization,
-  see `Micro-Batch Training with Batch-Channel Normalization and Weight Standardization <https://paperswithcode.com/paper/weight-standardization>`_.
-
-  Parameters
-  ----------
-  w : jax.typing.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
-  -------
-  jax.typing.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 = jnp.mean(w, axis=axes, keepdims=True)
-  var = jnp.var(w, axis=axes, keepdims=True)
-  scale = jax.lax.rsqrt(jnp.maximum(var * fan_in, eps))
-  if gain is not None:
-    scale = gain * scale
-  shift = mean * scale
-  return w * scale - shift

brainstate/nn/functional/_spikes.py

@@ -1,90 +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 __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}.")
-

brainstate/nn/init/__init__.py

@@ -1,26 +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 ._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

brainstate/nn/init/_base.py

@@ -1,36 +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 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}')

brainstate/nn/init/_generic.py

@@ -1,175 +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.
-# ==============================================================================
-
-# -*- coding: utf-8 -*-
-
-import numbers
-from typing import Union, Callable, Optional, Sequence
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-from brainstate._state import State
-from ._base import to_size
-
-__all__ = [
-  'param',
-  'state',
-  'noise',
-]
-
-
-def _is_scalar(x):
-  return isinstance(x, numbers.Number)
-
-
-def param(
-    param: Union[Callable, np.ndarray, jax.Array, float, int, bool],
-    sizes: Union[int, Sequence[int]],
-    batch_size: Optional[int] = 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_size: int
-    The batch size.
-  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
-  --------
-  noise, state
-  """
-  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
-
-  if batch_size is not None:
-    sizes.insert(0, 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=0), batch_size, axis=0)
-  elif isinstance(param, State):
-    param = param
-    if batch_size is not None:
-      param = type(param)(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_size: Optional[int] = 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_size, int):
-      sizes.insert(0, batch_size)
-    return 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}')
-    if isinstance(batch_size, int):
-      batch_size = batch_size
-      data = State(jnp.repeat(jnp.expand_dims(init, axis=0), batch_size, axis=0))
-    else:
-      data = State(init)
-  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 = param(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