From 0be42ae466ae34a0cfa01ccabe262ea0958c30fb Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Wed, 1 Sep 2021 14:12:29 +0000
Subject: [PATCH 01/14] Implement SplitOnGrid

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 .../pathology/transforms/spatial/array.py     | 71 +++++++++++++++++++
 1 file changed, 71 insertions(+)
 create mode 100644 monai/apps/pathology/transforms/spatial/array.py

diff --git a/monai/apps/pathology/transforms/spatial/array.py b/monai/apps/pathology/transforms/spatial/array.py
new file mode 100644
index 0000000000..7057605cd9
--- /dev/null
+++ b/monai/apps/pathology/transforms/spatial/array.py
@@ -0,0 +1,71 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING, Optional, Tuple, Union
+
+from torch.nn import Module
+
+from monai.transforms.transform import Transform
+
+if TYPE_CHECKING:
+    import torch
+
+__all__ = ["SplitOnGrid"]
+
+
+class SplitOnGrid(Transform):
+    """
+    Split the image into patches based on the provided grid shape.
+    This transform works only with torch.Tensor inputs.
+
+    Args:
+        grid_shape: a tuple or an integer define the shape of the grid upon which to extract patches.
+            If it's an integer, th evalue will be repeated for each dimension. Default is 2x2
+        patch_size: a tuple or an integer that defines the output patch sizes.
+            If it's an integer, the value will be repeated for each dimension.
+            If None (default), the patch size will be infered from the grid shape.
+    """
+
+    def __init__(
+        self,
+        grid_shape: Union[int, Tuple[int, int]] = (2, 2),
+        patch_size: Optional[Union[int, Tuple[int, int]]] = None,
+    ):
+        if isinstance(grid_shape, int):
+            self.grid_shape = (grid_shape, grid_shape)
+        else:
+            self.grid_shape = grid_shape
+        self.patch_size = None
+        if isinstance(patch_size, int):
+            self.patch_size = (patch_size, patch_size)
+        else:
+            self.patch_size = patch_size
+
+    def __call__(self, region: torch.Tensor) -> torch.Tensor:
+        _, h, w = region.shape
+        if self.patch_size is None:
+            if self.grid_shape == (1, 1):
+                return region
+            else:
+                self.patch_size = (h // self.grid_shape[0], w // self.grid_shape[1])
+
+        h_stride = (h - self.patch_size[0]) // (self.grid_shape[0] - 1)
+        w_stride = (w - self.patch_size[1]) // (self.grid_shape[1] - 1)
+
+        patches = (
+            region.unfold(1, self.patch_size[0], h_stride)
+            .unfold(2, self.patch_size[1], w_stride)
+            .flatten(1, 2)
+            .transpose(0, 1)
+            .contiguous()
+        )
+
+        return patches

From 40e0348ba6d3e59fb6e65c24810578d0cd065d30 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Wed, 1 Sep 2021 14:12:38 +0000
Subject: [PATCH 02/14] Implement dictionary-based SplitOnGrid

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 .../transforms/spatial/dictionary.py          | 55 +++++++++++++++++++
 1 file changed, 55 insertions(+)
 create mode 100644 monai/apps/pathology/transforms/spatial/dictionary.py

diff --git a/monai/apps/pathology/transforms/spatial/dictionary.py b/monai/apps/pathology/transforms/spatial/dictionary.py
new file mode 100644
index 0000000000..e20ee6bd9f
--- /dev/null
+++ b/monai/apps/pathology/transforms/spatial/dictionary.py
@@ -0,0 +1,55 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING, Dict, Hashable, Mapping, Optional, Tuple, Union
+
+from monai.config import KeysCollection
+from monai.transforms.transform import MapTransform
+
+from .array import SplitOnGrid
+
+if TYPE_CHECKING:
+    import torch
+
+__all__ = ["SplitOnGridd", "SplitOnGridD", "SplitOnGridDict"]
+
+
+class SplitOnGridd(MapTransform):
+    """
+    Split the image into patches based on the provided grid shape.
+    This transform works only with torch.Tensor inputs.
+
+    Args:
+        grid_shape: a tuple or an integer define the shape of the grid upon which to extract patches.
+            If it's an integer, th evalue will be repeated for each dimension. Default is 2x2
+        patch_size: a tuple or an integer that defines the output patch sizes.
+            If it's an integer, the value will be repeated for each dimension.
+            If None (default), the patch size will be infered from the grid shape.
+    """
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        grid_shape: Union[int, Tuple[int, int]] = (2, 2),
+        patch_size: Optional[Union[int, Tuple[int, int]]] = None,
+        allow_missing_keys: bool = False,
+    ):
+        super().__init__(keys, allow_missing_keys)
+        self.splitter = SplitOnGrid(grid_shape=grid_shape, patch_size=patch_size)
+
+    def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]:
+        d = dict(data)
+        for key in self.key_iterator(d):
+            d[key] = self.splitter(d[key])
+        return d
+
+
+SplitOnGridDict = SplitOnGridD = SplitOnGridd

