Make note of new syntax

RadostW · Oct 23, 2024 · 5ee2184 · 5ee2184
1 parent 146f1c4
commit 5ee2184
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pychastic/vectorized_I_generation.py b/pychastic/vectorized_I_generation.py
@@ -140,8 +140,9 @@ def make_D_mat(eta, zeta):
     def take(tensor, idx, fill=0):
         # Non jit-friendly implementation
         # illegal = jnp.logical_or(idx > p,idx < 1)
-        # return tensor[..., idx-1].at[..., illegal].set(fill)
+        # return tensor[..., idx-1].at[..., illegal].set(fill)        
         legalized_idx = jnp.clip(idx, a_min=1, a_max=p)
+        # legalized_idx = jnp.clip(idx, min=1, max=p) # python 3.9+
         illegal_mask = jnp.logical_or(idx > p, idx < 1)
         return (
             tensor[..., legalized_idx - 1] * (1 - 1 * illegal_mask)