fix bugs

Author: chaoming0625

brainpy/_src/analysis/constants.py

@@ -2,9 +2,6 @@
 
 
 __all__ = [
-  'CONTINUOUS',
-  'DISCRETE',
-
   'F_vmap_fx',
   'F_vmap_fy',
   'F_vmap_brentq_fx',

brainpy/_src/math/random.py

@@ -973,8 +973,8 @@ def wald(self, mean, scale, size=None, key=None):
     if size is None:
       size = lax.broadcast_shapes(jnp.shape(mean), jnp.shape(scale))
     size = _size2shape(size)
-    sampled_chi2 = jnp.square(self.randn(*size))
-    sampled_uniform = self.uniform(size=size, key=key)
+    sampled_chi2 = jnp.square(_as_jax_array(self.randn(*size)))
+    sampled_uniform = _as_jax_array(self.uniform(size=size, key=key))
     # Wikipedia defines an intermediate x with the formula
     #   x = loc + loc ** 2 * y / (2 * conc) - loc / (2 * conc) * sqrt(4 * loc * conc * y + loc ** 2 * y ** 2)
     # where y ~ N(0, 1)**2 (sampled_chi2 above) and conc is the concentration.

brainpy/analysis.py

@@ -17,6 +17,9 @@
   SlowPointFinder as SlowPointFinder,
 )
 
+from brainpy._src.analysis.constants import (CONTINUOUS as CONTINUOUS,
+                                             DISCRETE as DISCRETE)
+
 from brainpy._src.analysis import plotstyle, stability, constants
 C = constants