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

brainstate/_module.py

@@ -59,7 +59,7 @@ import numpy as np
 from . import environ
 from ._state import State, StateDictManager, visible_state_dict
 from ._utils import set_module_as
-from .mixin import Mixin, Mode, DelayedInit, AllOfTypes, Batching, UpdateReturn
+from .mixin import Mixin, Mode, DelayedInit, JointTypes, Batching, UpdateReturn
 from .transform import jit_error
 from .util import unique_name, DictManager, get_unique_name
 
@@ -809,7 +809,6 @@ class Dynamics(ExtendedUpdateWithBA, ReceiveInputProj, UpdateReturn):
       keep_size: bool = False,
       name: Optional[str] = None,
       mode: Optional[Mode] = None,
-      method: str = 'exp_auto'
   ):
     # size
     if isinstance(size, (list, tuple)):
@@ -831,9 +830,6 @@ class Dynamics(ExtendedUpdateWithBA, ReceiveInputProj, UpdateReturn):
     # number of neurons
     self.num = np.prod(size)
 
-    # integration method
-    self.method = method
-
     # -- Attribute for "InputProjMixIn" -- #
     # each instance of "SupportInputProj" should have
     # "_current_inputs" and "_delta_inputs" attributes
@@ -1213,10 +1209,10 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
     assert self.history is not None, 'The delay history is not initialized.'
     assert delay_step is not None, 'The delay step should be given.'
 
-    if environ.get(environ.JIT_ERROR_CHECK, True):
+    if environ.get(environ.JIT_ERROR_CHECK, False):
       def _check_delay(delay_len):
         raise ValueError(f'The request delay length should be less than the '
-                         f'maximum delay {self.max_length}. But we got {delay_len}')
+                         f'maximum delay {self.max_length - 1}. But we got {delay_len}')
 
       jit_error(delay_step >= self.max_length, _check_delay, delay_step)
 
@@ -1263,9 +1259,8 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
     current_time = environ.get(environ.T, desc='The current time.')
     dt = environ.get_dt()
 
-    if environ.get(environ.JIT_ERROR_CHECK, True):
-      def _check_delay(args):
-        t_now, t_delay = args
+    if environ.get(environ.JIT_ERROR_CHECK, False):
+      def _check_delay(t_now, t_delay):
         raise ValueError(f'The request delay time should be within '
                          f'[{t_now - self.max_time - dt}, {t_now}], '
                          f'but we got {t_delay}')
@@ -1273,7 +1268,7 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
       jit_error(jnp.logical_or(delay_time > current_time,
                                delay_time < current_time - self.max_time - dt),
                 _check_delay,
-                (current_time, delay_time))
+                current_time, delay_time)
 
     diff = current_time - delay_time
     float_time_step = diff / dt
@@ -1415,7 +1410,7 @@ class DelayAccess(Module):
     return self.refs['delay'].at(self._delay_entry, *self.indices)
 
 
-def register_delay_of_target(target: AllOfTypes[ExtendedUpdateWithBA, UpdateReturn]):
+def register_delay_of_target(target: JointTypes[ExtendedUpdateWithBA, UpdateReturn]):
   """Register delay class for the given target.
 
   Args:
@@ -1425,7 +1420,7 @@ def register_delay_of_target(target: AllOfTypes[ExtendedUpdateWithBA, UpdateRetu
     The delay registered for the given target.
   """
   if not target.has_after_update(delay_identifier):
-    assert isinstance(target, AllOfTypes[ExtendedUpdateWithBA, UpdateReturn])
+    assert isinstance(target, JointTypes[ExtendedUpdateWithBA, UpdateReturn])
     target.add_after_update(delay_identifier, Delay(target.update_return_info()))
   delay_cls = target.get_after_update(delay_identifier)
   return delay_cls

brainstate/_module_test.py

@@ -86,7 +86,7 @@ class TestDelay(unittest.TestCase):
     rotation_delay = bst.Delay(jnp.ones([1]), time=2., delay_method='concat', interp_method='round')
     rotation_delay.init_state()
 
