[FIX] ensure shapes are lists in TS code -- fixes things when PYTORCH_JIT=0

Author: balbasty

nitorch/_C/_ts/iso0.py

@@ -30,7 +30,7 @@ def pull3d(inp, g, bound: List[Bound], extrapolate: int = 1):
     """
     dim = 3
     boundx, boundy, boundz = bound
-    oshape = g.shape[-dim-1:-1]
+    oshape = list(g.shape[-dim-1:-1])
     g = g.reshape([g.shape[0], 1, -1, dim])
     gx, gy, gz = g.unbind(-1)
     batch = max(inp.shape[0], gx.shape[0])
@@ -47,7 +47,7 @@ def pull3d(inp, g, bound: List[Bound], extrapolate: int = 1):
     gz, signz = get_indices(gz, nz, boundz)
 
     # gather
-    inp = inp.reshape(inp.shape[:2] + [-1])
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
     idx = sub2ind_list([gx, gy, gz], shape)
     idx = idx.expand([batch, channel, idx.shape[-1]])
     out = inp.gather(-1, idx)
@@ -56,7 +56,7 @@ def pull3d(inp, g, bound: List[Bound], extrapolate: int = 1):
         out *= sign
     if mask is not None:
         out *= mask
-    out = out.reshape(out.shape[:2] + oshape)
+    out = out.reshape(list(out.shape[:2]) + oshape)
     return out
 
 
@@ -75,10 +75,10 @@ def push3d(inp, g, shape: Optional[List[int]], bound: List[Bound],
     boundx, boundy, boundz = bound
     if inp.shape[-dim:] != g.shape[-dim-1:-1]:
         raise ValueError('Input and grid should have the same spatial shape')
-    ishape = inp.shape[-dim:]
+    ishape = list(inp.shape[-dim:])
     g = g.reshape([g.shape[0], 1, -1, dim])
     gx, gy, gz = torch.unbind(g, -1)
-    inp = inp.reshape(inp.shape[:2] + [-1])
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
     batch = max(inp.shape[0], gx.shape[0])
     channel = inp.shape[1]
 
@@ -107,7 +107,7 @@ def push3d(inp, g, shape: Optional[List[int]], bound: List[Bound],
         inp *= mask
     out.scatter_add_(-1, idx, inp)
 
-    out = out.reshape(out.shape[:2] + shape)
+    out = out.reshape(list(out.shape[:2]) + shape)
     return out
 
 
@@ -127,12 +127,12 @@ def pull2d(inp, g, bound: List[Bound], extrapolate: int = 1):
     """
     dim = 2
     boundx, boundy = bound
-    oshape = g.shape[-dim-1:-1]
+    oshape = list(g.shape[-dim-1:-1])
     g = g.reshape([g.shape[0], 1, -1, dim])
     gx, gy = g.unbind(-1)
     batch = max(inp.shape[0], gx.shape[0])
     channel = inp.shape[1]
-    shape = inp.shape[-dim:]
+    shape = list(inp.shape[-dim:])
     nx, ny = shape
 
     # mask of inbounds voxels
@@ -171,10 +171,10 @@ def push2d(inp, g, shape: Optional[List[int]], bound: List[Bound],
     boundx, boundy = bound
     if inp.shape[-dim:] != g.shape[-dim-1:-1]:
         raise ValueError('Input and grid should have the same spatial shape')
-    ishape = inp.shape[-dim:]
+    ishape = list(inp.shape[-dim:])
     g = g.reshape([g.shape[0], 1, -1, dim])
     gx, gy = torch.unbind(g, -1)
-    inp = inp.reshape(inp.shape[:2] + [-1])
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
     batch = max(inp.shape[0], gx.shape[0])
     channel = inp.shape[1]
 
@@ -202,7 +202,7 @@ def push2d(inp, g, shape: Optional[List[int]], bound: List[Bound],
         inp *= mask
     out.scatter_add_(-1, idx, inp)
 
-    out = out.reshape(out.shape[:2] + shape)
+    out = out.reshape(list(out.shape[:2]) + shape)
     return out
 
 
@@ -222,12 +222,12 @@ def pull1d(inp, g, bound: List[Bound], extrapolate: int = 1):
     """
     dim = 1
     boundx = bound[0]
-    oshape = g.shape[-dim-1:-1]
+    oshape = list(g.shape[-dim-1:-1])
     g = g.reshape([g.shape[0], 1, -1, dim])
     gx = g.squeeze(-1)
     batch = max(inp.shape[0], gx.shape[0])
     channel = inp.shape[1]
-    shape = inp.shape[-dim:]
+    shape = list(inp.shape[-dim:])
     nx = shape[0]
 
     # mask of inbounds voxels
@@ -237,7 +237,7 @@ def pull1d(inp, g, bound: List[Bound], extrapolate: int = 1):
     gx, signx = get_indices(gx, nx, boundx)
 
     # gather
-    inp = inp.reshape(inp.shape[:2] + [-1])
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
     idx = gx
     idx = idx.expand([batch, channel, idx.shape[-1]])
     out = inp.gather(-1, idx)
@@ -246,7 +246,7 @@ def pull1d(inp, g, bound: List[Bound], extrapolate: int = 1):
         out *= sign
     if mask is not None:
         out *= mask
-    out = out.reshape(out.shape[:2] + oshape)
+    out = out.reshape(list(out.shape[:2]) + oshape)
     return out
 
 
@@ -265,10 +265,10 @@ def push1d(inp, g, shape: Optional[List[int]], bound: List[Bound],
     boundx = bound[0]
     if inp.shape[-dim:] != g.shape[-dim-1:-1]:
         raise ValueError('Input and grid should have the same spatial shape')
-    ishape = inp.shape[-dim:]
+    ishape = list(inp.shape[-dim:])
     g = g.reshape([g.shape[0], 1, -1, dim])
     gx = g.squeeze(-1)
-    inp = inp.reshape(inp.shape[:2] + [-1])
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
     batch = max(inp.shape[0], gx.shape[0])
     channel = inp.shape[1]
 
@@ -295,7 +295,7 @@ def push1d(inp, g, shape: Optional[List[int]], bound: List[Bound],
         inp *= mask
     out.scatter_add_(-1, idx, inp)
 
-    out = out.reshape(out.shape[:2] + shape)
+    out = out.reshape(list(out.shape[:2]) + shape)
     return out
 
 
@@ -336,7 +336,7 @@ def pushgrad(inp, g, shape: Optional[List[int]], bound: List[Bound],
     dim = g.shape[-1]
     if inp.shape[-dim-1:-1] != g.shape[-dim-1:-1]:
         raise ValueError('Input and grid should have the same spatial shape')
-    ishape = inp.shape[-dim-1:-1]
+    ishape = list(inp.shape[-dim-1:-1])
     batch = max(inp.shape[0], g.shape[0])
     channel = inp.shape[1]
 

nitorch/_C/_ts/nd.py

@@ -440,4 +440,4 @@ def hess(inp, grid, bound: List[Bound], spline: List[Spline],
             out.unbind(-1)[d2].unbind(-1)[d].copy_(out.unbind(-1)[d].unbind(-1)[d2])
 
     out = out.reshape(list(out.shape[:2]) + oshape + list(out.shape[-2:]))
-    return out
+    return out