From 6d46c3f268071bca3d20abdd0320bf9fc80fb2bd Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Fri, 6 Aug 2021 21:03:24 +0000
Subject: [PATCH 01/13] Add NVTX transforms

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

diff --git a/monai/transforms/nvtx.py b/monai/transforms/nvtx.py
new file mode 100644
index 0000000000..d6991914bb
--- /dev/null
+++ b/monai/transforms/nvtx.py
@@ -0,0 +1,111 @@
+# 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.
+"""
+Wrapper around NVIDIA Tools Extension for profiling MONAI transformations
+"""
+
+from monai.transforms.transform import RandomizableTransform, Transform
+from monai.utils import optional_import
+
+_nvtx, _ = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?")
+
+__all__ = ["RangePush", "RandRangePush", "RangePop", "RandRangePop", "Mark", "RandMark"]
+
+
+class RangePush(Transform):
+    """
+    Pushes a range onto a stack of nested range span.
+    Stores zero-based depth of the range that is started.
+
+    Args:
+        msg: ASCII message to associate with range
+    """
+
+    def __init__(self, msg: str) -> None:
+        self.msg = msg
+        self.depth = None
+
+    def __call__(self, data):
+        self.depth = _nvtx.rangePushA(self.msg)
+        return data
+
+
+class RandRangePush(RandomizableTransform):
+    """
+    Pushes a range onto a stack of nested range span (RandomizableTransform).
+    Stores zero-based depth of the range that is started.
+
+    Args:
+        msg: ASCII message to associate with range
+    """
+
+    def __init__(self, msg: str) -> None:
+        self.msg = msg
+        self.depth = None
+
+    def __call__(self, data):
+        self.depth = _nvtx.rangePushA(self.msg)
+        return data
+
+
+class RangePop(Transform):
+    """
+    Pops a range off of a stack of nested range spans.
+    Stores zero-based depth of the range that is ended.
+    """
+
+    def __call__(self, data):
+        _nvtx.rangePop()
+        return data
+
+
+class RandRangePop(RandomizableTransform):
+    """
+    Pops a range off of a stack of nested range spans (RandomizableTransform).
+    Stores zero-based depth of the range that is ended.
+    """
+
+    def __call__(self, data):
+        _nvtx.rangePop()
+        return data
+
+
+class Mark(Transform):
+    """
+    Mark an instantaneous event that occurred at some point.
+
+    Args:
+        msg: ASCII message to associate with the event.
+    """
+
+    def __init__(self, msg: str) -> None:
+        self.msg = msg
+
+    def __call__(self, data):
+        _nvtx.markA(self.msg)
+        return data
+
+
+class RandMark(RandomizableTransform):
+    """
+    Mark an instantaneous event that occurred at some point.
+    (RandomizableTransform)
+
+    Args:
+        msg: ASCII message to associate with the event.
+    """
+
+    def __init__(self, msg: str) -> None:
+        self.msg = msg
+
+    def __call__(self, data):
+        _nvtx.markA(self.msg)
+        return data

From 12c27235f4b89c9a9d772df53e2d84a8d78625b8 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Fri, 6 Aug 2021 21:03:51 +0000
Subject: [PATCH 02/13] Update transforms init

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

diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py
index cf9198dbf5..073b1065c9 100644
--- a/monai/transforms/__init__.py
+++ b/monai/transforms/__init__.py
@@ -191,6 +191,7 @@
 from .inverse_batch_transform import BatchInverseTransform, Decollated
 from .io.array import LoadImage, SaveImage
 from .io.dictionary import LoadImaged, LoadImageD, LoadImageDict, SaveImaged, SaveImageD, SaveImageDict
