Merge pull request #76 from sidward/main

Fixes to BlocksToArray and ArrayToBlocks
mikgroup · Apr 27, 2021 · 6373e21 · 6373e21
2 parents 5468aee + 920197c
commit 6373e21
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/sigpy/block.py b/sigpy/block.py
@@ -41,7 +41,7 @@ def array_to_blocks(input, blk_shape, blk_strides):
     batch_shape = list(input.shape[:-ndim])
     batch_size = util.prod(batch_shape)
     xp = backend.get_array_module(input)
-    output = xp.zeros([batch_size] + num_blks + blk_shape,
+    output = xp.zeros([batch_size] + num_blks + list(blk_shape),
                       dtype=input.dtype)
     input = input.reshape([batch_size] + list(input.shape[-ndim:]))
 
@@ -102,7 +102,7 @@ def array_to_blocks(input, blk_shape, blk_strides):
     else:
         raise ValueError('Only support ndim=1, 2, or 3, got {}'.format(ndim))
 
-    return output.reshape(batch_shape + num_blks + blk_shape)
+    return output.reshape(batch_shape + num_blks + list(blk_shape))
 
 
 def blocks_to_array(input, oshape, blk_shape, blk_strides):

diff --git a/sigpy/linop.py b/sigpy/linop.py
@@ -1320,7 +1320,7 @@ def __init__(self, ishape, blk_shape, blk_strides):
         D = len(blk_shape)
         num_blks = [(i - b + s) // s for i, b,
                     s in zip(ishape[-D:], blk_shape, blk_strides)]
-        oshape = num_blks + list(blk_shape)
+        oshape = list(ishape[:-D]) + num_blks + list(blk_shape)
 
         super().__init__(oshape, ishape)
 
@@ -1352,7 +1352,7 @@ def __init__(self, oshape, blk_shape, blk_strides):
         D = len(blk_shape)
         num_blks = [(i - b + s) // s for i, b,
                     s in zip(oshape[-D:], blk_shape, blk_strides)]
-        ishape = num_blks + list(blk_shape)
+        ishape = list(oshape[:-D]) + num_blks + list(blk_shape)
 
         super().__init__(oshape, ishape)