fix rotate behavior for ndim > 2 and add test case

python/cucim/src/cucim/skimage/transform/_warps.py

@@ -423,6 +423,10 @@ def rotate(image, angle, resize=False, center=None, order=None,
     symmetric, the result would be [0, 1, 2, 2, 1, 0, 0], while for reflect it
     would be [0, 1, 2, 1, 0, 1, 2].
 
+    If ``image.ndim > 2``, the rotation occurs for the first two dimensions of
+    the array. Unlike the scikit-image implementation, more than one additional
+    axis may be present on the array.
+
     Examples
     --------
     >>> from skimage import data
@@ -479,7 +483,7 @@ def rotate(image, angle, resize=False, center=None, order=None,
         maxr = corners[:, 1].max()
         out_rows = maxr - minr + 1
         out_cols = maxc - minc + 1
-        output_shape = np.around((out_rows, out_cols))
+        output_shape = (round(out_rows), round(out_cols))
 
         # fit output image in new shape
         translation = (minc, minr)
@@ -493,11 +497,21 @@ def rotate(image, angle, resize=False, center=None, order=None,
     tform.params[:2, :2] = tform.params[:2, :2].T
     tform.params[:2, 2] = tform.params[1::-1, 2]
 
+    if image.ndim == 2:
+        affine_params = tform.params
+    elif image.ndim > 2:
+        # note: only the first two dimensions are the ones being rotated
+        #       embed 2D affine into larger identity matrix.
+        affine_params = np.eye(image.ndim + 1)
+        affine_params[:3, :3] = tform.params
+        # keep original shape on the excess dimensions
+        output_shape = output_shape + image.shape[2:]
+
     # transfer the coordinate transform to the GPU
-    tform.params = cp.asarray(tform.params)
+    affine_params = cp.asarray(affine_params)
 
     return _ndimage_affine(
-        image, tform.params, output_shape=output_shape, order=order,
+        image, affine_params, output_shape=output_shape, order=order,
         mode=mode, cval=cval, clip=clip, preserve_range=preserve_range
     )
 

python/cucim/src/cucim/skimage/transform/tests/test_warps.py

@@ -259,6 +259,19 @@ def test_rotate_resize_90():
     assert x90.shape == (230, 470)
 
 
+@pytest.mark.parametrize('dtype', [cp.float16, cp.float32, cp.float64])
+@pytest.mark.parametrize('resize', [False, True])
+@pytest.mark.parametrize('ndim', [2, 3, 4])
+def test_rotate_nd(dtype, resize, ndim):
+    # verify fix for issue: https://github.com/rapidsai/cucim/issues/431
+    x = cp.zeros((5, 5) + (4,) * (ndim - 2), dtype=dtype)
+    x[1, 1] = 1
+    x_out = rotate(x, 15, resize=resize)
+    assert x_out.ndim == x.ndim
+    assert x_out.shape[2:] == x.shape[2:]
+    assert x_out.dtype == _supported_float_type(dtype)
+
+
 def test_rescale():
     # same scale factor
     x = cp.zeros((5, 5), dtype=cp.double)