+from .nvtx import RangePush, RandRangePush, RangePop, RandRangePop, Mark, RandMark
 from .post.array import (
     Activations,
     AsDiscrete,

From c607b1af7162568f8400a55d6448373a31f05d2a Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Fri, 6 Aug 2021 21:31:05 +0000
Subject: [PATCH 03/13] Update docs

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

diff --git a/docs/source/transforms.rst b/docs/source/transforms.rst
index 6a25c62c49..8a2460e099 100644
--- a/docs/source/transforms.rst
+++ b/docs/source/transforms.rst
@@ -335,6 +335,35 @@ IO
     :members:
     :special-members: __call__
 
+
+NVIDIA Tool Extension (NVTX)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+`RangePush`
+"""""""""""
+.. autoclass:: RangePush
+
+`RandRangePush`
+"""""""""""""""
+.. autoclass:: RandRangePush
+
+`RangePop`
+""""""""""
+.. autoclass:: RangePop
+
+`RandRangePop`
+""""""""""""""
+.. autoclass:: RandRangePop
+
+`Mark`
+""""""
+.. autoclass:: Mark
+
+`RandMark`
+""""""""""
+.. autoclass:: RandMark
+
+
 Post-processing
 ^^^^^^^^^^^^^^^
 

From 42c91c57e1c83382399eb506a3c42e810b940146 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+drbeh@users.noreply.github.com>
Date: Fri, 6 Aug 2021 21:41:12 +0000
Subject: [PATCH 04/13] Sort imports

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

diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py
index 073b1065c9..764359e5fe 100644
--- a/monai/transforms/__init__.py
+++ b/monai/transforms/__init__.py
@@ -191,7 +191,7 @@
 from .inverse_batch_transform import BatchInverseTransform, Decollated
 from .io.array import LoadImage, SaveImage
 from .io.dictionary import LoadImaged, LoadImageD, LoadImageDict, SaveImaged, SaveImageD, SaveImageDict
-from .nvtx import RangePush, RandRangePush, RangePop, RandRangePop, Mark, RandMark
+from .nvtx import Mark, RandMark, RandRangePop, RandRangePush, RangePop, RangePush
 from .post.array import (
     Activations,
     AsDiscrete,

From 8ba4f37c187c8ff03e14becc9cf1f702575260e9 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Tue, 10 Aug 2021 00:24:53 +0000
Subject: [PATCH 05/13] Update randomized transforms

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 monai/transforms/nvtx.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/monai/transforms/nvtx.py b/monai/transforms/nvtx.py
index d6991914bb..b3b1f3e0e3 100644
--- a/monai/transforms/nvtx.py
+++ b/monai/transforms/nvtx.py
@@ -48,6 +48,7 @@ class RandRangePush(RandomizableTransform):
     """
 
     def __init__(self, msg: str) -> None:
+        super().__init__()
         self.msg = msg
         self.depth = None
 
@@ -73,6 +74,9 @@ class RandRangePop(RandomizableTransform):
     Stores zero-based depth of the range that is ended.
     """
 
+    def __init__(self):
+        super().__init__()
+
     def __call__(self, data):
         _nvtx.rangePop()
         return data
@@ -104,6 +108,7 @@ class RandMark(RandomizableTransform):
     """
 
     def __init__(self, msg: str) -> None:
+        super().__init__()
         self.msg = msg
 
     def __call__(self, data):

From 20a76a8e5caf1a53449b5046e251b92d1407f91a Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Tue, 10 Aug 2021 00:37:44 +0000
Subject: [PATCH 06/13] Add unittests

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 tests/test_nvtx_transform.py | 104 +++++++++++++++++++++++++++++++++++
 1 file changed, 104 insertions(+)
 create mode 100644 tests/test_nvtx_transform.py

diff --git a/tests/test_nvtx_transform.py b/tests/test_nvtx_transform.py
new file mode 100644
index 0000000000..06c651cf57
--- /dev/null
+++ b/tests/test_nvtx_transform.py
@@ -0,0 +1,104 @@
+# 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 unittest
+import sys
+
+import numpy as np
+from numpy.lib.arraysetops import isin
+from parameterized import parameterized
+
+from torch import Tensor
+from monai.transforms import (
+    Compose,
+    Flip,
+    RandFlip,
+    RandFlipD,
+    ToTensor,
+    ToTensorD,
+)
+from monai.transforms.nvtx import Mark, RandMark, RandRangePop, RandRangePush, RangePop, RangePush
+from monai.utils import optional_import
+
+_, has_nvtx = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?")
+
+TEST_CASE_0 = [
+    np.random.randn(3, 3),
+]
+TEST_CASE_1 = [
+    {"image": np.random.randn(3, 3)},
+]
+
+
+class TestNVTXTransforms(unittest.TestCase):
+    @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
+    @parameterized.expand([TEST_CASE_0, TEST_CASE_1])
+    def test_nvtx_transfroms_alone(self, input):
+        transforms = Compose(
+            [
+                RandMark("Mark: Transform Starts!"),
+                RandRangePush("Range: RandFlipD"),
+                RandRangePop(),
+                RangePush("Range: ToTensorD"),
+                RangePop(),
+                Mark("Mark: Transform Ends!"),
+            ]
+        )
+        output = transforms(input)
+        self.assertEqual(id(input), id(output))
+
+    @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
+    @parameterized.expand([TEST_CASE_0])
+    def test_nvtx_transfroms(self, input):
+        transforms = Compose(
+            [
+                RandMark("Mark: Transform Starts!"),
+                RandRangePush("Range: RandFlip"),
+                RandFlip(prob=0.5),
+                RandRangePop(),
+                RangePush("Range: ToTensor"),
+                ToTensor(),
+                RangePop(),
+                Mark("Mark: Transform Ends!"),
+            ]
+        )
+        output = transforms(input)
+        self.assertIsInstance(output, Tensor)
+        try:
+            np.testing.assert_array_equal(input, output)
+        except AssertionError:
+            np.testing.assert_array_equal(input, Flip()(output.numpy()))
+
+    @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
+    @parameterized.expand([TEST_CASE_1])
+    def test_nvtx_transfromsd(self, input):
+        transforms = Compose(
+            [
+                RandMark("Mark: Transform Starts!"),
+                RandRangePush("Range: RandFlipD"),
+                RandFlipD(keys="image", prob=0.5),
+                RandRangePop(),
+                RangePush("Range: ToTensorD"),
+                ToTensorD(keys=("image")),
+                RangePop(),
+                Mark("Mark: Transform Ends!"),
+            ]
+        )
+        output = transforms(input)
+        self.assertIsInstance(output["image"], Tensor)
+        try:
+            np.testing.assert_array_equal(input["image"], output["image"])
+        except AssertionError:
+            np.testing.assert_array_equal(input["image"], Flip()(output["image"].numpy()))
+
+
+if __name__ == "__main__":
+    unittest.main()

From 09b2d462ddd5a4e073a7de04b644feef109f23e8 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Tue, 10 Aug 2021 00:43:33 +0000
Subject: [PATCH 07/13] sort imports

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 tests/test_nvtx_transform.py | 13 ++-----------
 1 file changed, 2 insertions(+), 11 deletions(-)

diff --git a/tests/test_nvtx_transform.py b/tests/test_nvtx_transform.py
index 06c651cf57..f6bfbe2fac 100644
--- a/tests/test_nvtx_transform.py
+++ b/tests/test_nvtx_transform.py
@@ -10,21 +10,12 @@
 # limitations under the License.
 
 import unittest
-import sys
 
 import numpy as np
-from numpy.lib.arraysetops import isin
 from parameterized import parameterized
-
 from torch import Tensor
-from monai.transforms import (
-    Compose,
-    Flip,
-    RandFlip,
-    RandFlipD,
-    ToTensor,
-    ToTensorD,
-)
+
+from monai.transforms import Compose, Flip, RandFlip, RandFlipD, ToTensor, ToTensorD
 from monai.transforms.nvtx import Mark, RandMark, RandRangePop, RandRangePush, RangePop, RangePush
 from monai.utils import optional_import
 

From b0c6020e3b204635ad5ab79c8f27cbadc184db39 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Tue, 10 Aug 2021 01:21:53 +0000
Subject: [PATCH 08/13] Add dictionary-based variants

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 monai/transforms/__init__.py | 27 ++++++++++++++++++++++++++-
 monai/transforms/nvtx.py     |  8 ++++++++
 tests/test_nvtx_transform.py | 29 +++++++++++++++++++++--------
 3 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py
index 0fa04cc932..9e50a6f3f9 100644
--- a/monai/transforms/__init__.py
+++ b/monai/transforms/__init__.py
@@ -191,7 +191,32 @@
 from .inverse_batch_transform import BatchInverseTransform, Decollated
 from .io.array import LoadImage, SaveImage
 from .io.dictionary import LoadImaged, LoadImageD, LoadImageDict, SaveImaged, SaveImageD, SaveImageDict
-from .nvtx import Mark, RandMark, RandRangePop, RandRangePush, RangePop, RangePush
+from .nvtx import (
+    Mark,
+    Markd,
+    MarkD,
+    MarkDict,
+    RandMark,
+    RandMarkd,
+    RandMarkD,
+    RandMarkDict,
+    RandRangePop,
+    RandRangePopd,
+    RandRangePopD,
+    RandRangePopDict,
+    RandRangePush,
+    RandRangePushd,
+    RandRangePushD,
+    RandRangePushDict,
+    RangePop,
+    RangePopd,
+    RangePopD,
+    RangePopDict,
+    RangePush,
+    RangePushd,
+    RangePushD,
+    RangePushDict,
+)
 from .post.array import (
     Activations,
     AsDiscrete,
diff --git a/monai/transforms/nvtx.py b/monai/transforms/nvtx.py
index b3b1f3e0e3..970d763265 100644
--- a/monai/transforms/nvtx.py
+++ b/monai/transforms/nvtx.py
@@ -114,3 +114,11 @@ def __init__(self, msg: str) -> None:
     def __call__(self, data):
         _nvtx.markA(self.msg)
         return data
+
+
+MarkDict = MarkD = Markd = Mark
+RandMarkDict = RandMarkD = RandMarkd = RandMark
+RandRangePopDict = RandRangePopD = RandRangePopd = RandRangePop
+RandRangePushDict = RandRangePushD = RandRangePushd = RandRangePush
+RangePopDict = RangePopD = RangePopd = RangePop
+RangePushDict = RangePushD = RangePushd = RangePush
diff --git a/tests/test_nvtx_transform.py b/tests/test_nvtx_transform.py
index f6bfbe2fac..6a9d23ee29 100644
--- a/tests/test_nvtx_transform.py
+++ b/tests/test_nvtx_transform.py
@@ -16,7 +16,20 @@
 from torch import Tensor
 
 from monai.transforms import Compose, Flip, RandFlip, RandFlipD, ToTensor, ToTensorD
-from monai.transforms.nvtx import Mark, RandMark, RandRangePop, RandRangePush, RangePop, RangePush
+from monai.transforms.nvtx import (
+    Mark,
+    MarkD,
+    RandMark,
+    RandMarkD,
+    RandRangePop,
+    RandRangePopD,
+    RandRangePush,
+    RandRangePushD,
+    RangePop,
+    RangePopD,
+    RangePush,
+    RangePushD,
+)
 from monai.utils import optional_import
 
 _, has_nvtx = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?")
@@ -48,7 +61,7 @@ def test_nvtx_transfroms_alone(self, input):
 
     @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
     @parameterized.expand([TEST_CASE_0])
-    def test_nvtx_transfroms(self, input):
+    def test_nvtx_transfroms_array(self, input):
         transforms = Compose(
             [
                 RandMark("Mark: Transform Starts!"),
@@ -73,14 +86,14 @@ def test_nvtx_transfroms(self, input):
     def test_nvtx_transfromsd(self, input):
         transforms = Compose(
             [
-                RandMark("Mark: Transform Starts!"),
-                RandRangePush("Range: RandFlipD"),
+                RandMarkD("Mark: Transform Starts!"),
+                RandRangePushD("Range: RandFlipD"),
                 RandFlipD(keys="image", prob=0.5),
-                RandRangePop(),
-                RangePush("Range: ToTensorD"),
+                RandRangePopD(),
+                RangePushD("Range: ToTensorD"),
                 ToTensorD(keys=("image")),
-                RangePop(),
-                Mark("Mark: Transform Ends!"),
+                RangePopD(),
+                MarkD("Mark: Transform Ends!"),
             ]
         )
         output = transforms(input)

From 97c26e41fd24798f055afe965c02140eede3c5ce Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Tue, 10 Aug 2021 01:31:03 +0000
Subject: [PATCH 09/13] Fix parametrized and skip order

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 tests/test_nvtx_transform.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tests/test_nvtx_transform.py b/tests/test_nvtx_transform.py
index 6a9d23ee29..b7d9543ec9 100644
--- a/tests/test_nvtx_transform.py
+++ b/tests/test_nvtx_transform.py
@@ -13,7 +13,7 @@
 
 import numpy as np
 from parameterized import parameterized
-from torch import Tensor
+import torch
 
 from monai.transforms import Compose, Flip, RandFlip, RandFlipD, ToTensor, ToTensorD
 from monai.transforms.nvtx import (
@@ -43,8 +43,8 @@
 
 
 class TestNVTXTransforms(unittest.TestCase):
-    @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
     @parameterized.expand([TEST_CASE_0, TEST_CASE_1])
+    @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
     def test_nvtx_transfroms_alone(self, input):
         transforms = Compose(
             [
@@ -59,8 +59,8 @@ def test_nvtx_transfroms_alone(self, input):
         output = transforms(input)
         self.assertEqual(id(input), id(output))
 
-    @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
     @parameterized.expand([TEST_CASE_0])
+    @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
     def test_nvtx_transfroms_array(self, input):
         transforms = Compose(
             [
@@ -75,14 +75,14 @@ def test_nvtx_transfroms_array(self, input):
             ]
         )
         output = transforms(input)
-        self.assertIsInstance(output, Tensor)
+        self.assertIsInstance(output, torch.Tensor)
         try:
             np.testing.assert_array_equal(input, output)
         except AssertionError:
             np.testing.assert_array_equal(input, Flip()(output.numpy()))
 
-    @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
     @parameterized.expand([TEST_CASE_1])
+    @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
     def test_nvtx_transfromsd(self, input):
         transforms = Compose(
             [
@@ -97,7 +97,7 @@ def test_nvtx_transfromsd(self, input):
             ]
         )
         output = transforms(input)
-        self.assertIsInstance(output["image"], Tensor)
+        self.assertIsInstance(output["image"], torch.Tensor)
         try:
             np.testing.assert_array_equal(input["image"], output["image"])
         except AssertionError:

From c6cd85b8e3ae1e045c4f44f7776068f781b77829 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Tue, 10 Aug 2021 01:50:59 +0000
Subject: [PATCH 10/13] Sort import

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 tests/test_nvtx_transform.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_nvtx_transform.py b/tests/test_nvtx_transform.py
index b7d9543ec9..6a5aa4feb6 100644
--- a/tests/test_nvtx_transform.py
+++ b/tests/test_nvtx_transform.py
@@ -12,8 +12,8 @@
 import unittest
 
 import numpy as np
-from parameterized import parameterized
 import torch
+from parameterized import parameterized
 
 from monai.transforms import Compose, Flip, RandFlip, RandFlipD, ToTensor, ToTensorD
 from monai.transforms.nvtx import (

From e7a8118852b3410ed360ac8c931a917c2d0fc060 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Tue, 10 Aug 2021 13:43:21 +0000
Subject: [PATCH 11/13] Update rand tranforms

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 monai/transforms/nvtx.py | 57 ++++++++++++++++++++--------------------
 1 file changed, 29 insertions(+), 28 deletions(-)

diff --git a/monai/transforms/nvtx.py b/monai/transforms/nvtx.py
index 970d763265..12c03dc028 100644
--- a/monai/transforms/nvtx.py
+++ b/monai/transforms/nvtx.py
@@ -17,7 +17,32 @@
 
 _nvtx, _ = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?")
 
-__all__ = ["RangePush", "RandRangePush", "RangePop", "RandRangePop", "Mark", "RandMark"]
+__all__ = [
+    "Mark",
+    "Markd",
+    "MarkD",
+    "MarkDict",
+    "RandMark",
+    "RandMarkd",
+    "RandMarkD",
+    "RandMarkDict",
+    "RandRangePop",
+    "RandRangePopd",
+    "RandRangePopD",
+    "RandRangePopDict",
+    "RandRangePush",
+    "RandRangePushd",
+    "RandRangePushD",
+    "RandRangePushDict",
+    "RangePop",
+    "RangePopd",
+    "RangePopD",
+    "RangePopDict",
+    "RangePush",
+    "RangePushd",
+    "RangePushD",
+    "RangePushDict",
+]
 
 
 class RangePush(Transform):
@@ -38,7 +63,7 @@ def __call__(self, data):
         return data
 
 
-class RandRangePush(RandomizableTransform):
+class RandRangePush(RangePush, RandomizableTransform):
     """
     Pushes a range onto a stack of nested range span (RandomizableTransform).
     Stores zero-based depth of the range that is started.
@@ -47,15 +72,6 @@ class RandRangePush(RandomizableTransform):
         msg: ASCII message to associate with range
     """
 
-    def __init__(self, msg: str) -> None:
-        super().__init__()
-        self.msg = msg
-        self.depth = None
-
-    def __call__(self, data):
-        self.depth = _nvtx.rangePushA(self.msg)
-        return data
-
 
 class RangePop(Transform):
     """
@@ -68,19 +84,12 @@ def __call__(self, data):
         return data
 
 
-class RandRangePop(RandomizableTransform):
+class RandRangePop(RangePop, RandomizableTransform):
     """
     Pops a range off of a stack of nested range spans (RandomizableTransform).
     Stores zero-based depth of the range that is ended.
     """
 
-    def __init__(self):
-        super().__init__()
-
-    def __call__(self, data):
-        _nvtx.rangePop()
-        return data
-
 
 class Mark(Transform):
     """
@@ -98,7 +107,7 @@ def __call__(self, data):
         return data
 
 
-class RandMark(RandomizableTransform):
+class RandMark(Mark, RandomizableTransform):
     """
     Mark an instantaneous event that occurred at some point.
     (RandomizableTransform)
@@ -107,14 +116,6 @@ class RandMark(RandomizableTransform):
         msg: ASCII message to associate with the event.
     """
 
-    def __init__(self, msg: str) -> None:
-        super().__init__()
-        self.msg = msg
-
-    def __call__(self, data):
-        _nvtx.markA(self.msg)
-        return data
-
 
 MarkDict = MarkD = Markd = Mark
 RandMarkDict = RandMarkD = RandMarkd = RandMark

From 05bd0b630c1852a5cc51df7a990e356a5a764349 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Tue, 10 Aug 2021 13:43:54 +0000
Subject: [PATCH 12/13] Update test cases

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 tests/test_nvtx_transform.py | 96 ++++++++++++++++++++++++++++--------
 1 file changed, 75 insertions(+), 21 deletions(-)

diff --git a/tests/test_nvtx_transform.py b/tests/test_nvtx_transform.py
index 6a5aa4feb6..353972dcc3 100644
--- a/tests/test_nvtx_transform.py
+++ b/tests/test_nvtx_transform.py
@@ -15,7 +15,15 @@
 import torch
 from parameterized import parameterized
 
-from monai.transforms import Compose, Flip, RandFlip, RandFlipD, ToTensor, ToTensorD
+from monai.transforms import (
+    Compose,
+    Flip,
+    RandFlip,
+    RandFlipD,
+    ToTensor,
+    ToTensorD,
+    Randomizable,
+)
 from monai.transforms.nvtx import (
     Mark,
     MarkD,
@@ -34,39 +42,59 @@
 
 _, has_nvtx = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?")
 
-TEST_CASE_0 = [
+
+TEST_CASE_ARRAY_0 = [
     np.random.randn(3, 3),
 ]
-TEST_CASE_1 = [
+TEST_CASE_ARRAY_1 = [
+    np.random.randn(3, 10, 10),
+]
+TEST_CASE_DICT_0 = [
     {"image": np.random.randn(3, 3)},
 ]
+TEST_CASE_DICT_1 = [
+    {"image": np.random.randn(3, 10, 10)},
+]
 
 
 class TestNVTXTransforms(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_0, TEST_CASE_1])
+    @parameterized.expand(
+        [
+            TEST_CASE_ARRAY_0,
+            TEST_CASE_ARRAY_1,
+            TEST_CASE_DICT_0,
+            TEST_CASE_DICT_1,
+        ]
+    )
     @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
     def test_nvtx_transfroms_alone(self, input):
         transforms = Compose(
             [
-                RandMark("Mark: Transform Starts!"),
-                RandRangePush("Range: RandFlipD"),
-                RandRangePop(),
-                RangePush("Range: ToTensorD"),
+                Mark("Mark: Transform Starts!"),
+                RangePush("Range: RandFlipD"),
                 RangePop(),
-                Mark("Mark: Transform Ends!"),
+                RandRangePush("Range: ToTensorD"),
+                RandRangePop(),
+                RandMark("Mark: Transform Ends!"),
             ]
         )
         output = transforms(input)
         self.assertEqual(id(input), id(output))
 
-    @parameterized.expand([TEST_CASE_0])
+        # Check if chain of randomizable/non-randomizable transforms is not broken
+        for tran in transforms.transforms:
+            if isinstance(tran, Randomizable):
+                self.assertIsInstance(tran, RangePush)
+                break
+
+    @parameterized.expand([TEST_CASE_ARRAY_0, TEST_CASE_ARRAY_1])
     @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
     def test_nvtx_transfroms_array(self, input):
         transforms = Compose(
             [
                 RandMark("Mark: Transform Starts!"),
                 RandRangePush("Range: RandFlip"),
-                RandFlip(prob=0.5),
+                RandFlip(prob=0.0),
                 RandRangePop(),
                 RangePush("Range: ToTensor"),
                 ToTensor(),
@@ -76,19 +104,48 @@ def test_nvtx_transfroms_array(self, input):
         )
         output = transforms(input)
         self.assertIsInstance(output, torch.Tensor)
-        try:
-            np.testing.assert_array_equal(input, output)
-        except AssertionError:
-            np.testing.assert_array_equal(input, Flip()(output.numpy()))
+        np.testing.assert_array_equal(input, output)
 
-    @parameterized.expand([TEST_CASE_1])
+        transforms = Compose(
+            [
+                RandMark("Mark: Transform Starts!"),
+                RandRangePush("Range: RandFlip"),
+                RandFlip(prob=1.0),
+                RandRangePop(),
+                RangePush("Range: ToTensor"),
+                ToTensor(),
+                RangePop(),
+                Mark("Mark: Transform Ends!"),
+            ]
+        )
+        output = transforms(input)
+        self.assertIsInstance(output, torch.Tensor)
+        np.testing.assert_array_equal(input, Flip()(output.numpy()))
+
+    @parameterized.expand([TEST_CASE_DICT_0, TEST_CASE_DICT_1])
     @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
     def test_nvtx_transfromsd(self, input):
         transforms = Compose(
             [
                 RandMarkD("Mark: Transform Starts!"),
                 RandRangePushD("Range: RandFlipD"),
-                RandFlipD(keys="image", prob=0.5),
+                RandFlipD(keys="image", prob=0.0),
+                RandRangePopD(),
+                RangePushD("Range: ToTensorD"),
+                ToTensorD(keys=("image")),
+                RangePopD(),
+                MarkD("Mark: Transform Ends!"),
+            ]
+        )
+        output = transforms(input)
+        self.assertIsInstance(output["image"], torch.Tensor)
+        np.testing.assert_array_equal(input["image"], output["image"])
+
+        transforms = Compose(
+            [
+                RandMarkD("Mark: Transform Starts!"),
+                RandRangePushD("Range: RandFlipD"),
+                RandFlipD(keys="image", prob=1.0),
                 RandRangePopD(),
                 RangePushD("Range: ToTensorD"),
                 ToTensorD(keys=("image")),
@@ -98,10 +155,7 @@ def test_nvtx_transfromsd(self, input):
         )
         output = transforms(input)
         self.assertIsInstance(output["image"], torch.Tensor)
-        try:
-            np.testing.assert_array_equal(input["image"], output["image"])
-        except AssertionError:
-            np.testing.assert_array_equal(input["image"], Flip()(output["image"].numpy()))
+        np.testing.assert_array_equal(input["image"], Flip()(output["image"].numpy()))
 
 
 if __name__ == "__main__":

From 3a2476c48c09f090a4b92e2b3b0798be4826f4d0 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Tue, 10 Aug 2021 13:46:10 +0000
Subject: [PATCH 13/13] Update test cases

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 tests/test_nvtx_transform.py | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/tests/test_nvtx_transform.py b/tests/test_nvtx_transform.py
index 353972dcc3..d1887377ba 100644
--- a/tests/test_nvtx_transform.py
+++ b/tests/test_nvtx_transform.py
@@ -15,15 +15,7 @@
 import torch
 from parameterized import parameterized
 
-from monai.transforms import (
-    Compose,
-    Flip,
-    RandFlip,
-    RandFlipD,
-    ToTensor,
-    ToTensorD,
-    Randomizable,
-)
+from monai.transforms import Compose, Flip, RandFlip, RandFlipD, Randomizable, ToTensor, ToTensorD
 from monai.transforms.nvtx import (
     Mark,
     MarkD,