update delays for compatible with brainunit (#14)

brainstate/_module.py

@@ -1597,6 +1597,6 @@ def _get_delay(delay_time, delay_step):
       delay_time = delay_step * environ.get_dt()
   else:
     assert delay_step is None, '"delay_step" should be None if "delay_time" is given.'
-    assert isinstance(delay_time, (int, float))
+    # assert isinstance(delay_time, (int, float))
     delay_step = math.ceil(delay_time / environ.get_dt())
   return delay_time, delay_step

brainstate/_state.py

@@ -234,6 +234,8 @@ class ShapeDtype:
   def __init__(self, shape, dtype):
     self.shape = shape
     self.dtype = dtype
+    self.ndim = len(shape)
+    self.size = np.prod(shape)
 
   def __repr__(self):
     return f'{self.dtype}{list(self.shape)}'

brainstate/nn/_projection/_align_post.py

@@ -15,6 +15,8 @@
 
 from typing import Optional, Union
 
+
+import brainunit as u
 from brainstate._module import (register_delay_of_target,
                                 Projection,
                                 Module,
@@ -278,11 +280,19 @@ def __init__(
     self.comm = comm
 
     # delay initialization
-    if delay is not None and delay > 0.:
-      delay_cls = register_delay_of_target(pre)
-      delay_cls.register_entry(self.name, delay)
-      self.delay = delay_cls
-      self.has_delay = True
+    if delay is not None:
+      if isinstance(delay, u.Quantity):
+        has_delay = delay.mantissa > 0.
+      else:
+        has_delay = delay > 0.
+      if has_delay:
+        delay_cls = register_delay_of_target(pre)
+        delay_cls.register_entry(self.name, delay)
+        self.delay = delay_cls
+        self.has_delay = True
+      else:
+        self.delay = None
+        self.has_delay = False
     else:
       self.delay = None
       self.has_delay = False
@@ -502,11 +512,19 @@ def __init__(
     self.out = out
 
     # delay initialization
-    if delay is not None and delay > 0.:
-      delay_cls = register_delay_of_target(pre)
-      delay_cls.register_entry(self.name, delay)
-      self.delay = delay_cls
-      self.has_delay = True
+    if delay is not None:
+      if isinstance(delay, u.Quantity):
+        has_delay = delay.mantissa > 0.
+      else:
+        has_delay = delay > 0.
+      if has_delay:
+        delay_cls = register_delay_of_target(pre)
+        delay_cls.register_entry(self.name, delay)
+        self.delay = delay_cls
+        self.has_delay = True
+      else:
+        self.delay = None
+        self.has_delay = False
     else:
       self.delay = None
       self.has_delay = False

docs/index.rst

@@ -42,7 +42,7 @@ See also the BDP ecosystem
 ^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 
-We are building the `BDP ecosystem <https://ecosystem-for-brain-dynamics.readthedocs.io/>`_:
+We are building the `BDP ecosystem <https://ecosystem-for-brain-dynamics.readthedocs.io/>`_.