From 7b4fcd2d73374472d03ab29eaf3d26ee67b76853 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Wed, 1 Sep 2021 14:13:05 +0000
Subject: [PATCH 03/14] Update inits

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 monai/apps/pathology/transforms/__init__.py         |  3 +++
 monai/apps/pathology/transforms/spatial/__init__.py | 13 +++++++++++++
 2 files changed, 16 insertions(+)
 create mode 100644 monai/apps/pathology/transforms/spatial/__init__.py

diff --git a/monai/apps/pathology/transforms/__init__.py b/monai/apps/pathology/transforms/__init__.py
index 0df016244b..8ff7edd976 100644
--- a/monai/apps/pathology/transforms/__init__.py
+++ b/monai/apps/pathology/transforms/__init__.py
@@ -18,3 +18,6 @@
     NormalizeHEStainsD,
     NormalizeHEStainsDict,
 )
+
+from .spatial.array import SplitOnGrid
+from .spatial.dictionary import SplitOnGrid, SplitOnGridD, SplitOnGridDict
diff --git a/monai/apps/pathology/transforms/spatial/__init__.py b/monai/apps/pathology/transforms/spatial/__init__.py
new file mode 100644
index 0000000000..07ba222ab0
--- /dev/null
+++ b/monai/apps/pathology/transforms/spatial/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .array import SplitOnGrid
+from .dictionary import SplitOnGridd, SplitOnGridD, SplitOnGridDict

From 15ee5d67c095068ece3a8b3e5b5332150ef21749 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Wed, 1 Sep 2021 14:28:10 +0000
Subject: [PATCH 04/14] Update docs

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 docs/source/apps.rst | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/docs/source/apps.rst b/docs/source/apps.rst
index 1a2efeff48..11d60767ec 100644
--- a/docs/source/apps.rst
+++ b/docs/source/apps.rst
@@ -110,3 +110,11 @@ Clara MMARs
     :members:
 .. autoclass:: NormalizeHEStainsd
     :members:
+
+.. automodule:: monai.apps.pathology.transforms.spatial.array
+.. autoclass:: SplitOnGrid
+    :members:
+
+.. automodule:: monai.apps.pathology.transforms.spatial.dictionary
+.. autoclass:: SplitOnGridd
+    :members:

From b33da39107b1f2387bbad3062e605511a13ac095 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Wed, 1 Sep 2021 14:28:28 +0000
Subject: [PATCH 05/14] Change imports

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 monai/apps/pathology/transforms/spatial/array.py      | 7 ++-----
 monai/apps/pathology/transforms/spatial/dictionary.py | 7 +++----
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/monai/apps/pathology/transforms/spatial/array.py b/monai/apps/pathology/transforms/spatial/array.py
index 7057605cd9..8c8fac1e00 100644
--- a/monai/apps/pathology/transforms/spatial/array.py
+++ b/monai/apps/pathology/transforms/spatial/array.py
@@ -9,15 +9,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 
-from torch.nn import Module
+import torch
 
 from monai.transforms.transform import Transform
 
-if TYPE_CHECKING:
-    import torch
-
 __all__ = ["SplitOnGrid"]
 
 
diff --git a/monai/apps/pathology/transforms/spatial/dictionary.py b/monai/apps/pathology/transforms/spatial/dictionary.py
index e20ee6bd9f..75944e833c 100644
--- a/monai/apps/pathology/transforms/spatial/dictionary.py
+++ b/monai/apps/pathology/transforms/spatial/dictionary.py
@@ -9,16 +9,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Dict, Hashable, Mapping, Optional, Tuple, Union
+from typing import Dict, Hashable, Mapping, Optional, Tuple, Union
+
+import torch
 
 from monai.config import KeysCollection
 from monai.transforms.transform import MapTransform
 
 from .array import SplitOnGrid
 
-if TYPE_CHECKING:
-    import torch
-
 __all__ = ["SplitOnGridd", "SplitOnGridD", "SplitOnGridDict"]
 
 

From e796f46857c5fee218cae678021f42c903a46156 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Wed, 1 Sep 2021 21:03:57 +0000
Subject: [PATCH 06/14] Update input logic in SplitOnGrid)

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 .../pathology/transforms/spatial/array.py     | 55 ++++++++++++-------
 .../transforms/spatial/dictionary.py          |  6 +-
 2 files changed, 38 insertions(+), 23 deletions(-)

