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

brainstate/environ.py

@@ -18,12 +18,12 @@ from .util import MemScaling, IdMemScaling
 __all__ = [
   'set', 'context', 'get', 'all',
   'set_host_device_count', 'set_platform',
-  'get_host_device_count', 'get_platform', 'get_dt', 'get_mode', 'get_mem_scaling', 'get_precision',
+  'get_host_device_count', 'get_platform',
+  'get_dt', 'get_mode', 'get_mem_scaling', 'get_precision',
   'tolerance',
   'dftype', 'ditype', 'dutype', 'dctype',
 ]
 
-
 # Default, there are several shared arguments in the global context.
 I = 'i'  # the index of the current computation.
 T = 't'  # the current time of the current computation.

brainstate/functional/__init__.py

@@ -18,8 +18,9 @@ from ._activations import *
 from ._activations import __all__ as __activations_all__
 from ._normalization import *
 from ._normalization import __all__ as __others_all__
+from ._others import *
+from ._others import __all__ as __others_all__
 from ._spikes import *
 from ._spikes import __all__ as __spikes_all__
 
-__all__ = __spikes_all__ + __others_all__ + __activations_all__
-
+__all__ = __spikes_all__ + __others_all__ + __activations_all__ + __others_all__

brainstate/functional/_activations.py

@@ -27,7 +27,7 @@ import jax.numpy as jnp
 from jax.scipy.special import logsumexp
 from jax.typing import ArrayLike
 