-    with bst.environ.context(i=0, t=0):
+    with bst.environ.context(i=0, t=0, jit_error_check=True):
       rotation_delay.retrieve_at_time(-2.0)
       with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
         rotation_delay.retrieve_at_time(-2.1)

brainstate/_random_for_unit.py

@@ -0,0 +1,48 @@
+# 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.
+# ==============================================================================
+import brainunit as bu
+import jax
+import jax.random as jr
+
+from .typing import ArrayLike, Size, DTypeLike
+
+
+def uniform_for_unit(
+    key,
+    shape: Size = (),
+    dtype: DTypeLike = float,
+    minval: ArrayLike = 0.,
+    maxval: ArrayLike = 1.
+) -> jax.Array | bu.Quantity:
+  if isinstance(minval, bu.Quantity) and isinstance(maxval, bu.Quantity):
+    return bu.Quantity(jr.uniform(key, shape, dtype, minval.value, maxval.value), dim=minval.dim)
+  elif isinstance(minval, bu.Quantity):
+    assert minval.is_unitless, f'minval must be unitless when maxval is not a Quantity, got {minval}'
+    minval = minval.value
+  elif isinstance(maxval, bu.Quantity):
+    assert maxval.is_unitless, f'maxval must be unitless when minval is not a Quantity, got {maxval}'
+    maxval = maxval.value
+  return jr.uniform(key, shape, dtype, minval, maxval)
+
+
+def permutation_for_unit(
+    key,
+    x: int | ArrayLike,
+    axis: int = 0,
+    independent: bool = False
+) -> jax.Array | bu.Quantity:
+  if isinstance(x, bu.Quantity):
+    return bu.Quantity(jr.permutation(key, x.value, axis, independent=independent), dim=x.dim)
+  return jr.permutation(key, x, axis, independent=independent)

brainstate/_state.py

@@ -29,7 +29,9 @@ max_int = np.iinfo(np.int32)
 
 __all__ = [
   'State', 'ShortTermState', 'LongTermState', 'ParamState',
-  'StateDictManager', 'visible_state_dict',
+  'StateDictManager',
+  'StateTrace',
+  'visible_state_dict',
   'check_state_value_tree',
 ]
 
@@ -141,7 +143,7 @@ class State(object):
     """
     # value checking
     v = v.value if isinstance(v, State) else v
-    self._check_value(v)
+    self._check_value_tree(v)
     # write the value by the stack (>= level)
     trace: StateTrace
     for trace in thread_local_stack.stack[self._level:]:
@@ -149,9 +151,9 @@ class State(object):
     # set the value
     self._value = v
 
-  def _check_value(self, v):
+  def _check_value_tree(self, v):
     if self._check_tree or _global_context_to_check_state_tree[-1]:
-      in_tree = jax.tree_util.tree_structure(v)
+      in_tree = jax.tree.structure(v)
       if in_tree != self._tree:
         self._raise_error_with_source_info(
           ValueError(f'The given value {in_tree} does not '
@@ -370,12 +372,13 @@ class StateTrace(object):
       self.types[index] = 'write'
       self._written_ids.add(id_)
 
-  def collect_values(self, *categories: str) -> Tuple:
+  def collect_values(self, *categories: str, check_val_tree: bool = False) -> Tuple:
     """
     Collect the values by the given categories.
 
     Args:
       *categories: The categories.
+      check_val_tree: Whether to check the tree structure of the value.
 
     Returns:
       results: The values.
@@ -383,7 +386,10 @@ class StateTrace(object):
     results = []
     for st, ty in zip(self.states, self.types):
       if ty in categories:
-        results.append(st.value)
+        val = st.value
+        if check_val_tree:
+          st._check_value_tree(val)
+        results.append(val)
     return tuple(results)
 
   def recovery_original_values(self) -> None:

brainstate/init/_generic.py

@@ -15,15 +15,15 @@
 
 # -*- coding: utf-8 -*-
 
-import numbers
 from typing import Union, Callable, Optional, Sequence
 
+import brainunit as bu
 import jax
-import jax.numpy as jnp
 import numpy as np
 
 from brainstate._state import State
 from ._base import to_size