diff --git a/monai/apps/pathology/transforms/spatial/array.py b/monai/apps/pathology/transforms/spatial/array.py
index 8c8fac1e00..cd54f2bd05 100644
--- a/monai/apps/pathology/transforms/spatial/array.py
+++ b/monai/apps/pathology/transforms/spatial/array.py
@@ -29,40 +29,53 @@ class SplitOnGrid(Transform):
         patch_size: a tuple or an integer that defines the output patch sizes.
             If it's an integer, the value will be repeated for each dimension.
             If None (default), the patch size will be infered from the grid shape.
+
+    Note: the shape of the input image is infered based on the first image used.
     """
 
     def __init__(
-        self,
-        grid_shape: Union[int, Tuple[int, int]] = (2, 2),
-        patch_size: Optional[Union[int, Tuple[int, int]]] = None,
+        self, grid_size: Union[int, Tuple[int, int]] = (2, 2), patch_size: Optional[Union[int, Tuple[int, int]]] = None
     ):
-        if isinstance(grid_shape, int):
-            self.grid_shape = (grid_shape, grid_shape)
+        # Grid size
+        if isinstance(grid_size, int):
+            self.grid_size = (grid_size, grid_size)
         else:
-            self.grid_shape = grid_shape
-        self.patch_size = None
+            self.grid_size = grid_size
+        # Patch size
         if isinstance(patch_size, int):
             self.patch_size = (patch_size, patch_size)
+        elif patch_size is None:
+            self.patch_size = (0, 0)
         else:
             self.patch_size = patch_size
+        # Set steps to a default to be overriden
+        self.steps = (0, 0)
+        self.ready = False
+        # Set skip flags to bypass if the input and output should be the same.
+        self.skip = False
+        if self.grid_size == (1, 1) and self.patch_size == (0, 0):
+            self.skip = True
 
-    def __call__(self, region: torch.Tensor) -> torch.Tensor:
-        _, h, w = region.shape
-        if self.patch_size is None:
-            if self.grid_shape == (1, 1):
-                return region
-            else:
-                self.patch_size = (h // self.grid_shape[0], w // self.grid_shape[1])
-
-        h_stride = (h - self.patch_size[0]) // (self.grid_shape[0] - 1)
-        w_stride = (w - self.patch_size[1]) // (self.grid_shape[1] - 1)
-
+    def __call__(self, image: torch.Tensor) -> torch.Tensor:
+        if self.skip:
+            return torch.stack([image])
+        if not self.ready:
+            self.prepare_params(image.shape[1:])
         patches = (
-            region.unfold(1, self.patch_size[0], h_stride)
-            .unfold(2, self.patch_size[1], w_stride)
+            image.unfold(1, self.patch_size[0], self.steps[0])
+            .unfold(2, self.patch_size[1], self.steps[1])
             .flatten(1, 2)
             .transpose(0, 1)
             .contiguous()
         )
-
         return patches
+
+    def prepare_params(self, image_size):
+        if self.patch_size == (0, 0):
+            self.patch_size = tuple(image_size[i] // self.grid_size[i] for i in range(2))
+
+        self.steps = tuple(
+            (image_size[i] - self.patch_size[i]) // (self.grid_size[i] - 1) if self.grid_size[i] > 1 else image_size[i]
+            for i in range(2)
+        )
+        self.ready = True
diff --git a/monai/apps/pathology/transforms/spatial/dictionary.py b/monai/apps/pathology/transforms/spatial/dictionary.py
index 75944e833c..e5782e36a4 100644
--- a/monai/apps/pathology/transforms/spatial/dictionary.py
+++ b/monai/apps/pathology/transforms/spatial/dictionary.py
@@ -32,17 +32,19 @@ class SplitOnGridd(MapTransform):
         patch_size: a tuple or an integer that defines the output patch sizes.
             If it's an integer, the value will be repeated for each dimension.
             If None (default), the patch size will be infered from the grid shape.
+
+    Note: the shape of the input image is infered based on the first image used.
     """
 
     def __init__(
         self,
         keys: KeysCollection,
-        grid_shape: Union[int, Tuple[int, int]] = (2, 2),
+        grid_size: Union[int, Tuple[int, int]] = (2, 2),
         patch_size: Optional[Union[int, Tuple[int, int]]] = None,
         allow_missing_keys: bool = False,
     ):
         super().__init__(keys, allow_missing_keys)
-        self.splitter = SplitOnGrid(grid_shape=grid_shape, patch_size=patch_size)
+        self.splitter = SplitOnGrid(grid_size=grid_size, patch_size=patch_size)
 
     def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]:
         d = dict(data)