-from brainstate import math, random
+from .. import random
 
 __all__ = [
   "tanh",
@@ -136,10 +136,7 @@ def prelu(x, a=0.25):
   parameter :math:`a` across all input channels. If called with `nn.PReLU(nChannels)`,
   a separate :math:`a` is used for each input channel.
   """
-  dtype = math.get_dtype(x)
-  return jnp.where(x >= jnp.asarray(0., dtype),
-                   x,
-                   jnp.asarray(a, dtype) * x)
+  return jnp.where(x >= 0., x, a * x)
 
 
 def soft_shrink(x, lambd=0.5):
@@ -161,11 +158,7 @@ def soft_shrink(x, lambd=0.5):
       - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
       - Output: :math:`(*)`, same shape as the input.
   """
-  dtype = math.get_dtype(x)
-  lambd = jnp.asarray(lambd, dtype)
-  return jnp.where(x > lambd,
-                   x - lambd,
-                   jnp.where(x < -lambd, x + lambd, jnp.asarray(0., dtype)))
+  return jnp.where(x > lambd, x - lambd, jnp.where(x < -lambd, x + lambd, 0.))
 
 
 def mish(x):
@@ -217,9 +210,8 @@ def rrelu(x, lower=0.125, upper=0.3333333333333333):
   .. _`Empirical Evaluation of Rectified Activations in Convolutional Network`:
       https://arxiv.org/abs/1505.00853
   """
-  dtype = math.get_dtype(x)
-  a = random.uniform(lower, upper, size=jnp.shape(x), dtype=dtype)
-  return jnp.where(x >= jnp.asarray(0., dtype), x, jnp.asarray(a, dtype) * x)
+  a = random.uniform(lower, upper, size=jnp.shape(x), dtype=x.dtype)
+  return jnp.where(x >= 0., x, a * x)
 
 
 def hard_shrink(x, lambd=0.5):
@@ -243,11 +235,7 @@ def hard_shrink(x, lambd=0.5):
       - Output: :math:`(*)`, same shape as the input.
 
   """
-  dtype = math.get_dtype(x)
-  lambd = jnp.asarray(lambd, dtype)
-  return jnp.where(x > lambd,
-                   x,
-                   jnp.where(x < -lambd, x, jnp.asarray(0., dtype)))
+  return jnp.where(x > lambd, x, jnp.where(x < -lambd, x, 0.))
 
 
 def relu(x: ArrayLike) -> jax.Array:
@@ -298,8 +286,7 @@ def squareplus(x: ArrayLike, b: ArrayLike = 4) -> jax.Array:
     x : input array
     b : smoothness parameter
   """
-  dtype = math.get_dtype(x)
-  return jax.nn.squareplus(x, jnp.asarray(b, dtype))
+  return jax.nn.squareplus(x, b)
 
 
 def softplus(x: ArrayLike) -> jax.Array:
@@ -417,8 +404,6 @@ def elu(x: ArrayLike, alpha: ArrayLike = 1.0) -> jax.Array:
   See also:
     :func:`selu`
   """
-  dtype = math.get_dtype(x)
-  alpha = jnp.asarray(alpha, dtype)
   return jax.nn.elu(x, alpha)
 
 
@@ -445,8 +430,6 @@ def leaky_relu(x: ArrayLike, negative_slope: ArrayLike = 1e-2) -> jax.Array:
   See also:
     :func:`relu`
   """
-  dtype = math.get_dtype(x)
-  negative_slope = jnp.asarray(negative_slope, dtype)
   return jax.nn.leaky_relu(x, negative_slope=negative_slope)
 
 
@@ -493,8 +476,6 @@ def celu(x: ArrayLike, alpha: ArrayLike = 1.0) -> jax.Array:
   Returns:
     An array.
   """
-  dtype = math.get_dtype(x)
-  alpha = jnp.asarray(alpha, dtype)
   return jax.nn.celu(x, alpha)
 
 

brainstate/functional/_normalization.py

@@ -20,11 +20,14 @@ from typing import Optional
 import jax
 import jax.numpy as jnp
 
+from .._utils import set_module_as
+
 __all__ = [
   'weight_standardization',
 ]
 
 
+@set_module_as('brainstate.functional')
 def weight_standardization(
     w: jax.typing.ArrayLike,
     eps: float = 1e-4,

brainstate/functional/_others.py

@@ -0,0 +1,49 @@
+# 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 functools import partial
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+
+PyTree = Any
+
+__all__ = [
+  'clip_grad_norm',
+]
+
+
+def clip_grad_norm(
+    grad: PyTree,
+    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)

brainstate/functional/_spikes.py

@@ -87,4 +87,3 @@ def spike_bitwise(x, y, op: str):
     return spike_bitwise_ixor(x, y)
   else:
     raise NotImplementedError(f"Unsupported bitwise operation: {op}.")
-

brainstate/mixin.py

@@ -68,7 +68,7 @@ class DelayedInit(Mixin):
   Note this Mixin can be applied in any Python object.
   """
 
-  non_hash_params: Optional[Sequence[str]] = None
+  non_hashable_params: Optional[Sequence[str]] = None
 
   @classmethod
   def delayed(cls, *args, **kwargs) -> 'DelayedInitializer':
@@ -94,7 +94,7 @@ class DelayedInitializer(metaclass=NoSubclassMeta):
   """
 
   def __init__(self, cls: T, *desc_tuple, **desc_dict):
-    self.cls = cls
+    self.cls: type = cls
 
     # arguments
     self.args = desc_tuple

brainstate/nn/__init__.py

@@ -21,6 +21,8 @@ from ._dynamics import *
 from ._dynamics import __all__ as dynamics_all
 from ._elementwise import *
 from ._elementwise import __all__ as elementwise_all
+from ._embedding import *
+from ._embedding import __all__ as embed_all
 from ._misc import *
 from ._misc import __all__ as _misc_all
 from ._normalizations import *
@@ -43,6 +45,7 @@ __all__ = (
     connections_all +
     dynamics_all +
     elementwise_all +
+    embed_all +
     normalizations_all +
     others_all +
     poolings_all +
@@ -58,6 +61,7 @@ del (
   connections_all,
   dynamics_all,
   elementwise_all,
+  embed_all,
   normalizations_all,
   others_all,
   poolings_all,

brainstate/nn/_base.py

@@ -55,22 +55,24 @@ class ExplicitInOutSize(Mixin):
 
   @property
   def in_size(self) -> Tuple[int, ...]:
-    if self._in_size is None:
-      raise ValueError(f"The input shape is not set in this node: {self} ")
     return self._in_size
 
   @in_size.setter
-  def in_size(self, in_size: Sequence[int]):
+  def in_size(self, in_size: Sequence[int] | int):
+    if isinstance(in_size, int):
+      in_size = (in_size,)
+    assert isinstance(in_size, (tuple, list)), f"Invalid type of in_size: {type(in_size)}"
     self._in_size = tuple(in_size)
 
   @property
   def out_size(self) -> Tuple[int, ...]:
-    if self._out_size is None:
-      raise ValueError(f"The output shape is not set in this node: {self}")
     return self._out_size
 
   @out_size.setter
-  def out_size(self, out_size: Sequence[int]):
+  def out_size(self, out_size: Sequence[int] | int):
+    if isinstance(out_size, int):
+      out_size = (out_size,)
+    assert isinstance(out_size, (tuple, list)), f"Invalid type of out_size: {type(out_size)}"
     self._out_size = tuple(out_size)
 
 
@@ -152,7 +154,8 @@ class Sequential(Module, UpdateReturn, Container, ExplicitInOutSize):
     self.children = visible_module_dict(self.format_elements(object, first, *tuple_modules, **dict_modules))
 
     # the input and output shape
-    self.in_size = tuple(first.in_size)
+    if first.in_size is not None:
+      self.in_size = first.in_size
     self.out_size = tuple(in_size)
 
   def _format_module(self, module, in_size):

brainstate/nn/_dynamics.py

@@ -103,6 +103,9 @@ class IF(Neuron):
   def init_state(self, batch_size: int = None, **kwargs):
     self.V = ShortTermState(init.param(jnp.zeros, self.varshape, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.V.value = init.param(jnp.zeros, self.varshape, batch_size)
+
   def get_spike(self, V=None):
     V = self.V.value if V is None else V
     v_scaled = (V - self.V_th) / self.V_th
@@ -160,6 +163,9 @@ class LIF(Neuron):
   def init_state(self, batch_size: int = None, **kwargs):
     self.V = ShortTermState(init.param(init.Constant(self.V_reset), self.varshape, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.V.value = init.param(init.Constant(self.V_reset), self.varshape, batch_size)
+
   def get_spike(self, V=None):
     V = self.V.value if V is None else V
     v_scaled = (V - self.V_th) / self.V_th
@@ -214,6 +220,10 @@ class ALIF(Neuron):
     self.V = ShortTermState(init.param(init.Constant(0.), self.varshape, batch_size))
     self.a = ShortTermState(init.param(init.Constant(0.), self.varshape, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.V.value = init.param(init.Constant(0.), self.varshape, batch_size)
+    self.a.value = init.param(init.Constant(0.), self.varshape, batch_size)
+
   def get_spike(self, V=None, a=None):
     V = self.V.value if V is None else V
     a = self.a.value if a is None else a
@@ -275,6 +285,9 @@ class Expon(Synapse):
   def init_state(self, batch_size: int = None, **kwargs):
     self.g = ShortTermState(init.param(init.Constant(0.), self.varshape, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.g.value = init.param(init.Constant(0.), self.varshape, batch_size)
+
   def update(self, x=None):
     self.g.value = exp_euler_step(self.dg, self.g.value, environ.get('t'))
     if x is not None:
@@ -325,6 +338,10 @@ class STP(Synapse):
     self.x = ShortTermState(init.param(init.Constant(1.), self.varshape, batch_size))
     self.u = ShortTermState(init.param(init.Constant(self.U), self.varshape, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
+    self.u.value = init.param(init.Constant(self.U), self.varshape, batch_size)
+
   def du(self, u, t):
     return self.U - u / self.tau_f
 
@@ -390,6 +407,9 @@ class STD(Synapse):
   def init_state(self, batch_size: int = None, **kwargs):
     self.x = ShortTermState(init.param(init.Constant(1.), self.varshape, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
+
   def update(self, pre_spike):
     t = environ.get('t')
     x = exp_euler_step(self.dx, self.x.value, t)

brainstate/nn/_elementwise.py

@@ -19,11 +19,12 @@ from __future__ import annotations
 
 from typing import Optional
 
+import brainunit as bu
 import jax.numpy as jnp
 import jax.typing
 
 from ._base import ElementWiseBlock
-from .. import math, environ, random, functional as F
+from .. import environ, random, functional as F
 from .._module import Module
 from .._state import ParamState
 from ..mixin import Mode
@@ -82,7 +83,7 @@ class Threshold(Module, ElementWiseBlock):
     self.value = value
 
   def __call__(self, x: ArrayLike) -> ArrayLike:
-    dtype = math.get_dtype(x)
+    dtype = bu.math.get_dtype(x)
     return jnp.where(x > jnp.asarray(self.threshold, dtype=dtype),
                      x,
                      jnp.asarray(self.value, dtype=dtype))
@@ -1142,7 +1143,7 @@ class Dropout(Module, ElementWiseBlock):
     self.prob = prob
 
   def __call__(self, x):
-    dtype = math.get_dtype(x)
+    dtype = bu.math.get_dtype(x)
     fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
     if fit_phase:
       keep_mask = random.bernoulli(self.prob, x.shape)
@@ -1172,7 +1173,7 @@ class _DropoutNd(Module, ElementWiseBlock):
     self.channel_axis = channel_axis
 
   def __call__(self, x):
-    dtype = math.get_dtype(x)
+    dtype = bu.math.get_dtype(x)
     # get fit phase
     fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
 

brainstate/nn/_embedding.py

@@ -0,0 +1,66 @@
+# 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, Callable, Union
+
+from ._base import DnnLayer
+from .. import init
+from .._state import ParamState
+from ..mixin import Mode, Training
+from ..typing import ArrayLike
+
+__all__ = [
+  'Embedding',
+]
+
+
+class Embedding(DnnLayer):
+  r"""
+  A simple lookup table that stores embeddings of a fixed size.
+
+  Args:
+    num_embeddings: Size of embedding dictionary. Must be non-negative.
+    embedding_size: Size of each embedding vector. Must be non-negative.
+    embed_init: The initializer for the embedding lookup table, of shape `(num_embeddings, embedding_size)`.
+
+  """
+
+  def __init__(
+      self,
+      num_embeddings: int,
+      embedding_size: int,
+      embed_init: Union[Callable, ArrayLike] = init.LecunUniform(),
+      name: Optional[str] = None,
+      mode: Optional[Mode] = None,
+  ):
+    super().__init__(name=name, mode=mode)
+    if num_embeddings < 0:
+      raise ValueError("num_embeddings must not be negative.")
+    if embedding_size < 0:
+      raise ValueError("embedding_size must not be negative.")
+    self.num_embeddings = num_embeddings
+    self.embedding_size = embedding_size
+    self.out_size = (embedding_size,)
+
+    weight = init.param(embed_init, (self.num_embeddings, self.embedding_size))
+    if self.mode.has(Training):
+      self.weight = ParamState(weight)
+    else:
+      self.weight = weight
+
+  def update(self, indices: ArrayLike):
+    if self.mode.has(Training):
+      return self.weight.value[indices]
+    return self.weight[indices]

brainstate/nn/_misc.py

@@ -20,9 +20,10 @@ from enum import Enum
 from functools import wraps
 from typing import Sequence, Callable
 
+import brainunit as bu
 import jax.numpy as jnp
 
-from .. import environ, math
+from .. import environ
 from .._state import State
 from ..transform import vector_grad
 
@@ -96,7 +97,7 @@ def exp_euler(fun):
       )
     dt = environ.get('dt')
     linear, derivative = vector_grad(fun, argnums=0, return_value=True)(*args, **kwargs)
-    phi = math.exprel(dt * linear)
+    phi = bu.math.exprel(dt * linear)
     return args[0] + dt * phi * derivative
 
   return integral
@@ -128,5 +129,5 @@ def exp_euler_step(fun: Callable, *args, **kwargs):
     )
   dt = environ.get('dt')
   linear, derivative = vector_grad(fun, argnums=0, return_value=True)(*args, **kwargs)
-  phi = math.exprel(dt * linear)
+  phi = bu.math.exprel(dt * linear)
   return args[0] + dt * phi * derivative

brainstate/nn/_others.py

@@ -19,10 +19,11 @@ from __future__ import annotations
 from functools import partial
 from typing import Optional
 
+import brainunit as bu
 import jax.numpy as jnp
 
 from ._base import DnnLayer
-from .. import random, math, environ, typing, init
+from .. import random, environ, typing, init
 from ..mixin import Mode
 
 __all__ = [
@@ -88,7 +89,7 @@ class DropoutFixed(DnnLayer):
     self.mask = init.param(partial(random.bernoulli, self.prob), self.in_size, batch_size)
 
   def update(self, x):
-    dtype = math.get_dtype(x)
+    dtype = bu.math.get_dtype(x)
     fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
     if fit_phase:
       assert self.mask.shape == x.shape, (f"Input shape {x.shape} does not match the mask shape {self.mask.shape}. "

brainstate/nn/_poolings.py

@@ -21,12 +21,13 @@ import functools
 from typing import Sequence, Optional
 from typing import Union, Tuple, Callable, List
 
+import brainunit as bu
 import jax
 import jax.numpy as jnp
 import numpy as np
 
 from ._base import DnnLayer, ExplicitInOutSize
-from .. import environ, math
+from .. import environ
 from ..mixin import Mode
 from ..typing import Size
 
@@ -53,8 +54,8 @@ class Flatten(DnnLayer, ExplicitInOutSize):
 
   Args:
       in_size: Sequence of int. The shape of the input tensor.
-      start_dim: first dim to flatten (default = 1).
-      end_dim: last dim to flatten (default = -1).
+      start_axis: first dim to flatten (default = 1).
+      end_axis: last dim to flatten (default = -1).
 
   Examples::
       >>> import brainstate as bst
@@ -74,36 +75,36 @@ class Flatten(DnnLayer, ExplicitInOutSize):
 
   def __init__(
       self,
-      start_dim: int = 0,
-      end_dim: int = -1,
+      start_axis: int = 0,
+      end_axis: int = -1,
       in_size: Optional[Size] = None
   ) -> None:
     super().__init__()
-    self.start_dim = start_dim
-    self.end_dim = end_dim
+    self.start_axis = start_axis
+    self.end_axis = end_axis
 
     if in_size is not None:
       self.in_size = tuple(in_size)
-      y = jax.eval_shape(functools.partial(math.flatten, start_dim=start_dim, end_dim=end_dim),
+      y = jax.eval_shape(functools.partial(bu.math.flatten, start_axis=start_axis, end_axis=end_axis),
                          jax.ShapeDtypeStruct(self.in_size, environ.dftype()))
       self.out_size = y.shape
 
   def update(self, x):
     if self._in_size is None:
-      start_dim = self.start_dim if self.start_dim >= 0 else x.ndim + self.start_dim
+      start_axis = self.start_axis if self.start_axis >= 0 else x.ndim + self.start_axis
     else:
       assert x.ndim >= len(self.in_size), 'Input tensor has fewer dimensions than the expected shape.'
       dim_diff = x.ndim - len(self.in_size)
       if self.in_size != x.shape[dim_diff:]:
         raise ValueError(f'Input tensor has shape {x.shape}, but expected shape {self.in_size}.')
-      if self.start_dim >= 0:
-        start_dim = self.start_dim + dim_diff
+      if self.start_axis >= 0:
+        start_axis = self.start_axis + dim_diff
       else:
-        start_dim = x.ndim + self.start_dim
-    return math.flatten(x, start_dim, self.end_dim)
+        start_axis = x.ndim + self.start_axis
+    return bu.math.flatten(x, start_axis, self.end_axis)
 
   def __repr__(self) -> str:
-    return f'{self.__class__.__name__}(start_dim={self.start_dim}, end_dim={self.end_dim})'
+    return f'{self.__class__.__name__}(start_axis={self.start_axis}, end_axis={self.end_axis})'
 
 
 class Unflatten(DnnLayer, ExplicitInOutSize):
@@ -124,7 +125,7 @@ class Unflatten(DnnLayer, ExplicitInOutSize):
         :math:`\prod_{i=1}^n U_i = S_{\text{dim}}`.
 
   Args:
-      dim: int, Dimension to be unflattened.
+      axis: int, Dimension to be unflattened.
       sizes: Sequence of int. New shape of the unflattened dimension.
       in_size: Sequence of int. The shape of the input tensor.
   """
@@ -132,7 +133,7 @@ class Unflatten(DnnLayer, ExplicitInOutSize):
 
   def __init__(
       self,
-      dim: int,
+      axis: int,
       sizes: Size,
       mode: Mode = None,
       name: str = None,
@@ -140,7 +141,7 @@ class Unflatten(DnnLayer, ExplicitInOutSize):
   ) -> None:
     super().__init__(mode=mode, name=name)
 
-    self.dim = dim
+    self.axis = axis
     self.sizes = sizes
     if isinstance(sizes, (tuple, list)):
       for idx, elem in enumerate(sizes):
@@ -152,15 +153,15 @@ class Unflatten(DnnLayer, ExplicitInOutSize):
 
     if in_size is not None:
       self.in_size = tuple(in_size)
-      y = jax.eval_shape(functools.partial(math.unflatten, dim=dim, sizes=sizes),
+      y = jax.eval_shape(functools.partial(bu.math.unflatten, axis=axis, sizes=sizes),
                          jax.ShapeDtypeStruct(self.in_size, environ.dftype()))
       self.out_size = y.shape
 
   def update(self, x):
-    return math.unflatten(x, self.dim, self.sizes)
+    return bu.math.unflatten(x, self.axis, self.sizes)
 
   def __repr__(self):
-    return f'{self.__class__.__name__}(dim={self.dim}, sizes={self.sizes})'
+    return f'{self.__class__.__name__}(axis={self.axis}, sizes={self.sizes})'
 
 
 class _MaxPool(DnnLayer, ExplicitInOutSize):

brainstate/nn/_poolings_test.py

@@ -18,7 +18,7 @@ class TestFlatten(parameterized.TestCase):
       (10, 20, 30),
     ]:
       arr = bst.random.rand(*size)
-      f = nn.Flatten(start_dim=0)
+      f = nn.Flatten(start_axis=0)
       out = f(arr)
       self.assertTrue(out.shape == (np.prod(size),))
 
@@ -29,21 +29,21 @@ class TestFlatten(parameterized.TestCase):
       (10, 20, 30),
     ]:
       arr = bst.random.rand(*size)
-      f = nn.Flatten(start_dim=1)
+      f = nn.Flatten(start_axis=1)
       out = f(arr)
       self.assertTrue(out.shape == (size[0], np.prod(size[1:])))
 
   def test_flatten3(self):
     size = (16, 32, 32, 8)
     arr = bst.random.rand(*size)
-    f = nn.Flatten(start_dim=0, in_size=(32, 8))
+    f = nn.Flatten(start_axis=0, in_size=(32, 8))
     out = f(arr)
     self.assertTrue(out.shape == (16, 32, 32 * 8))
 
   def test_flatten4(self):
     size = (16, 32, 32, 8)
     arr = bst.random.rand(*size)
-    f = nn.Flatten(start_dim=1, in_size=(32, 32, 8))
+    f = nn.Flatten(start_axis=1, in_size=(32, 32, 8))
     out = f(arr)
     self.assertTrue(out.shape == (16, 32, 32 * 8))