+from ..typing import ArrayLike
 
 __all__ = [
   'param',
@@ -33,11 +33,57 @@ __all__ = [
 
 
 def _is_scalar(x):
-  return isinstance(x, numbers.Number)
+  return bu.math.isscalar(x)
+
+
+def are_shapes_broadcastable(shape1, shape2):
+  """
+  Check if two shapes are broadcastable.
+
+  Parameters:
+  - shape1: Tuple[int], the shape of the first array.
+  - shape2: Tuple[int], the shape of the second array.
+
+  Returns:
+  - bool: True if shapes are broadcastable, False otherwise.
+  """
+  # Reverse the shapes to compare from the last dimension
+  shape1_reversed = shape1[::-1]
+  shape2_reversed = shape2[::-1]
+
+  # Iterate over the dimensions of the shorter shape
+  for dim1, dim2 in zip(shape1_reversed, shape2_reversed):
+    # Check if the dimensions are not equal and neither is 1
+    if dim1 != dim2 and 1 not in (dim1, dim2):
+      return False
+
+  # If all dimensions are compatible, the shapes are broadcastable
+  return True
+
+
+def _expand_params_to_match_sizes(params, sizes):
+  """
+  Expand the dimensions of params to match the dimensions of sizes.
+
+  Parameters:
+  - params: jax.Array or np.ndarray, the parameter array to be expanded.
+  - sizes: tuple[int] or list[int], the target shape dimensions.
+
+  Returns:
+  - Expanded params with dimensions matching sizes.
+  """
+  params_dim = params.ndim
+  sizes_dim = len(sizes)
+  dim_diff = sizes_dim - params_dim
+
+  # Add new axes to params if it has fewer dimensions than sizes
+  for _ in range(dim_diff):
+    params = bu.math.expand_dims(params, axis=0)  # Add new axis at the last dimension
+  return params
 
 
 def param(
-    param: Union[Callable, np.ndarray, jax.Array, float, int, bool],
+    parameter: Union[Callable, ArrayLike],
     sizes: Union[int, Sequence[int]],
     batch_size: Optional[int] = None,
     allow_none: bool = True,
@@ -47,7 +93,7 @@ def param(
 
   Parameters
   ----------
-  param: callable, Initializer, bm.ndarray, jnp.ndarray, onp.ndarray, float, int, bool
+  parameter: callable, ArrayLike, State
     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.
@@ -71,42 +117,55 @@ def param(
   --------
   noise, state
   """
-  if param is None:
+  # Check if the parameter is None
+  if parameter 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)
+  # Check if the parameter is a scalar value
+  if allow_scalar and _is_scalar(parameter):
+    return parameter
 
-  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
+  # Convert sizes to a tuple
+  sizes = tuple(to_size(sizes))
+
+  # Check if the parameter is a callable function
+  if callable(parameter):
     if batch_size is not None:
-      param = type(param)(jnp.repeat(jnp.expand_dims(param.value, axis=batch_axis), batch_size, axis=batch_axis))
+      sizes = (batch_size,) + sizes
+    return parameter(sizes)
+  elif isinstance(parameter, (np.ndarray, jax.Array, bu.Quantity, State)):
+    parameter = parameter
   else:
-    raise ValueError(f'Unknown parameter type: {type(param)}')
+    raise ValueError(f'Unknown parameter type: {type(parameter)}')
+
+  # Check if the shape of the parameter matches the given size
+  if not are_shapes_broadcastable(parameter.shape, sizes):
+    raise ValueError(f'The shape of the parameter {parameter.shape} does not match with the given size {sizes}')
 
-  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
+  # Expand the parameter to match the given batch size
+  param_value = parameter.value if isinstance(parameter, State) else parameter
+  if batch_size is not None:
+    if param_value.ndim <= len(sizes):
+      # add a new axis to the params so that it matches the dimensionality of the given shape ``sizes``
+      param_value = _expand_params_to_match_sizes(param_value, sizes)
+      param_value = bu.math.repeat(
+        bu.math.expand_dims(param_value, axis=0),
+        batch_size,
+        axis=0
+      )
+    else:
+      if param_value.shape[0] != batch_size:
+        raise ValueError(f'The batch size of the parameter {param_value.shape[0]} '
+                         f'does not match with the given batch size {batch_size}')
+  return type(parameter)(param_value) if isinstance(parameter, State) else param_value
 
 
 def state(
-    init: Union[Callable, np.ndarray, jax.Array],
+    init: Union[Callable, jax.typing.ArrayLike],
     sizes: Union[int, Sequence[int]] = None,
     batch_size: Optional[int] = None,
 ):
@@ -124,18 +183,24 @@ def state(
 
   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 bu.math.shape(init) != sizes:
+        raise ValueError(f'The shape of "data" {bu.math.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))
+      data = State(
+        bu.math.repeat(
+          bu.math.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]],
+    noises: Optional[Union[ArrayLike, Callable]],
     size: Union[int, Sequence[int]],
     num_vars: int = 1,
     noise_idx: int = 0,

brainstate/init/_random_inits.py

@@ -17,11 +17,13 @@
 
 import math
 
+import brainunit as bu
 import jax.numpy as jnp
 import numpy as np
 
 from brainstate import environ, random
 from ._base import Initializer, to_size
+from ..typing import ArrayLike
 
 __all__ = [
   'Normal',
@@ -260,7 +262,7 @@ class Uniform(Initializer):
 class VarianceScaling(Initializer):
   def __init__(
       self,
-      scale: float,
+      scale: ArrayLike,
       mode: str,
       distribution: str,
       in_axis: int = -2,
@@ -287,7 +289,9 @@ class VarianceScaling(Initializer):
       denominator = (fan_in + fan_out) / 2
     else:
       raise ValueError("invalid mode for variance scaling initializer: {}".format(self.mode))
-    variance = (self.scale / denominator).astype(self.dtype)
+    scale = self.scale.value if isinstance(self.scale, bu.Quantity) else self.scale
+    dim = self.scale.dim if isinstance(self.scale, bu.Quantity) else bu.DIMENSIONLESS
+    variance = (scale / denominator).astype(self.dtype)
     if self.distribution == "truncated_normal":
       stddev = (jnp.sqrt(variance) / .87962566103423978).astype(self.dtype)
       res = random.truncated_normal(-2, 2, shape, dtype=self.dtype) * stddev
@@ -298,7 +302,7 @@ class VarianceScaling(Initializer):
              jnp.sqrt(3 * variance).astype(self.dtype))
     else:
       raise ValueError("invalid distribution for variance scaling initializer")
-    return res
+    return res if dim == bu.DIMENSIONLESS else res * dim
 
   def __repr__(self):
     name = self.__class__.__name__
@@ -425,7 +429,7 @@ class Orthogonal(Initializer):
 
   def __init__(
       self,
-      scale: float = 1.,
+      scale: ArrayLike = 1.,
       axis: int = -1,
       dtype=None
   ):
@@ -440,6 +444,9 @@ class Orthogonal(Initializer):
     n_cols = np.prod(shape) // n_rows
     matrix_shape = (n_rows, n_cols) if n_rows > n_cols else (n_cols, n_rows)
     norm_dst = random.normal(size=matrix_shape, dtype=self.dtype)
+
+    scale = self.scale.value if isinstance(self.scale, bu.Quantity) else self.scale
+    dim = self.scale.dim if isinstance(self.scale, bu.Quantity) else bu.DIMENSIONLESS
     q_mat, r_mat = jnp.linalg.qr(norm_dst)
     # Enforce Q is uniformly distributed
     q_mat *= jnp.sign(jnp.diag(r_mat))
@@ -447,7 +454,8 @@ class Orthogonal(Initializer):
       q_mat = q_mat.T
     q_mat = jnp.reshape(q_mat, (n_rows,) + tuple(np.delete(shape, self.axis)))
     q_mat = jnp.moveaxis(q_mat, 0, self.axis)
-    return jnp.asarray(self.scale, dtype=self.dtype) * q_mat
+    r = jnp.asarray(scale, dtype=self.dtype) * q_mat
+    return r if dim == bu.DIMENSIONLESS else r * dim
 
   def __repr__(self):
     return f'{self.__class__.__name__}(scale={self.scale}, axis={self.axis}, dtype={self.dtype})'
@@ -472,7 +480,9 @@ class DeltaOrthogonal(Initializer):
       raise ValueError("Delta orthogonal initializer requires a 3D, 4D or 5D shape.")
     if shape[-1] < shape[-2]:
       raise ValueError("`fan_in` must be less or equal than `fan_out`. ")
-    ortho_matrix = Orthogonal(scale=self.scale, axis=self.axis, dtype=self.dtype)(*shape[-2:])
+    scale = self.scale.value if isinstance(self.scale, bu.Quantity) else self.scale
+    dim = self.scale.dim if isinstance(self.scale, bu.Quantity) else bu.DIMENSIONLESS
+    ortho_matrix = Orthogonal(scale=scale, axis=self.axis, dtype=self.dtype)(*shape[-2:])
     W = jnp.zeros(shape, dtype=self.dtype)
     if len(shape) == 3:
       k = shape[0]
@@ -483,7 +493,7 @@ class DeltaOrthogonal(Initializer):
     else:
       k1, k2, k3 = shape[:3]
       W = W.at[(k1 - 1) // 2, (k2 - 1) // 2, (k3 - 1) // 2].set(ortho_matrix)
-    return W
+    return W if dim == bu.DIMENSIONLESS else W * dim
 
   def __repr__(self):
     return f'{self.__class__.__name__}(scale={self.scale}, axis={self.axis}, dtype={self.dtype})'

brainstate/init/_regular_inits.py

@@ -15,7 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
-import jax.numpy as jnp
+
+import brainunit as bu
 
 from brainstate import environ
 from ._base import Initializer, to_size
@@ -39,7 +40,7 @@ class ZeroInit(Initializer):
 
   def __call__(self, shape):
     shape = to_size(shape)
-    return jnp.zeros(shape, dtype=self.dtype)
+    return bu.math.zeros(shape, dtype=self.dtype)
 
   def __repr__(self):
     return f"{self.__class__.__name__}(dtype={self.dtype})"
@@ -59,11 +60,11 @@ class Constant(Initializer):
   def __init__(self, value=1., dtype=None):
     super(Constant, self).__init__()
     self.dtype = dtype or environ.dftype()
-    self.value = jnp.asarray(value, dtype=self.dtype)
+    self.value = bu.math.asarray(value, dtype=self.dtype)
 
   def __call__(self, shape):
     shape = to_size(shape)
-    return jnp.full(shape, self.value, dtype=self.dtype)
+    return bu.math.full(shape, self.value, dtype=self.dtype)
 
   def __repr__(self):
     return f'{self.__class__.__name__}(value={self.value}, dtype={self.dtype})'
@@ -94,15 +95,15 @@ class Identity(Initializer):
   def __init__(self, value=1., dtype=None):
     super(Identity, self).__init__()
     self.dtype = dtype or environ.dftype()
-    self.value = jnp.asarray(value, dtype=self.dtype)
+    self.value = bu.math.asarray(value, dtype=self.dtype)
 
   def __call__(self, shape):
     shape = to_size(shape)
     if isinstance(shape, (tuple, list)):
       if len(shape) > 2:
         raise ValueError(f'Only support initialize 2D weights for {self.__class__.__name__}.')
-    r = jnp.eye(shape, dtype=self.dtype)
-    r = jnp.fill_diagonal(r, self.value)
+    r = bu.math.eye(shape, dtype=self.dtype)
+    r = bu.math.fill_diagonal(r, self.value)
     return r
 
   def __repr__(self):

brainstate/mixin.py

@@ -32,7 +32,7 @@ __all__ = [
   'UpdateReturn',
 
   # types
-  'AllOfTypes',
+  'JointTypes',
   'OneOfTypes',
 
   # behavior modes
@@ -206,7 +206,7 @@ class _JointGenericAlias(_UnionGenericAlias, _root=True):
 
 
 @_SpecialForm
-def AllOfTypes(self, parameters):
+def JointTypes(self, parameters):
   """All of types; AllOfTypes[X, Y] means both X and Y.
 
   To define a union, use e.g. Union[int, str].
@@ -341,7 +341,7 @@ class JointMode(Mode):
     """
     Check whether the mode is exactly the desired mode.
     """
-    return AllOfTypes[tuple(self.types)] == cls
+    return JointTypes[tuple(self.types)] == cls
 
   def __getattr__(self, item):
     """

brainstate/mixin_test.py

@@ -23,7 +23,7 @@ class TestMixin(unittest.TestCase):
     self.assertTrue(bc.mixin.Mixin)
     self.assertTrue(bc.mixin.DelayedInit)
     self.assertTrue(bc.mixin.DelayedInitializer)
-    self.assertTrue(bc.mixin.AllOfTypes)
+    self.assertTrue(bc.mixin.JointTypes)
     self.assertTrue(bc.mixin.OneOfTypes)
     self.assertTrue(bc.mixin.Mode)
     self.assertTrue(bc.mixin.Batching)
@@ -33,29 +33,29 @@ class TestMixin(unittest.TestCase):
 class TestMode(unittest.TestCase):
   def test_JointMode(self):
     a = bc.mixin.JointMode(bc.mixin.Batching(), bc.mixin.Training())
-    self.assertTrue(a.is_a(bc.mixin.AllOfTypes[bc.mixin.Batching, bc.mixin.Training]))
+    self.assertTrue(a.is_a(bc.mixin.JointTypes[bc.mixin.Batching, bc.mixin.Training]))
     self.assertTrue(a.has(bc.mixin.Batching))
     self.assertTrue(a.has(bc.mixin.Training))
     b = bc.mixin.JointMode(bc.mixin.Batching())
-    self.assertTrue(b.is_a(bc.mixin.AllOfTypes[bc.mixin.Batching]))
+    self.assertTrue(b.is_a(bc.mixin.JointTypes[bc.mixin.Batching]))
     self.assertTrue(b.is_a(bc.mixin.Batching))
     self.assertTrue(b.has(bc.mixin.Batching))
 
   def test_Training(self):
     a = bc.mixin.Training()
     self.assertTrue(a.is_a(bc.mixin.Training))
-    self.assertTrue(a.is_a(bc.mixin.AllOfTypes[bc.mixin.Training]))
+    self.assertTrue(a.is_a(bc.mixin.JointTypes[bc.mixin.Training]))
     self.assertTrue(a.has(bc.mixin.Training))
-    self.assertTrue(a.has(bc.mixin.AllOfTypes[bc.mixin.Training]))
+    self.assertTrue(a.has(bc.mixin.JointTypes[bc.mixin.Training]))
     self.assertFalse(a.is_a(bc.mixin.Batching))
     self.assertFalse(a.has(bc.mixin.Batching))
 
   def test_Batching(self):
     a = bc.mixin.Batching()
     self.assertTrue(a.is_a(bc.mixin.Batching))
-    self.assertTrue(a.is_a(bc.mixin.AllOfTypes[bc.mixin.Batching]))
+    self.assertTrue(a.is_a(bc.mixin.JointTypes[bc.mixin.Batching]))
     self.assertTrue(a.has(bc.mixin.Batching))
-    self.assertTrue(a.has(bc.mixin.AllOfTypes[bc.mixin.Batching]))
+    self.assertTrue(a.has(bc.mixin.JointTypes[bc.mixin.Batching]))
 
     self.assertFalse(a.is_a(bc.mixin.Training))
     self.assertFalse(a.has(bc.mixin.Training))
@@ -63,9 +63,9 @@ class TestMode(unittest.TestCase):
   def test_Mode(self):
     a = bc.mixin.Mode()
     self.assertTrue(a.is_a(bc.mixin.Mode))
-    self.assertTrue(a.is_a(bc.mixin.AllOfTypes[bc.mixin.Mode]))
+    self.assertTrue(a.is_a(bc.mixin.JointTypes[bc.mixin.Mode]))
     self.assertTrue(a.has(bc.mixin.Mode))
-    self.assertTrue(a.has(bc.mixin.AllOfTypes[bc.mixin.Mode]))
+    self.assertTrue(a.has(bc.mixin.JointTypes[bc.mixin.Mode]))
 
     self.assertFalse(a.is_a(bc.mixin.Training))
     self.assertFalse(a.has(bc.mixin.Training))