From b204a872240cc3dba25cc2ad0fc728a2bdeedb75 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Wed, 1 Sep 2021 21:34:19 +0000
Subject: [PATCH 07/14] Add unittests for SplitOnGrid and SplitOnGridDict

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 tests/test_split_on_grid.py      | 134 +++++++++++++++++++++++++++++++
 tests/test_split_on_grid_dict.py | 134 +++++++++++++++++++++++++++++++
 2 files changed, 268 insertions(+)
 create mode 100644 tests/test_split_on_grid.py
 create mode 100644 tests/test_split_on_grid_dict.py

diff --git a/tests/test_split_on_grid.py b/tests/test_split_on_grid.py
new file mode 100644
index 0000000000..4415a9ed6a
--- /dev/null
+++ b/tests/test_split_on_grid.py
@@ -0,0 +1,134 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import unittest
+from unittest import skipUnless
+
+import numpy as np
+import torch
+from numpy.testing import assert_array_equal
+from parameterized import parameterized
+
+from monai.apps.pathology.transforms import SplitOnGrid
+
+A11 = torch.randn(3, 2, 2)
+A12 = torch.randn(3, 2, 2)
+A21 = torch.randn(3, 2, 2)
+A22 = torch.randn(3, 2, 2)
+
+A1 = torch.cat([A11, A12], 2)
+A2 = torch.cat([A21, A22], 2)
+A = torch.cat([A1, A2], 1)
+
+TEST_CASE_0 = [
+    {"grid_size": (2, 2)},
+    A,
+    torch.stack([A11, A12, A21, A22]),
+]
+
+TEST_CASE_1 = [
+    {"grid_size": (2, 1)},
+    A,
+    torch.stack([A1, A2]),
+]
+
+TEST_CASE_2 = [
+    {"grid_size": (1, 2)},
+    A1,
+    torch.stack([A11, A12]),
+]
+
+TEST_CASE_3 = [
+    {"grid_size": (1, 2)},
+    A2,
+    torch.stack([A21, A22]),
+]
+
+TEST_CASE_4 = [
+    {"grid_size": (1, 1), "patch_size": (2, 2)},
+    A,
+    torch.stack([A11]),
+]
+
+TEST_CASE_5 = [
+    {"grid_size": 1, "patch_size": 4},
+    A,
+    torch.stack([A]),
+]
+
+TEST_CASE_6 = [
+    {"grid_size": 2, "patch_size": 2},
+    A,
+    torch.stack([A11, A12, A21, A22]),
+]
+
+TEST_CASE_7 = [
+    {"grid_size": 1},
+    A,
+    torch.stack([A]),
+]
+
+TEST_CASE_MC_0 = [
+    {"grid_size": (2, 2)},
+    [A, A],
+    [torch.stack([A11, A12, A21, A22]), torch.stack([A11, A12, A21, A22])],
+]
+
+
+TEST_CASE_MC_1 = [
+    {"grid_size": (2, 1)},
+    [A] * 5,
+    [torch.stack([A1, A2])] * 5,
+]
+
+
+TEST_CASE_MC_2 = [
+    {"grid_size": (1, 2)},
+    [A1, A2],
+    [torch.stack([A11, A12]), torch.stack([A21, A22])],
+]
+
+
+class TestSplitOnGrid(unittest.TestCase):
+    @parameterized.expand(
+        [
+            TEST_CASE_0,
+            TEST_CASE_1,
+            TEST_CASE_2,
+            TEST_CASE_3,
+            TEST_CASE_4,
+            TEST_CASE_5,
+            TEST_CASE_6,
+            TEST_CASE_7,
+        ]
+    )
+    def test_split_pathce_single_call(self, input_parameters, img, expected):
+        splitter = SplitOnGrid(**input_parameters)
+        output = splitter(img)
+        np.testing.assert_equal(output.numpy(), expected.numpy())
+
+    @parameterized.expand(
+        [
+            TEST_CASE_MC_0,
+            TEST_CASE_MC_1,
+            TEST_CASE_MC_2,
+        ]
+    )
+    def test_split_pathce_multiple_call(self, input_parameters, img_list, expected_list):
+        splitter = SplitOnGrid(**input_parameters)
+        for img, expected in zip(img_list, expected_list):
+            output = splitter(img)
+            np.testing.assert_equal(output.numpy(), expected.numpy())
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_split_on_grid_dict.py b/tests/test_split_on_grid_dict.py
new file mode 100644
index 0000000000..d6b96b0b74
--- /dev/null
+++ b/tests/test_split_on_grid_dict.py
@@ -0,0 +1,134 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import unittest
+from unittest import skipUnless
+
+import numpy as np
+import torch
+from numpy.testing import assert_array_equal
+from parameterized import parameterized
+
+from monai.apps.pathology.transforms import SplitOnGridDict
+
+A11 = torch.randn(3, 2, 2)
+A12 = torch.randn(3, 2, 2)
+A21 = torch.randn(3, 2, 2)
+A22 = torch.randn(3, 2, 2)
+
+A1 = torch.cat([A11, A12], 2)
+A2 = torch.cat([A21, A22], 2)
+A = torch.cat([A1, A2], 1)
+
+TEST_CASE_0 = [
+    {"keys": "image", "grid_size": (2, 2)},
+    {"image": A},
+    torch.stack([A11, A12, A21, A22]),
+]
+
+TEST_CASE_1 = [
+    {"keys": "image", "grid_size": (2, 1)},
+    {"image": A},
+    torch.stack([A1, A2]),
+]
+
+TEST_CASE_2 = [
+    {"keys": "image", "grid_size": (1, 2)},
+    {"image": A1},
+    torch.stack([A11, A12]),
+]
+
+TEST_CASE_3 = [
+    {"keys": "image", "grid_size": (1, 2)},
+    {"image": A2},
+    torch.stack([A21, A22]),
+]
+
+TEST_CASE_4 = [
+    {"keys": "image", "grid_size": (1, 1), "patch_size": (2, 2)},
+    {"image": A},
+    torch.stack([A11]),
+]
+
+TEST_CASE_5 = [
+    {"keys": "image", "grid_size": 1, "patch_size": 4},
+    {"image": A},
+    torch.stack([A]),
+]
+
+TEST_CASE_6 = [
+    {"keys": "image", "grid_size": 2, "patch_size": 2},
+    {"image": A},
+    torch.stack([A11, A12, A21, A22]),
+]
+
+TEST_CASE_7 = [
+    {"keys": "image", "grid_size": 1},
+    {"image": A},
+    torch.stack([A]),
+]
+
+TEST_CASE_MC_0 = [
+    {"keys": "image", "grid_size": (2, 2)},
+    [{"image": A}, {"image": A}],
+    [torch.stack([A11, A12, A21, A22]), torch.stack([A11, A12, A21, A22])],
+]
+
+
+TEST_CASE_MC_1 = [
+    {"keys": "image", "grid_size": (2, 1)},
+    [{"image": A}] * 5,
+    [torch.stack([A1, A2])] * 5,
+]
+
+
+TEST_CASE_MC_2 = [
+    {"keys": "image", "grid_size": (1, 2)},
+    [{"image": A1}, {"image": A2}],
+    [torch.stack([A11, A12]), torch.stack([A21, A22])],
+]
+
+
+class TestSplitOnGridDict(unittest.TestCase):
+    @parameterized.expand(
+        [
+            TEST_CASE_0,
+            TEST_CASE_1,
+            TEST_CASE_2,
+            TEST_CASE_3,
+            TEST_CASE_4,
+            TEST_CASE_5,
+            TEST_CASE_6,
+            TEST_CASE_7,
+        ]
+    )
+    def test_split_pathce_single_call(self, input_parameters, img_dict, expected):
+        splitter = SplitOnGridDict(**input_parameters)
+        output = splitter(img_dict)[input_parameters["keys"]]
+        np.testing.assert_equal(output.numpy(), expected.numpy())
+
+    @parameterized.expand(
+        [
+            TEST_CASE_MC_0,
+            TEST_CASE_MC_1,
+            TEST_CASE_MC_2,
+        ]
+    )
+    def test_split_pathce_multiple_call(self, input_parameters, img_list, expected_list):
+        splitter = SplitOnGridDict(**input_parameters)
+        for img_dict, expected in zip(img_list, expected_list):
+            output = splitter(img_dict)[input_parameters["keys"]]
+            np.testing.assert_equal(output.numpy(), expected.numpy())
+
+
+if __name__ == "__main__":
+    unittest.main()

