update v1.1.0

2024-10-21 00:35:53 +08:00
parent 7d6451efad
commit 0bba2c319d
148 changed files with 18476 additions and 1356 deletions
@@ -0,0 +1,10 @@
+import torch
+
+def _i(tensor, t, x):
+    """
+    Index tensor using t and format the output according to x.
+    """
+    shape = (x.size(0),) + (1,) * (x.ndim - 1)
+    if isinstance(t, torch.Tensor):
+        t = t.to(tensor.device)
+    return tensor[t].view(shape).to(x.device)