Merge pull request #170 from chaoming0625/master

synapse models support heterogeneuos weights

brainpy/dyn/base.py

@@ -483,16 +483,19 @@ def register_delay(
   def get_delay(
       self,
       name: str,
-      delay_step: Union[int, bm.JaxArray, bm.ndarray]
+      delay_step: Union[int, bm.JaxArray, bm.ndarray],
+      indices=None,
   ):
-    """Get delay data according to the delay times.
+    """Get delay data according to the provided delay steps.
 
     Parameters
     ----------
     name: str
       The delay variable name.
     delay_step: int, JaxArray, ndarray
       The delay length.
+    indices: optional, JaxArray, ndarray
+      The indices of the delay.
 
     Returns
     -------
@@ -501,14 +504,18 @@ def get_delay(
     """
     if name in self.global_delay_vars:
       if isinstance(delay_step, int):
-        return self.global_delay_vars[name](delay_step)
+        return self.global_delay_vars[name](delay_step, indices)
       else:
-        return self.global_delay_vars[name](delay_step, jnp.arange(delay_step.size))
+        if indices is None:
+          indices = jnp.arange(delay_step.size)
+        return self.global_delay_vars[name](delay_step, indices)
     elif name in self.local_delay_vars:
       if isinstance(delay_step, int):
         return self.local_delay_vars[name](delay_step)
       else:
-        return self.local_delay_vars[name](delay_step, jnp.arange(delay_step.size))
+        if indices is None:
+          indices = jnp.arange(delay_step.size)
+        return self.local_delay_vars[name](delay_step, indices)
     else:
       raise ValueError(f'{name} is not defined in delay variables.')