From 6af67d7a32c0f9e543a7178e7ac59ff69e7885fd Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Wed, 1 Sep 2021 21:35:25 +0000
Subject: [PATCH 08/14] Sort import

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 monai/apps/pathology/transforms/__init__.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/monai/apps/pathology/transforms/__init__.py b/monai/apps/pathology/transforms/__init__.py
index 8ff7edd976..fe9c18447d 100644
--- a/monai/apps/pathology/transforms/__init__.py
+++ b/monai/apps/pathology/transforms/__init__.py
@@ -9,6 +9,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .spatial.array import SplitOnGrid
+from .spatial.dictionary import SplitOnGrid, SplitOnGridD, SplitOnGridDict
 from .stain.array import ExtractHEStains, NormalizeHEStains
 from .stain.dictionary import (
     ExtractHEStainsd,
@@ -18,6 +20,3 @@
     NormalizeHEStainsD,
     NormalizeHEStainsDict,
 )
-
-from .spatial.array import SplitOnGrid
-from .spatial.dictionary import SplitOnGrid, SplitOnGridD, SplitOnGridDict

From efc0e1c939ce37b4bfe28ad5143c7ac202016e6b Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Wed, 1 Sep 2021 22:23:59 +0000
Subject: [PATCH 09/14] Remove imports

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 tests/test_split_on_grid.py      | 3 ---
 tests/test_split_on_grid_dict.py | 3 ---
 2 files changed, 6 deletions(-)

