brainpy.math.random.shuffle is numpy like

brainpy/math/__init__.py

@@ -33,7 +33,7 @@
 
 # data structure
 from .jaxarray import *
-from .delay_vars import *
+from .delayvars import *
 
 # high-level numpy operations
 from .numpy_ops import *

brainpy/math/compat/__init__.py

@@ -6,5 +6,5 @@
 ]
 
 from . import optimizers, losses
-from .delay_vars import *
+from .delayvars import *
 

brainpy/math/compat/delay_vars.py → brainpy/math/compat/delayvars.py

@@ -7,7 +7,7 @@
 
 from brainpy.math.jaxarray import ndarray
 from brainpy.math.numpy_ops import zeros
-from brainpy.math.delay_vars import TimeDelay
+from brainpy.math.delayvars import TimeDelay
 
 
 __all__ = [

brainpy/math/random.py

@@ -33,9 +33,7 @@ def _size2shape(size):
 
 
 class RandomState(Variable):
-  """RandomState are variables that track the
-  random generator state. They are meant to be used internally.
-  Currently only the random.Generator module uses them."""
+  """RandomState that track the random generator state. """
   __slots__ = ()
 
   def __init__(self, seed=None):
@@ -126,9 +124,8 @@ def permutation(self, x):
     return JaxArray(jr.permutation(self.split_key(), x))
 
   def shuffle(self, x, axis=0):
-    x = x.value if isinstance(x, JaxArray) else x
-    # return JaxArray(jr.permutation(self.split_key(), x, axis=axis, independent=True))
-    return JaxArray(jr.permutation(self.split_key(), x, axis=axis))
+    assert isinstance(x, JaxArray), f'Must be a JaxArray, but got {type(x)}'
+    x.value = jr.permutation(self.split_key(), x.value, axis=axis)
 
   def beta(self, a, b, size=None):
     a = a.value if isinstance(a, JaxArray) else a
@@ -251,8 +248,8 @@ def permutation(x):
 
 
 def shuffle(x, axis=0):
-  x = x.value if isinstance(x, JaxArray) else x
-  return JaxArray(jr.permutation(DEFAULT.split_key(), x, axis=axis, independent=True))
+  assert isinstance(x, JaxArray), f'Must be a JaxArray, but got {type(x)}'
+  x.value = jr.permutation(DEFAULT.split_key(), x.value, axis=axis)
 
 
 def beta(a, b, size=None):

docs/auto_generater.py

@@ -5,14 +5,14 @@
 import os
 
 from brainpy.math import (activations, autograd, controls, function,
-                          jit, operators, parallels, setting, delay_vars,
+                          jit, operators, parallels, setting, delayvars,
                           compat)
 
 block_list = ['test', 'register_pytree_node']
 for module in [jit, autograd, function,
                controls, activations,
                operators, parallels, setting,
-               delay_vars, compat]:
+               delayvars, compat]:
   for k in dir(module):
     if (not k.startswith('_')) and (not inspect.ismodule(getattr(module, k))):
       block_list.append(k)
@@ -417,7 +417,7 @@ def generate_math_docs(path='apis/auto/math/'):
   write_module(module_name='brainpy.math.function',
                filename=os.path.join(path, 'function.rst'),
                header='Function')
-  write_module(module_name='brainpy.math.delay_vars',
+  write_module(module_name='brainpy.math.delayvars',
                filename=os.path.join(path, 'delay_vars.rst'),
                header='Delay Variables')