diff --git a/tests/test_split_on_grid.py b/tests/test_split_on_grid.py
index 4415a9ed6a..a187835e7b 100644
--- a/tests/test_split_on_grid.py
+++ b/tests/test_split_on_grid.py
@@ -9,13 +9,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import unittest
-from unittest import skipUnless
 
 import numpy as np
 import torch
-from numpy.testing import assert_array_equal
 from parameterized import parameterized
 
 from monai.apps.pathology.transforms import SplitOnGrid
diff --git a/tests/test_split_on_grid_dict.py b/tests/test_split_on_grid_dict.py
index d6b96b0b74..96ec095423 100644
--- a/tests/test_split_on_grid_dict.py
+++ b/tests/test_split_on_grid_dict.py
@@ -9,13 +9,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import unittest
-from unittest import skipUnless
 
 import numpy as np
 import torch
-from numpy.testing import assert_array_equal
 from parameterized import parameterized
 
 from monai.apps.pathology.transforms import SplitOnGridDict

From 11ba76b1feca7894ad47dec68c6fd306e2dabab4 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Thu, 2 Sep 2021 15:21:46 +0000
Subject: [PATCH 10/14] Address comments

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 monai/apps/pathology/transforms/__init__.py    |  2 +-
 .../apps/pathology/transforms/spatial/array.py | 18 ++++++------------
 .../pathology/transforms/spatial/dictionary.py |  6 +++---
 3 files changed, 10 insertions(+), 16 deletions(-)

diff --git a/monai/apps/pathology/transforms/__init__.py b/monai/apps/pathology/transforms/__init__.py
index fe9c18447d..1be96b8e34 100644
--- a/monai/apps/pathology/transforms/__init__.py
+++ b/monai/apps/pathology/transforms/__init__.py
@@ -10,7 +10,7 @@
 # limitations under the License.
 
 from .spatial.array import SplitOnGrid
-from .spatial.dictionary import SplitOnGrid, SplitOnGridD, SplitOnGridDict
+from .spatial.dictionary import SplitOnGridd, SplitOnGridD, SplitOnGridDict
 from .stain.array import ExtractHEStains, NormalizeHEStains
 from .stain.dictionary import (
     ExtractHEStainsd,
diff --git a/monai/apps/pathology/transforms/spatial/array.py b/monai/apps/pathology/transforms/spatial/array.py
index cd54f2bd05..9e30659986 100644
--- a/monai/apps/pathology/transforms/spatial/array.py
+++ b/monai/apps/pathology/transforms/spatial/array.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple, Union
+from typing import Tuple, Union
 
 import torch
 
@@ -25,16 +25,16 @@ class SplitOnGrid(Transform):
 
     Args:
         grid_shape: a tuple or an integer define the shape of the grid upon which to extract patches.
-            If it's an integer, th evalue will be repeated for each dimension. Default is 2x2
+            If it's an integer, the value will be repeated for each dimension. Default is 2x2
         patch_size: a tuple or an integer that defines the output patch sizes.
             If it's an integer, the value will be repeated for each dimension.
-            If None (default), the patch size will be infered from the grid shape.
+            The default is (0, 0), where the patch size will be infered from the grid shape.
 
     Note: the shape of the input image is infered based on the first image used.
     """
 
     def __init__(
-        self, grid_size: Union[int, Tuple[int, int]] = (2, 2), patch_size: Optional[Union[int, Tuple[int, int]]] = None
+        self, grid_size: Union[int, Tuple[int, int]] = (2, 2), patch_size: Union[int, Tuple[int, int]] = (0, 0)
     ):
         # Grid size
         if isinstance(grid_size, int):
@@ -44,23 +44,18 @@ def __init__(
         # Patch size
         if isinstance(patch_size, int):
             self.patch_size = (patch_size, patch_size)
-        elif patch_size is None:
-            self.patch_size = (0, 0)
         else:
             self.patch_size = patch_size
         # Set steps to a default to be overriden
         self.steps = (0, 0)
         self.ready = False
-        # Set skip flags to bypass if the input and output should be the same.
-        self.skip = False
-        if self.grid_size == (1, 1) and self.patch_size == (0, 0):
-            self.skip = True
 
     def __call__(self, image: torch.Tensor) -> torch.Tensor:
-        if self.skip:
+        if self.grid_size == (1, 1) and self.patch_size == (0, 0):
             return torch.stack([image])
         if not self.ready:
             self.prepare_params(image.shape[1:])
+            self.ready = True
         patches = (
             image.unfold(1, self.patch_size[0], self.steps[0])
             .unfold(2, self.patch_size[1], self.steps[1])
@@ -78,4 +73,3 @@ def prepare_params(self, image_size):
             (image_size[i] - self.patch_size[i]) // (self.grid_size[i] - 1) if self.grid_size[i] > 1 else image_size[i]
             for i in range(2)
         )
-        self.ready = True
diff --git a/monai/apps/pathology/transforms/spatial/dictionary.py b/monai/apps/pathology/transforms/spatial/dictionary.py
index e5782e36a4..c550ce76fd 100644
--- a/monai/apps/pathology/transforms/spatial/dictionary.py
+++ b/monai/apps/pathology/transforms/spatial/dictionary.py
@@ -28,10 +28,10 @@ class SplitOnGridd(MapTransform):
 
     Args:
         grid_shape: a tuple or an integer define the shape of the grid upon which to extract patches.
-            If it's an integer, th evalue will be repeated for each dimension. Default is 2x2
+            If it's an integer, the value will be repeated for each dimension. Default is 2x2
         patch_size: a tuple or an integer that defines the output patch sizes.
             If it's an integer, the value will be repeated for each dimension.
-            If None (default), the patch size will be infered from the grid shape.
+            The default is (0, 0), where the patch size will be infered from the grid shape.
 
     Note: the shape of the input image is infered based on the first image used.
     """
@@ -40,7 +40,7 @@ def __init__(
         self,
         keys: KeysCollection,
         grid_size: Union[int, Tuple[int, int]] = (2, 2),
-        patch_size: Optional[Union[int, Tuple[int, int]]] = None,
+        patch_size: Union[int, Tuple[int, int]] = (0, 0),
         allow_missing_keys: bool = False,
     ):
         super().__init__(keys, allow_missing_keys)

From d51a6a495ecfc8d0dcc6ddc63dbe8397ddf3a853 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Thu, 2 Sep 2021 15:41:27 +0000
Subject: [PATCH 11/14] Remove optional

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 monai/apps/pathology/transforms/spatial/dictionary.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/monai/apps/pathology/transforms/spatial/dictionary.py b/monai/apps/pathology/transforms/spatial/dictionary.py
index c550ce76fd..306558a6ae 100644
--- a/monai/apps/pathology/transforms/spatial/dictionary.py
+++ b/monai/apps/pathology/transforms/spatial/dictionary.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Hashable, Mapping, Optional, Tuple, Union
+from typing import Dict, Hashable, Mapping, Tuple, Union
 
 import torch
 

From fc58668b29376f69e95ecd2ab5d6830dfbf88007 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Fri, 3 Sep 2021 13:25:29 +0000
Subject: [PATCH 12/14] Address thread safety issues

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 .../pathology/transforms/spatial/array.py     | 34 ++++++++++---------
 .../transforms/spatial/dictionary.py          |  4 +--
 2 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/monai/apps/pathology/transforms/spatial/array.py b/monai/apps/pathology/transforms/spatial/array.py
index 9e30659986..53e0c63715 100644
--- a/monai/apps/pathology/transforms/spatial/array.py
+++ b/monai/apps/pathology/transforms/spatial/array.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Tuple, Union
+from typing import Optional, Tuple, Union
 
 import torch
 
@@ -34,7 +34,9 @@ class SplitOnGrid(Transform):
     """
 
     def __init__(
-        self, grid_size: Union[int, Tuple[int, int]] = (2, 2), patch_size: Union[int, Tuple[int, int]] = (0, 0)
+        self,
+        grid_size: Union[int, Tuple[int, int]] = (2, 2),
+        patch_size: Optional[Union[int, Tuple[int, int]]] = None,
     ):
         # Grid size
         if isinstance(grid_size, int):
@@ -42,34 +44,34 @@ def __init__(
         else:
             self.grid_size = grid_size
         # Patch size
+        self.patch_size = None
         if isinstance(patch_size, int):
             self.patch_size = (patch_size, patch_size)
         else:
             self.patch_size = patch_size
-        # Set steps to a default to be overriden
-        self.steps = (0, 0)
-        self.ready = False
 
     def __call__(self, image: torch.Tensor) -> torch.Tensor:
-        if self.grid_size == (1, 1) and self.patch_size == (0, 0):
+        if self.grid_size == (1, 1) and self.patch_size is None:
             return torch.stack([image])
-        if not self.ready:
-            self.prepare_params(image.shape[1:])
-            self.ready = True
+        patch_size, steps = self.get_params(image.shape[1:])
         patches = (
-            image.unfold(1, self.patch_size[0], self.steps[0])
-            .unfold(2, self.patch_size[1], self.steps[1])
+            image.unfold(1, patch_size[0], steps[0])
+            .unfold(2, patch_size[1], steps[1])
             .flatten(1, 2)
             .transpose(0, 1)
             .contiguous()
         )
         return patches
 
-    def prepare_params(self, image_size):
-        if self.patch_size == (0, 0):
-            self.patch_size = tuple(image_size[i] // self.grid_size[i] for i in range(2))
+    def get_params(self, image_size):
+        if self.patch_size is None:
+            patch_size = tuple(image_size[i] // self.grid_size[i] for i in range(2))
+        else:
+            patch_size = self.patch_size
 
-        self.steps = tuple(
-            (image_size[i] - self.patch_size[i]) // (self.grid_size[i] - 1) if self.grid_size[i] > 1 else image_size[i]
+        steps = tuple(
+            (image_size[i] - patch_size[i]) // (self.grid_size[i] - 1) if self.grid_size[i] > 1 else image_size[i]
             for i in range(2)
         )
+
+        return patch_size, steps
diff --git a/monai/apps/pathology/transforms/spatial/dictionary.py b/monai/apps/pathology/transforms/spatial/dictionary.py
index 306558a6ae..10b01a39de 100644
--- a/monai/apps/pathology/transforms/spatial/dictionary.py
+++ b/monai/apps/pathology/transforms/spatial/dictionary.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Hashable, Mapping, Tuple, Union
+from typing import Dict, Hashable, Mapping, Optional, Tuple, Union
 
 import torch
 
@@ -40,7 +40,7 @@ def __init__(
         self,
         keys: KeysCollection,
         grid_size: Union[int, Tuple[int, int]] = (2, 2),
-        patch_size: Union[int, Tuple[int, int]] = (0, 0),
+        patch_size: Optional[Union[int, Tuple[int, int]]] = None,
         allow_missing_keys: bool = False,
     ):
         super().__init__(keys, allow_missing_keys)

From 76e8c9363a887c117d738dacd6554f4146459407 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Fri, 3 Sep 2021 20:12:44 +0000
Subject: [PATCH 13/14] Update special method replacement

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 monai/utils/nvtx.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/monai/utils/nvtx.py b/monai/utils/nvtx.py
index 2dfbd03529..8ea5956736 100644
--- a/monai/utils/nvtx.py
+++ b/monai/utils/nvtx.py
@@ -109,7 +109,16 @@ def range_wrapper(*args, **kwargs):
             return output
 
         # Replace the method with the wrapped version
-        setattr(owner, method, range_wrapper)
+        if method.startswith("_"):
+            # If it is a special method, it requires special attention
+            class NVTXRangeDecoratedClass(owner):
+                ...
+
+            setattr(NVTXRangeDecoratedClass, method, range_wrapper)
+            obj.__class__ = NVTXRangeDecoratedClass
+
+        else:
+            setattr(owner, method, range_wrapper)
 
     def _get_method(self, obj: Any) -> tuple:
         if isinstance(obj, Module):

From 435fa7857bab29fd482461127cb9db77c24514ce Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Mon, 6 Sep 2021 20:54:56 +0000
Subject: [PATCH 14/14] Update special method check

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
---
 monai/utils/nvtx.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/monai/utils/nvtx.py b/monai/utils/nvtx.py
index 8ea5956736..1980ceef71 100644
--- a/monai/utils/nvtx.py
+++ b/monai/utils/nvtx.py
@@ -92,7 +92,7 @@ def _decorate_method(self, obj, method, append_method_name):
             name = self.name
 
         # Get the class for special functions
-        if method.startswith("_"):
+        if method.startswith("__"):
             owner = type(obj)
         else:
             owner = obj
@@ -109,7 +109,7 @@ def range_wrapper(*args, **kwargs):
             return output
 
         # Replace the method with the wrapped version
-        if method.startswith("_"):
+        if method.startswith("__"):
             # If it is a special method, it requires special attention
             class NVTXRangeDecoratedClass(owner):
                 ...