From 9d92e6b45d29d5cd7630eb4d4c9abb096c2e13ab Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Sat, 4 Jun 2022 16:17:42 -0600
Subject: [PATCH 01/30] Improve LaplacianImage

* Improved `LaplacianImage` code as part of removing the ToDo comment.
* Added `"same"` padding option to `GaussianSmoothing`.
---
 captum/optim/_param/image/images.py     | 83 ++++++++++---------------
 captum/optim/_param/image/transforms.py | 11 +++-
 tests/optim/param/test_images.py        | 52 +++++++++++++---
 tests/optim/param/test_transforms.py    | 43 ++++++++++---
 4 files changed, 122 insertions(+), 67 deletions(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index fa313b38af..28921fce65 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -1,4 +1,3 @@
-from copy import deepcopy
 from types import MethodType
 from typing import Callable, List, Optional, Tuple, Type, Union
 
@@ -374,22 +373,24 @@ def forward(self) -> torch.Tensor:
 
 class LaplacianImage(ImageParameterization):
     """
-    TODO: Fix divison by 6 in setup_input when init is not None.
     Parameterize an image tensor with a laplacian pyramid.
     """
 
     def __init__(
         self,
-        size: Tuple[int, int] = None,
+        size: Tuple[int, int] = (224, 225),
         channels: int = 3,
         batch: int = 1,
         init: Optional[torch.Tensor] = None,
+        power: float = 0.1,
+        scale_list: List[float] = [1.0, 2.0, 4.0, 8.0, 16.0, 32.0],
     ) -> None:
         """
         Args:
 
             size (Tuple[int, int]): The height & width dimensions to use for the
                 parameterized output image tensor.
+                Default: (224, 224)
             channels (int, optional): The number of channels to use for each image.
                 Default: 3
             batch (int, optional): The number of images to stack along the batch
@@ -398,72 +399,56 @@ def __init__(
             init (torch.tensor, optional): Optionally specify a tensor to
                 use instead of creating one.
                 Default: None
+            power (float, optional): The desired power value to use.
+                Default: 0.1
+            scale_list (list of float, optional): The desired list of scale values to
+                use in the laplacian pyramid. The height & width dimensions specified
+                in size or used in the init tensor should be divisable by every scale
+                value in the scale list with no remainder left over. The default
+                scale_list values are set to work with a size of (224, 224).
+                Default: [1.0, 2.0, 4.0, 8.0, 16.0, 32.0]
         """
         super().__init__()
-        power = 0.1
-
-        if init is None:
-            tensor_params, self.scaler = self._setup_input(size, channels, power, init)
-
-            self.tensor_params = torch.nn.ModuleList(
-                [deepcopy(tensor_params) for b in range(batch)]
-            )
-        else:
+        if init is not None:
+            assert init.dim() in [3, 4]
             init = init.unsqueeze(0) if init.dim() == 3 else init
-            P = []
-            for b in range(init.size(0)):
-                tensor_params, self.scaler = self._setup_input(
-                    size, channels, power, init[b].unsqueeze(0)
-                )
-                P.append(tensor_params)
-            self.tensor_params = torch.nn.ModuleList(P)
+            size = list(init.shape[2:])
 
-    def _setup_input(
-        self,
-        size: Tuple[int, int],
-        channels: int,
-        power: float = 0.1,
-        init: Optional[torch.Tensor] = None,
-    ) -> Tuple[List[torch.Tensor], List[torch.nn.Upsample]]:
         tensor_params, scaler = [], []
-        scale_list = [1, 2, 4, 8, 16, 32]
         for scale in scale_list:
+            assert size[0] % scale == 0 and size[1] % scale == 0, (
+                "The chosen image height & width dimensions"
+                + " must be divisable by all scale values "
+                + " with no remainder left over."
+            )
+
             h, w = int(size[0] // scale), int(size[1] // scale)
             if init is None:
-                x = torch.randn([1, channels, h, w]) / 10
+                x = torch.randn([batch, channels, h, w]) / 10
             else:
                 x = F.interpolate(init.clone(), size=(h, w), mode="bilinear")
-                x = x / 6  # Prevents output from being all white
+                x = x / 10
             upsample = torch.nn.Upsample(scale_factor=scale, mode="nearest")
-            x = x * (scale**power) / (32**power)
+            x = x * (scale**power) / (max(scale_list) ** power)
             x = torch.nn.Parameter(x)
             tensor_params.append(x)
             scaler.append(upsample)
-        tensor_params = torch.nn.ParameterList(tensor_params)
-        return tensor_params, scaler
+        self.tensor_params = torch.nn.ParameterList(tensor_params)
+        self.scaler = scaler
 
-    def _create_tensor(self, params_list: torch.nn.ParameterList) -> torch.Tensor:
+    def forward(self) -> torch.Tensor:
         """
-        Resize tensor parameters to the target size.
-
-        Args:
-
-            params_list (torch.nn.ParameterList): List of tensors to resize.
-
         Returns:
-            **tensor** (torch.Tensor): The sum of all tensor parameters.
+            **output** (torch.tensor): A tensor created from a laplacian pyramid.
         """
-        A: List[torch.Tensor] = []
-        for xi, upsamplei in zip(params_list, self.scaler):
+        A = []
+        for xi, upsamplei in zip(self.tensor_params, self.scaler):
             A.append(upsamplei(xi))
-        return torch.sum(torch.cat(A), 0) + 0.5
+        output = sum(A) + 0.5
 
-    def forward(self) -> torch.Tensor:
-        A: List[torch.Tensor] = []
-        for params_list in self.tensor_params:
-            tensor = self._create_tensor(params_list)
-            A.append(tensor)
-        return torch.stack(A).refine_names("B", "C", "H", "W")
+        if torch.jit.is_scripting():
+            return output
+        return output.refine_names("B", "C", "H", "W")
 
 
 class SimpleTensorParameterization(ImageParameterization):
diff --git a/captum/optim/_param/image/transforms.py b/captum/optim/_param/image/transforms.py
index 4ec8762637..fbf7b01cbc 100644
--- a/captum/optim/_param/image/transforms.py
+++ b/captum/optim/_param/image/transforms.py
@@ -979,7 +979,7 @@ class GaussianSmoothing(nn.Module):
     in the input using a depthwise convolution.
     """
 
-    __constants__ = ["groups"]
+    __constants__ = ["groups", "padding"]
 
     def __init__(
         self,
@@ -987,6 +987,7 @@ def __init__(
         kernel_size: Union[int, Sequence[int]],
         sigma: Union[float, Sequence[float]],
         dim: int = 2,
+        use_same_padding: bool = True,
     ) -> None:
         """
         Args:
@@ -997,6 +998,9 @@ def __init__(
             sigma (float, sequence): Standard deviation of the gaussian kernel.
             dim (int, optional): The number of dimensions of the data.
                 Default value is 2 (spatial).
+            use_same_padding (bool, optional): Whether or not to use "same" padding so
+                that the output shape is the same as the input shape.
+                Default: True
         """
         super().__init__()
         if isinstance(kernel_size, numbers.Number):
@@ -1027,6 +1031,7 @@ def __init__(
 
         self.register_buffer("weight", kernel)
         self.groups = channels
+        self.padding = "same" if use_same_padding else 0
 
         if dim == 1:
             self.conv = F.conv1d
@@ -1050,7 +1055,9 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         Returns:
             **filtered** (torch.Tensor): Filtered output.
         """
-        return self.conv(input, weight=self.weight, groups=self.groups)
+        return self.conv(
+            input, weight=self.weight, groups=self.groups, padding=self.padding
+        )
 
 
 class SymmetricPadding(torch.autograd.Function):
diff --git a/tests/optim/param/test_images.py b/tests/optim/param/test_images.py
index 617d34a3a3..525d232998 100644
--- a/tests/optim/param/test_images.py
+++ b/tests/optim/param/test_images.py
@@ -384,21 +384,55 @@ def test_subclass(self) -> None:
     def test_laplacianimage_random_forward(self) -> None:
         size = (224, 224)
         channels = 3
-        image_param = images.LaplacianImage(size=size, channels=channels)
+        batch = 1
+        image_param = images.LaplacianImage(size=size, channels=channels, batch=batch)
         test_tensor = image_param.forward().rename(None)
+        self.assertEqual(list(test_tensor.shape), [batch, channels, size[0], size[1]])
+        self.assertTrue(test_tensor.requires_grad)
 
-        self.assertEqual(test_tensor.dim(), 4)
-        self.assertEqual(test_tensor.size(0), 1)
-        self.assertEqual(test_tensor.size(1), channels)
-        self.assertEqual(test_tensor.size(2), size[0])
-        self.assertEqual(test_tensor.size(3), size[1])
+    def test_laplacianimage_random_forward_batch_5(self) -> None:
+        size = (224, 224)
+        channels = 3
+        batch = 5
+        image_param = images.LaplacianImage(size=size, channels=channels, batch=batch)
+        test_tensor = image_param.forward().rename(None)
+        self.assertEqual(list(test_tensor.shape), [batch, channels, size[0], size[1]])
+
+    def test_laplacianimage_random_forward_scale_list(self) -> None:
+        size = (224, 224)
+        channels = 3
+        batch = 1
+        scale_list = [1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 56.0, 112.0]
+        image_param = images.LaplacianImage(
+            size=size, channels=channels, batch=batch, scale_list=scale_list
+        )
+        test_tensor = image_param.forward().rename(None)
+        self.assertEqual(list(test_tensor.shape), [batch, channels, size[0], size[1]])
+
+    def test_laplacianimage_random_forward_scale_list_error(self) -> None:
+        scale_list = [1.0, 2.0, 4.0, 8.0, 16.0, 64.0, 144.0]
+        with self.assertRaises(AssertionError):
+            images.LaplacianImage(
+                size=(224, 224), channels=3, batch=1, scale_list=scale_list
+            )
 
-    def test_laplacianimage_init(self) -> None:
-        init_t = torch.zeros(1, 224, 224)
-        image_param = images.LaplacianImage(size=(224, 224), channels=3, init=init_t)
+    def test_laplacianimage_init_tensor(self) -> None:
+        init_tensor = torch.zeros(1, 3, 224, 224)
+        image_param = images.LaplacianImage(init=init_tensor)
         output = image_param.forward().detach().rename(None)
         assertTensorAlmostEqual(self, torch.ones_like(output) * 0.5, output, mode="max")
 
+    def test_laplacianimage_random_forward_cuda(self) -> None:
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest(
+                "Skipping LaplacianImage CUDA test due to not supporting CUDA."
+            )
+        image_param = images.LaplacianImage(size=(224, 224), channels=3, batch=1).cuda()
+        test_tensor = image_param.forward().rename(None)
+        self.assertTrue(test_tensor.is_cuda)
+        self.assertEqual(list(test_tensor.shape), [1, 3, 224, 224])
+        self.assertTrue(test_tensor.requires_grad)
+
 
 class TestSimpleTensorParameterization(BaseTest):
     def test_subclass(self) -> None:
diff --git a/tests/optim/param/test_transforms.py b/tests/optim/param/test_transforms.py
index 385006a7ac..911330cb75 100644
--- a/tests/optim/param/test_transforms.py
+++ b/tests/optim/param/test_transforms.py
@@ -1582,11 +1582,16 @@ def test_gaussian_smoothing_init_1d(self) -> None:
         sigma = 2.0
         dim = 1
         smoothening_module = transforms.GaussianSmoothing(
-            channels, kernel_size, sigma, dim
+            channels,
+            kernel_size,
+            sigma,
+            dim,
+            use_same_padding=False,
         )
         self.assertEqual(smoothening_module.groups, channels)
         weight = torch.tensor([[0.3192, 0.3617, 0.3192]]).repeat(6, 1, 1)
         assertTensorAlmostEqual(self, smoothening_module.weight, weight, 0.001)
+        self.assertFalse(smoothening_module.padding)
 
     def test_gaussian_smoothing_init_2d(self) -> None:
         channels = 3
@@ -1594,7 +1599,11 @@ def test_gaussian_smoothing_init_2d(self) -> None:
         sigma = 2.0
         dim = 2
         smoothening_module = transforms.GaussianSmoothing(
-            channels, kernel_size, sigma, dim
+            channels,
+            kernel_size,
+            sigma,
+            dim,
+            use_same_padding=False,
         )
         self.assertEqual(smoothening_module.groups, channels)
         weight = torch.tensor(
@@ -1614,7 +1623,11 @@ def test_gaussian_smoothing_init_3d(self) -> None:
         sigma = 1.021
         dim = 3
         smoothening_module = transforms.GaussianSmoothing(
-            channels, kernel_size, sigma, dim
+            channels,
+            kernel_size,
+            sigma,
+            dim,
+            use_same_padding=False,
         )
         self.assertEqual(smoothening_module.groups, channels)
         weight = torch.tensor(
@@ -1654,7 +1667,11 @@ def test_gaussian_smoothing_1d(self) -> None:
         sigma = 2.0
         dim = 1
         smoothening_module = transforms.GaussianSmoothing(
-            channels, kernel_size, sigma, dim
+            channels,
+            kernel_size,
+            sigma,
+            dim,
+            use_same_padding=False,
         )
 
         test_tensor = torch.tensor([1.0, 5.0]).repeat(6, 2).unsqueeze(0)
@@ -1671,7 +1688,11 @@ def test_gaussian_smoothing_2d(self) -> None:
         sigma = 2.0
         dim = 2
         smoothening_module = transforms.GaussianSmoothing(
-            channels, kernel_size, sigma, dim
+            channels,
+            kernel_size,
+            sigma,
+            dim,
+            use_same_padding=False,
         )
 
         test_tensor = torch.tensor([1.0, 5.0]).repeat(3, 6, 3).unsqueeze(0)
@@ -1688,7 +1709,11 @@ def test_gaussian_smoothing_3d(self) -> None:
         sigma = 1.021
         dim = 3
         smoothening_module = transforms.GaussianSmoothing(
-            channels, kernel_size, sigma, dim
+            channels,
+            kernel_size,
+            sigma,
+            dim,
+            use_same_padding=False,
         )
 
         test_tensor = torch.tensor([1.0, 5.0, 1.0]).repeat(4, 6, 6, 2).unsqueeze(0)
@@ -1712,7 +1737,11 @@ def test_gaussian_smoothing_2d_jit_module(self) -> None:
         sigma = 2.0
         dim = 2
         smoothening_module = transforms.GaussianSmoothing(
-            channels, kernel_size, sigma, dim
+            channels,
+            kernel_size,
+            sigma,
+            dim,
+            use_same_padding=False,
         )
         jit_smoothening_module = torch.jit.script(smoothening_module)
 

From a8fa243387fab35bd2552ec3e5afb02404760674 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Mon, 6 Jun 2022 14:05:09 -0600
Subject: [PATCH 02/30] Improve NaturalImage docs

* Improved the `NaturalImage` docs.
* Also removed built-in clamp squash function. It's inclusion was an aesthetic choice, and it's easy enough for users to add on their own.
---
 captum/optim/_param/image/images.py | 33 +++++++++++++++-----------
 tests/optim/param/test_images.py    | 36 ++++++++++++++++++-----------
 2 files changed, 42 insertions(+), 27 deletions(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index 28921fce65..da697c66c1 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -798,24 +798,37 @@ def __init__(
         Args:
 
             size (Tuple[int, int], optional): The height and width to use for the
-                nn.Parameter image tensor.
+                nn.Parameter image tensor. This parameter is not used if
+                parameterization is an instance.
                 Default: (224, 224)
             channels (int, optional): The number of channels to use when creating the
-                nn.Parameter tensor.
+                nn.Parameter tensor. This parameter is not used if parameterization is
+                an instance.
                 Default: 3
             batch (int, optional): The number of channels to use when creating the
-                nn.Parameter tensor, or stacking init images.
+                nn.Parameter tensor, or stacking init images. This parameter is not
+                used if parameterization is an instance.
                 Default: 1
+            init (torch.tensor, optional): Optionally specify a tensor to use instead
+                of creating one from random noise. This parameter is not used if
+                parameterization is an instance. Set to None for random init.
+                Default: None
             parameterization (ImageParameterization, optional): An image
                 parameterization class, or instance of an image parameterization class.
                 Default: FFTImage
             squash_func (Callable[[torch.Tensor], torch.Tensor]], optional): The squash
-                function to use after color recorrelation. A funtion or lambda function.
+                function to use after color recorrelation. A function, lambda function,
+                or callable class instance.
                 Default: None
-            decorrelation_module (nn.Module, optional): A ToRGB instance.
+            decorrelation_module (nn.Module, optional): A module instance that
+                recorrelates the colors of an input image. Custom modules can make use
+                of the decorrelate_init parameter by having a second inverse parameter
+                in their forward functions that performs the inverse operation when it
+                is set to True. Set to None for no recorrelation.
                 Default: ToRGB
             decorrelate_init (bool, optional): Whether or not to apply color
-                decorrelation to the init tensor input.
+                decorrelation to the init tensor input. This parameter is not used if
+                parameterization is an instance or if init is None.
                 Default: True
         """
         super().__init__()
@@ -836,9 +849,6 @@ def __init__(
                 )
                 init = self.decorrelate(init, inverse=True).rename(None)
 
-            if squash_func is None:
-                squash_func = self._clamp_image
-
         self.squash_func = torch.sigmoid if squash_func is None else squash_func
         if not isinstance(parameterization, ImageParameterization):
             parameterization = parameterization(
@@ -846,11 +856,6 @@ def __init__(
             )
         self.parameterization = parameterization
 
-    @torch.jit.export
-    def _clamp_image(self, x: torch.Tensor) -> torch.Tensor:
-        """JIT supported squash function."""
-        return x.clamp(0, 1)
-
     @torch.jit.ignore
     def _to_image_tensor(self, x: torch.Tensor) -> torch.Tensor:
         """
diff --git a/tests/optim/param/test_images.py b/tests/optim/param/test_images.py
index 525d232998..567af0e5e1 100644
--- a/tests/optim/param/test_images.py
+++ b/tests/optim/param/test_images.py
@@ -1105,11 +1105,19 @@ def test_natural_image_init_func_pixelimage(self) -> None:
         self.assertIsInstance(image_param.decorrelate, ToRGB)
         self.assertEqual(image_param.squash_func, torch.sigmoid)
 
-    def test_natural_image_init_func_default_init_tensor(self) -> None:
-        image_param = images.NaturalImage(init=torch.ones(1, 3, 1, 1))
+    def test_natural_image_custom_squash_func(self) -> None:
+        init_tensor = torch.randn(1, 3, 1, 1)
+
+        def clamp_image(x: torch.Tensor) -> torch.Tensor:
+            return x.clamp(0, 1)
+
+        image_param = images.NaturalImage(init=init_tensor, squash_func=clamp_image)
+        image = image_param.forward().detach()
+
         self.assertIsInstance(image_param.parameterization, images.FFTImage)
         self.assertIsInstance(image_param.decorrelate, ToRGB)
-        self.assertEqual(image_param.squash_func, image_param._clamp_image)
+        self.assertEqual(image_param.squash_func, clamp_image)
+        assertTensorAlmostEqual(self, image, init_tensor.clamp(0, 1))
 
     def test_natural_image_init_tensor_pixelimage_sf_sigmoid(self) -> None:
         if version.parse(torch.__version__) <= version.parse("1.8.0"):
@@ -1137,9 +1145,10 @@ def test_natural_image_0(self) -> None:
         )
 
     def test_natural_image_1(self) -> None:
-        image_param = images.NaturalImage(init=torch.ones(3, 1, 1))
+        init_tensor = torch.ones(3, 1, 1)
+        image_param = images.NaturalImage(init=init_tensor)
         image = image_param.forward().detach()
-        assertTensorAlmostEqual(self, image, torch.ones_like(image), mode="max")
+        assertTensorAlmostEqual(self, image, torch.sigmoid(init_tensor).unsqueeze(0))
 
     def test_natural_image_cuda(self) -> None:
         if not torch.cuda.is_available():
@@ -1166,10 +1175,11 @@ def test_natural_image_jit_module_init_tensor(self) -> None:
                 "Skipping NaturalImage init tensor JIT module test due to"
                 + " insufficient Torch version."
             )
-        image_param = images.NaturalImage(init=torch.ones(1, 3, 1, 1))
+        init_tensor = torch.ones(1, 3, 1, 1)
+        image_param = images.NaturalImage(init=init_tensor)
         jit_image_param = torch.jit.script(image_param)
         output_tensor = jit_image_param()
-        assertTensorAlmostEqual(self, output_tensor, torch.ones_like(output_tensor))
+        assertTensorAlmostEqual(self, output_tensor, torch.sigmoid(init_tensor))
 
     def test_natural_image_jit_module_init_tensor_pixelimage(self) -> None:
         if version.parse(torch.__version__) <= version.parse("1.8.0"):
@@ -1177,12 +1187,13 @@ def test_natural_image_jit_module_init_tensor_pixelimage(self) -> None:
                 "Skipping NaturalImage PixelImage init tensor JIT module"
                 + " test due to insufficient Torch version."
             )
+        init_tensor = torch.ones(1, 3, 1, 1)
         image_param = images.NaturalImage(
-            init=torch.ones(1, 3, 1, 1), parameterization=images.PixelImage
+            init=init_tensor, parameterization=images.PixelImage
         )
         jit_image_param = torch.jit.script(image_param)
         output_tensor = jit_image_param()
-        assertTensorAlmostEqual(self, output_tensor, torch.ones_like(output_tensor))
+        assertTensorAlmostEqual(self, output_tensor, torch.sigmoid(init_tensor))
 
     def test_natural_image_decorrelation_module_none(self) -> None:
         if version.parse(torch.__version__) <= version.parse("1.8.0"):
@@ -1190,9 +1201,8 @@ def test_natural_image_decorrelation_module_none(self) -> None:
                 "Skipping NaturalImage no decorrelation module"
                 + " test due to insufficient Torch version."
             )
-        image_param = images.NaturalImage(
-            init=torch.ones(1, 3, 4, 4), decorrelation_module=None
-        )
+        init_tensor = torch.ones(1, 3, 1, 1)
+        image_param = images.NaturalImage(init=init_tensor, decorrelation_module=None)
         image = image_param.forward().detach()
         self.assertIsNone(image_param.decorrelate)
-        assertTensorAlmostEqual(self, image, torch.ones_like(image))
+        assertTensorAlmostEqual(self, image, torch.sigmoid(init_tensor))

From c8f9cc54f5ff7578dbb68fd4e4880fbfe9662692 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Thu, 9 Jun 2022 13:21:07 -0600
Subject: [PATCH 03/30] Remove old commented out TransformationRobustness
 version

---
 captum/optim/_param/image/transforms.py | 18 ------------------
 1 file changed, 18 deletions(-)

diff --git a/captum/optim/_param/image/transforms.py b/captum/optim/_param/image/transforms.py
index fbf7b01cbc..6828c17670 100644
--- a/captum/optim/_param/image/transforms.py
+++ b/captum/optim/_param/image/transforms.py
@@ -939,24 +939,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x[:, [2, 1, 0]]
 
 
-# class TransformationRobustness(nn.Module):
-#     def __init__(self, jitter=False, scale=False):
-#         super().__init__()
-#         if jitter:
-#             self.jitter = RandomSpatialJitter(4)
-#         if scale:
-#             self.scale = RandomScale()
-
-#     def forward(self, x):
-#         original_shape = x.shape
-#         if hasattr(self, "jitter"):
-#             x = self.jitter(x)
-#         if hasattr(self, "scale"):
-#             x = self.scale(x)
-#         cropped = center_crop(x, original_shape)
-#         return cropped
-
-
 # class RandomHomography(nn.Module):
 #     def __init__(self):
 #         super().__init__()

From 44203fa8b9fa27798811c816119260842138c919 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Sat, 18 Jun 2022 14:42:35 -0600
Subject: [PATCH 04/30] Fix torch.meshgrid warning

---
 captum/optim/_param/image/transforms.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/captum/optim/_param/image/transforms.py b/captum/optim/_param/image/transforms.py
index 6828c17670..714eb20236 100644
--- a/captum/optim/_param/image/transforms.py
+++ b/captum/optim/_param/image/transforms.py
@@ -993,9 +993,18 @@ def __init__(
         # The gaussian kernel is the product of the
         # gaussian function of each dimension.
         kernel = 1
-        meshgrids = torch.meshgrid(
-            [torch.arange(size, dtype=torch.float32) for size in kernel_size]
-        )
+
+        # PyTorch v1.10.0 adds a new indexing argument
+        if version.parse(torch.__version__) >= version.parse("1.10.0"):
+            meshgrids = torch.meshgrid(
+                [torch.arange(size, dtype=torch.float32) for size in kernel_size],
+                indexing="ij",
+            )
+        else:
+            meshgrids = torch.meshgrid(
+                [torch.arange(size, dtype=torch.float32) for size in kernel_size]
+            )
+
         for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
             mean = (size - 1) / 2
             kernel *= (

From 9e2f9537810937ba34400f999bd00e1930e30cd3 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Sat, 18 Jun 2022 15:04:32 -0600
Subject: [PATCH 05/30] self.assertEquals -> self.assertEqual

---
 tests/optim/utils/test_reducer.py | 38 +++++++++++++++----------------
 1 file changed, 19 insertions(+), 19 deletions(-)

diff --git a/tests/optim/utils/test_reducer.py b/tests/optim/utils/test_reducer.py
index f2baa7675c..4c97ef2440 100644
--- a/tests/optim/utils/test_reducer.py
+++ b/tests/optim/utils/test_reducer.py
@@ -34,10 +34,10 @@ def test_channelreducer_pytorch(self) -> None:
         test_input = torch.randn(1, 32, 224, 224).abs()
         c_reducer = reducer.ChannelReducer(n_components=3, max_iter=100)
         test_output = c_reducer.fit_transform(test_input)
-        self.assertEquals(test_output.size(0), 1)
-        self.assertEquals(test_output.size(1), 3)
-        self.assertEquals(test_output.size(2), 224)
-        self.assertEquals(test_output.size(3), 224)
+        self.assertEqual(test_output.size(0), 1)
+        self.assertEqual(test_output.size(1), 3)
+        self.assertEqual(test_output.size(2), 224)
+        self.assertEqual(test_output.size(3), 224)
 
     def test_channelreducer_pytorch_dim_three(self) -> None:
         try:
@@ -52,9 +52,9 @@ def test_channelreducer_pytorch_dim_three(self) -> None:
         test_input = torch.randn(32, 224, 224).abs()
         c_reducer = reducer.ChannelReducer(n_components=3, max_iter=100)
         test_output = c_reducer.fit_transform(test_input)
-        self.assertEquals(test_output.size(0), 3)
-        self.assertEquals(test_output.size(1), 224)
-        self.assertEquals(test_output.size(2), 224)
+        self.assertEqual(test_output.size(0), 3)
+        self.assertEqual(test_output.size(1), 224)
+        self.assertEqual(test_output.size(2), 224)
 
     def test_channelreducer_pytorch_pca(self) -> None:
         try:
@@ -70,10 +70,10 @@ def test_channelreducer_pytorch_pca(self) -> None:
         c_reducer = reducer.ChannelReducer(n_components=3, reduction_alg="PCA")
 
         test_output = c_reducer.fit_transform(test_input)
-        self.assertEquals(test_output.size(0), 1)
-        self.assertEquals(test_output.size(1), 3)
-        self.assertEquals(test_output.size(2), 224)
-        self.assertEquals(test_output.size(3), 224)
+        self.assertEqual(test_output.size(0), 1)
+        self.assertEqual(test_output.size(1), 3)
+        self.assertEqual(test_output.size(2), 224)
+        self.assertEqual(test_output.size(3), 224)
 
     def test_channelreducer_pytorch_custom_alg(self) -> None:
         test_input = torch.randn(1, 32, 224, 224).abs()
@@ -82,10 +82,10 @@ def test_channelreducer_pytorch_custom_alg(self) -> None:
             n_components=3, reduction_alg=reduction_alg, max_iter=100
         )
         test_output = c_reducer.fit_transform(test_input)
-        self.assertEquals(test_output.size(0), 1)
-        self.assertEquals(test_output.size(1), 3)
-        self.assertEquals(test_output.size(2), 224)
-        self.assertEquals(test_output.size(3), 224)
+        self.assertEqual(test_output.size(0), 1)
+        self.assertEqual(test_output.size(1), 3)
+        self.assertEqual(test_output.size(2), 224)
+        self.assertEqual(test_output.size(3), 224)
 
     def test_channelreducer_pytorch_custom_alg_components(self) -> None:
         reduction_alg = FakeReductionAlgorithm
@@ -149,10 +149,10 @@ def test_channelreducer_noreshape_pytorch(self) -> None:
         test_input = torch.randn(1, 224, 224, 32).abs()
         c_reducer = reducer.ChannelReducer(n_components=3, max_iter=100)
         test_output = c_reducer.fit_transform(test_input, swap_2nd_and_last_dims=False)
-        self.assertEquals(test_output.size(0), 1)
-        self.assertEquals(test_output.size(1), 224)
-        self.assertEquals(test_output.size(2), 224)
-        self.assertEquals(test_output.size(3), 3)
+        self.assertEqual(test_output.size(0), 1)
+        self.assertEqual(test_output.size(1), 224)
+        self.assertEqual(test_output.size(2), 224)
+        self.assertEqual(test_output.size(3), 3)
 
     def test_channelreducer_error(self) -> None:
         if not torch.cuda.is_available():

From f867bf3cea02f077196f767c7fbdcf156011b9a1 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Wed, 22 Jun 2022 14:05:39 -0600
Subject: [PATCH 06/30] Resolve register_backward_hook ->
 register_full_backward_hook depreciation

---
 tests/optim/models/test_models_common.py | 6 +++++-
 tests/optim/param/test_transforms.py     | 9 ++++++---
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/tests/optim/models/test_models_common.py b/tests/optim/models/test_models_common.py
index 5c10060760..04e80f3a24 100644
--- a/tests/optim/models/test_models_common.py
+++ b/tests/optim/models/test_models_common.py
@@ -5,6 +5,7 @@
 import captum.optim.models._common as model_utils
 import torch
 import torch.nn.functional as F
+from packaging import version
 from captum.optim.models import googlenet
 from tests.helpers.basic import BaseTest, assertTensorAlmostEqual
 
@@ -37,7 +38,10 @@ def check_grad(self, grad_input, grad_output):
 
         rr_layer = model_utils.RedirectedReluLayer()
         x = torch.zeros(1, 3, 4, 4, requires_grad=True)
-        rr_layer.register_backward_hook(check_grad)
+        if version.parse(torch.__version__) >= version.parse("1.8.0"):
+            rr_layer.register_full_backward_hook(check_grad)
+        else:
+            rr_layer.register_backward_hook(check_grad)
         rr_loss = rr_layer(x * 1).mean()
         rr_loss.backward()
 
diff --git a/tests/optim/param/test_transforms.py b/tests/optim/param/test_transforms.py
index 911330cb75..7afa0b7723 100644
--- a/tests/optim/param/test_transforms.py
+++ b/tests/optim/param/test_transforms.py
@@ -1830,12 +1830,15 @@ def check_grad(self, grad_input, grad_output):
 
         class SymmetricPaddingLayer(torch.nn.Module):
             def forward(
-                self, x: torch.Tensor, padding: List[List[int]]
+                self, x_input: torch.Tensor, padding: List[List[int]]
             ) -> torch.Tensor:
-                return transforms.SymmetricPadding.apply(x_pt, padding)
+                return transforms.SymmetricPadding.apply(x_input, padding)
 
         sym_pad = SymmetricPaddingLayer()
-        sym_pad.register_backward_hook(check_grad)
+        if version.parse(torch.__version__) >= version.parse("1.8.0"):
+            sym_pad.register_full_backward_hook(check_grad)
+        else:
+            sym_pad.register_backward_hook(check_grad)
         x_out = sym_pad(x_pt, offset_pad)
         (x_out.sum() * 1).backward()
 

From 54b652d71d767548f01be3af2446e430b5b1415b Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Wed, 22 Jun 2022 14:18:39 -0600
Subject: [PATCH 07/30] Fix lint errors

---
 tests/optim/models/test_models_common.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/optim/models/test_models_common.py b/tests/optim/models/test_models_common.py
index 04e80f3a24..11856e44e8 100644
--- a/tests/optim/models/test_models_common.py
+++ b/tests/optim/models/test_models_common.py
@@ -5,8 +5,8 @@
 import captum.optim.models._common as model_utils
 import torch
 import torch.nn.functional as F
-from packaging import version
 from captum.optim.models import googlenet
+from packaging import version
 from tests.helpers.basic import BaseTest, assertTensorAlmostEqual
 
 

From dd58b750f072569e48465f15d80a153fa6d160ba Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Mon, 27 Jun 2022 15:00:22 -0600
Subject: [PATCH 08/30] Improve ImageParameterization docs for Sphinx

---
 captum/optim/_param/image/__init__.py |   2 +-
 captum/optim/_param/image/images.py   | 215 +++++++++++++++++++++-----
 tests/optim/param/test_images.py      |   7 +
 3 files changed, 181 insertions(+), 43 deletions(-)

diff --git a/captum/optim/_param/image/__init__.py b/captum/optim/_param/image/__init__.py
index a2311f7c46..5c36c0c80f 100755
--- a/captum/optim/_param/image/__init__.py
+++ b/captum/optim/_param/image/__init__.py
@@ -1 +1 @@
-"""(Differentiable) Input Parameterizations. Currently only 3-channel images"""
+"""(Differentiable) Input Parameterizations. Currently only images"""
diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index da697c66c1..86d1d6eeb3 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -180,12 +180,32 @@ def forward(self) -> torch.Tensor:
 
 
 class ImageParameterization(InputParameterization):
+    r"""The base class for all Image Parameterizations"""
     pass
 
 
 class FFTImage(ImageParameterization):
     """
     Parameterize an image using inverse real 2D FFT
+
+    Example::
+
+        >>> fft_image = opt.images.FFTImage(size=(224, 224))
+        >>> output_image = fft_image()
+        >>> print(output_image.required_grad)
+        True
+        >>> print(output_image.shape)
+        torch.Size([1, 3, 224, 224])
+
+    Example for using an initialization tensor::
+
+        >>> init = torch.randn(1, 3, 224, 224)
+        >>> fft_image = opt.images.FFTImage(init=init)
+        >>> output_image = fft_image()
+        >>> print(output_image.required_grad)
+        True
+        >>> print(output_image.shape)
+        torch.Size([1, 3, 224, 224])
     """
 
     __constants__ = ["size"]
@@ -203,13 +223,13 @@ def __init__(
             size (Tuple[int, int]): The height & width dimensions to use for the
                 parameterized output image tensor.
             channels (int, optional): The number of channels to use for each image.
-                Default: 3
+                Default: ``3``
             batch (int, optional): The number of images to stack along the batch
                 dimension.
-                Default: 1
+                Default: ``1``
             init (torch.tensor, optional): Optionally specify a tensor to
                 use instead of creating one.
-                Default: None
+                Default: ``None``
         """
         super().__init__()
         if init is None:
@@ -332,6 +352,25 @@ def forward(self) -> torch.Tensor:
 class PixelImage(ImageParameterization):
     """
     Parameterize a simple pixel image tensor that requires no additional transforms.
+
+    Example::
+
+        >>> pixel_image = opt.images.PixelImage(size=(224, 224))
+        >>> output_image = pixel_image()
+        >>> print(output_image.required_grad)
+        True
+        >>> print(output_image.shape)
+        torch.Size([1, 3, 224, 224])
+
+    Example for using an initialization tensor::
+
+        >>> init = torch.randn(1, 3, 224, 224)
+        >>> pixel_image = opt.images.PixelImage(init=init)
+        >>> output_image = pixel_image()
+        >>> print(output_image.required_grad)
+        True
+        >>> print(output_image.shape)
+        torch.Size([1, 3, 224, 224])
     """
 
     def __init__(
@@ -347,13 +386,13 @@ def __init__(
             size (Tuple[int, int]): The height & width dimensions to use for the
                 parameterized output image tensor.
             channels (int, optional): The number of channels to use for each image.
-                Default: 3
+                Default: ``3``
             batch (int, optional): The number of images to stack along the batch
                 dimension.
-                Default: 1
+                Default: ``1``
             init (torch.tensor, optional): Optionally specify a tensor to
                 use instead of creating one.
-                Default: None
+                Default: ``None``
         """
         super().__init__()
         if init is None:
@@ -374,6 +413,25 @@ def forward(self) -> torch.Tensor:
 class LaplacianImage(ImageParameterization):
     """
     Parameterize an image tensor with a laplacian pyramid.
+
+    Example::
+
+        >>> laplacian_image = opt.images.LaplacianImage(size=(224, 224))
+        >>> output_image = laplacian_image()
+        >>> print(output_image.required_grad)
+        True
+        >>> print(output_image.shape)
+        torch.Size([1, 3, 224, 224])
+
+    Example for using an initialization tensor::
+
+        >>> init = torch.randn(1, 3, 224, 224)
+        >>> laplacian_image = opt.images.LaplacianImage(init=init)
+        >>> output_image = laplacian_image()
+        >>> print(output_image.required_grad)
+        True
+        >>> print(output_image.shape)
+        torch.Size([1, 3, 224, 224])
     """
 
     def __init__(
@@ -388,25 +446,25 @@ def __init__(
         """
         Args:
 
-            size (Tuple[int, int]): The height & width dimensions to use for the
-                parameterized output image tensor.
-                Default: (224, 224)
+            size (Tuple[int, int], optional): The height & width dimensions to use for
+                the parameterized output image tensor.
+                Default: ``(224, 224)``
             channels (int, optional): The number of channels to use for each image.
-                Default: 3
+                Default: ``3``
             batch (int, optional): The number of images to stack along the batch
                 dimension.
-                Default: 1
+                Default: ``1``
             init (torch.tensor, optional): Optionally specify a tensor to
                 use instead of creating one.
-                Default: None
+                Default: ``None``
             power (float, optional): The desired power value to use.
-                Default: 0.1
+                Default: ``0.1``
             scale_list (list of float, optional): The desired list of scale values to
                 use in the laplacian pyramid. The height & width dimensions specified
-                in size or used in the init tensor should be divisable by every scale
-                value in the scale list with no remainder left over. The default
-                scale_list values are set to work with a size of (224, 224).
-                Default: [1.0, 2.0, 4.0, 8.0, 16.0, 32.0]
+                in ``size`` or used in the ``init`` tensor should be divisable by every
+                scale value in the scale list with no remainder left over. The default
+                scale_list values are set to work with a ``size`` of ``(224, 224)``.
+                Default: ``[1.0, 2.0, 4.0, 8.0, 16.0, 32.0]``
         """
         super().__init__()
         if init is not None:
@@ -494,6 +552,17 @@ class SharedImage(ImageParameterization):
 
     Mordvintsev, et al., "Differentiable Image Parameterizations", Distill, 2018.
     https://distill.pub/2018/differentiable-parameterizations/
+
+    Example::
+
+        >>> fft_image = opt.images.FFTImage(size=(224, 224), batch=2)
+        >>> shared_shapes = ((1, 3, 64, 64), (4, 3, 32, 32))
+        >>> shared_image = opt.images.SharedImage(shared_shapes, fft_image)
+        >>> output_image = shared_image()
+        >>> print(output_image.required_grad)
+        True
+        >>> print(output_image.shape)
+        torch.Size([2, 3, 224, 224])
     """
 
     __constants__ = ["offset"]
@@ -513,7 +582,7 @@ def __init__(
                 instance.
             offset (int or list of int or list of list of ints , optional): The offsets
                 to use for the shared tensors.
-                Default: None
+                Default: ``None``
         """
         super().__init__()
         assert shapes is not None
@@ -706,6 +775,28 @@ def forward(self) -> torch.Tensor:
 class StackImage(ImageParameterization):
     """
     Stack multiple NCHW image parameterizations along their batch dimensions.
+
+    Example::
+
+        >>> fft_image_1 = opt.images.FFTImage(size=(224, 224), batch=1)
+        >>> fft_image_2 = opt.images.FFTImage(size=(224, 224), batch=1)
+        >>> stack_image = opt.images.StackImage([fft_image_1, fft_image_2])
+        >>> output_image = stack_image()
+        >>> print(output_image.required_grad)
+        True
+        >>> print(output_image.shape)
+        torch.Size([2, 3, 224, 224])
+
+    Example with ``ImageParameterization`` & ``torch.Tensor``::
+
+        >>> fft_image = opt.images.FFTImage(size=(224, 224), batch=1)
+        >>> tensor_image = torch.randn(1, 3, 224, 224)
+        >>> stack_image = opt.images.StackImage([fft_image, tensor_image])
+        >>> output_image = stack_image()
+        >>> print(output_image.required_grad)
+        True
+        >>> print(output_image.shape)
+        torch.Size([2, 3, 224, 224])
     """
 
     __constants__ = ["dim", "output_device"]
@@ -723,12 +814,12 @@ def __init__(
                  of image parameterizations to stack across their batch dimensions.
             dim (int, optional): Optionally specify the dim to concatinate
                  parameterization outputs on. Default is set to the batch dimension.
-                 Default: 0
+                 Default: ``0``
             output_device (torch.device, optional): If the parameterizations are on
                 different devices, then their outputs will be moved to the device
-                specified by this variable. Default is set to None with the expectation
-                that all parameterizations are on the same device.
-                Default: None
+                specified by this variable. Default is set to ``None`` with the
+                expectation that all parameterization outputs are on the same device.
+                Default: ``None``
         """
         super().__init__()
         assert len(parameterizations) > 0
@@ -771,16 +862,50 @@ def forward(self) -> torch.Tensor:
 
 
 class NaturalImage(ImageParameterization):
-    r"""Outputs an optimizable input image.
+    r"""Outputs an optimizable input image wrapped in :class:`.ImageTensor`.
 
-    By convention, single images are CHW and float32s in [0,1].
-    The underlying parameterization can be decorrelated via a ToRGB transform.
+    By convention, single images are CHW and float32s in [0, 1].
+    The underlying parameterization can be decorrelated via a ``ToRGB`` transform.
     When used with the (default) FFT parameterization, this results in a fully
     uncorrelated image parameterization. :-)
 
     If a model requires a normalization step, such as normalizing imagenet RGB values,
-    or rescaling to [0,255], it can perform those steps with the provided transforms or
-    inside its computation.
+    or rescaling to [0, 255], it can perform those steps with the provided transforms or
+    inside its module class.
+
+    Example::
+
+        >>> image = opt.images.NaturalImage(size=(224, 224), channels=3, batch=1)
+        >>> image_tensor = image()
+        >>> print(image_tensor.required_grad)
+        True
+        >>> print(image_tensor.shape)
+        torch.Size([1, 3, 224, 224])
+
+    Example for using an initialization tensor::
+
+        >>> init = torch.randn(1, 3, 224, 224)
+        >>> image = opt.images.NaturalImage(init=init)
+        >>> image_tensor = image()
+        >>> print(image_tensor.required_grad)
+        True
+        >>> print(image_tensor.shape)
+        torch.Size([1, 3, 224, 224])
+
+    Example for using a parameterization::
+
+        >>> fft_image = opt.images.FFTImage(size=(224, 224), channels=3, batch=1)
+        >>> image = opt.images.NaturalImage(parameterization=fft_image)
+        >>> image_tensor = image()
+        >>> print(image_tensor.required_grad)
+        True
+        >>> print(image_tensor.shape)
+        torch.Size([1, 3, 224, 224])
+
+    :ivar parameterization: initial value (ImageParameterization): The given image
+        parameterization instance given when initializing ``NaturalImage``.
+    :ivar decorrelation_module: initial value (nn.Module): The given decorrelation
+        module instance given when initializing ``NaturalImage``.
     """
 
     def __init__(
@@ -800,36 +925,37 @@ def __init__(
             size (Tuple[int, int], optional): The height and width to use for the
                 nn.Parameter image tensor. This parameter is not used if
                 parameterization is an instance.
-                Default: (224, 224)
+                Default: ``(224, 224)``
             channels (int, optional): The number of channels to use when creating the
                 nn.Parameter tensor. This parameter is not used if parameterization is
                 an instance.
-                Default: 3
+                Default: ``3``
             batch (int, optional): The number of channels to use when creating the
-                nn.Parameter tensor, or stacking init images. This parameter is not
-                used if parameterization is an instance.
-                Default: 1
+                nn.Parameter tensor. This parameter is not used if ``parameterization``
+                is an instance.
+                Default: ``1``
             init (torch.tensor, optional): Optionally specify a tensor to use instead
                 of creating one from random noise. This parameter is not used if
-                parameterization is an instance. Set to None for random init.
-                Default: None
+                ``parameterization`` is an instance. Set to ``None`` for random init.
+                Default: ``None``
             parameterization (ImageParameterization, optional): An image
                 parameterization class, or instance of an image parameterization class.
                 Default: FFTImage
             squash_func (Callable[[torch.Tensor], torch.Tensor]], optional): The squash
                 function to use after color recorrelation. A function, lambda function,
                 or callable class instance.
-                Default: None
+                Default: ``None``
             decorrelation_module (nn.Module, optional): A module instance that
                 recorrelates the colors of an input image. Custom modules can make use
-                of the decorrelate_init parameter by having a second inverse parameter
-                in their forward functions that performs the inverse operation when it
-                is set to True. Set to None for no recorrelation.
-                Default: ToRGB
+                of the ``decorrelate_init`` parameter by having a second ``inverse``
+                parameter in their forward functions that performs the inverse
+                operation when it is set to ``True`` (see ``ToRGB`` for an example).
+                Set to ``None`` for no recorrelation.
+                Default: ``ToRGB``
             decorrelate_init (bool, optional): Whether or not to apply color
                 decorrelation to the init tensor input. This parameter is not used if
-                parameterization is an instance or if init is None.
-                Default: True
+                ``parameterization`` is an instance or if init is ``None``.
+                Default: ``True``
         """
         super().__init__()
         if not isinstance(parameterization, ImageParameterization):
@@ -866,11 +992,16 @@ def _to_image_tensor(self, x: torch.Tensor) -> torch.Tensor:
             x (torch.tensor): An input tensor.
 
         Returns:
-            x (ImageTensor): An instance of ImageTensor with the input tensor.
+            x (ImageTensor): An instance of ``ImageTensor`` with the input tensor.
         """
         return ImageTensor(x)
 
     def forward(self) -> torch.Tensor:
+        """
+        Returns:
+            image_tensor (torch.Tensor): The parameterization output wrapped in
+                ``ImageTensor``, that has optionally had its colors recorrelated.
+        """
         image = self.parameterization()
         if self.decorrelate is not None:
             image = self.decorrelate(image)
diff --git a/tests/optim/param/test_images.py b/tests/optim/param/test_images.py
index 567af0e5e1..eac6a1050b 100644
--- a/tests/optim/param/test_images.py
+++ b/tests/optim/param/test_images.py
@@ -33,6 +33,13 @@ def test_new_list(self) -> None:
         self.assertTrue(torch.is_tensor(test_tensor))
         self.assertEqual(x.shape, test_tensor.shape)
 
+    def test_new_with_grad(self) -> None:
+        x = torch.ones(5, requires_grad=True)
+        test_tensor = images.ImageTensor(x)
+        self.assertTrue(test_tensor.requires_grad)
+        self.assertTrue(torch.is_tensor(test_tensor))
+        self.assertEqual(x.shape, test_tensor.shape)
+
     def test_torch_function(self) -> None:
         x = torch.ones(5)
         image_tensor = images.ImageTensor(x)

From aafc4f7dfe9bf5bb433c074d188750f59abe16d3 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Tue, 28 Jun 2022 08:49:07 -0600
Subject: [PATCH 09/30] Improve GaussianSmoothing

---
 captum/optim/_param/image/transforms.py | 13 +++++++------
 tests/optim/param/test_transforms.py    | 15 ++++++++-------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/captum/optim/_param/image/transforms.py b/captum/optim/_param/image/transforms.py
index 714eb20236..36318f9709 100644
--- a/captum/optim/_param/image/transforms.py
+++ b/captum/optim/_param/image/transforms.py
@@ -969,7 +969,7 @@ def __init__(
         kernel_size: Union[int, Sequence[int]],
         sigma: Union[float, Sequence[float]],
         dim: int = 2,
-        use_same_padding: bool = True,
+        padding: Union[str, int, Tuple[int, int]] = "same",
     ) -> None:
         """
         Args:
@@ -979,10 +979,11 @@ def __init__(
             kernel_size (int, sequence): Size of the gaussian kernel.
             sigma (float, sequence): Standard deviation of the gaussian kernel.
             dim (int, optional): The number of dimensions of the data.
-                Default value is 2 (spatial).
-            use_same_padding (bool, optional): Whether or not to use "same" padding so
-                that the output shape is the same as the input shape.
-                Default: True
+                Default value is ``2`` for (spatial)
+            padding (str, int or list of tuple, optional): The desired padding amount
+                or mode to use. One of; ``"valid"``, ``"same"``, a single number, or a
+                tuple in the format of: (padH, padW).
+                Default: ``"same"``
         """
         super().__init__()
         if isinstance(kernel_size, numbers.Number):
@@ -1022,7 +1023,7 @@ def __init__(
 
         self.register_buffer("weight", kernel)
         self.groups = channels
-        self.padding = "same" if use_same_padding else 0
+        self.padding = padding
 
         if dim == 1:
             self.conv = F.conv1d
diff --git a/tests/optim/param/test_transforms.py b/tests/optim/param/test_transforms.py
index 7afa0b7723..38b436d4ff 100644
--- a/tests/optim/param/test_transforms.py
+++ b/tests/optim/param/test_transforms.py
@@ -1586,9 +1586,10 @@ def test_gaussian_smoothing_init_1d(self) -> None:
             kernel_size,
             sigma,
             dim,
-            use_same_padding=False,
+            padding=0,
         )
         self.assertEqual(smoothening_module.groups, channels)
+        self.assertEqual(smoothening_module.padding, 0)
         weight = torch.tensor([[0.3192, 0.3617, 0.3192]]).repeat(6, 1, 1)
         assertTensorAlmostEqual(self, smoothening_module.weight, weight, 0.001)
         self.assertFalse(smoothening_module.padding)
@@ -1603,7 +1604,7 @@ def test_gaussian_smoothing_init_2d(self) -> None:
             kernel_size,
             sigma,
             dim,
-            use_same_padding=False,
+            padding=0,
         )
         self.assertEqual(smoothening_module.groups, channels)
         weight = torch.tensor(
@@ -1627,7 +1628,7 @@ def test_gaussian_smoothing_init_3d(self) -> None:
             kernel_size,
             sigma,
             dim,
-            use_same_padding=False,
+            padding=0,
         )
         self.assertEqual(smoothening_module.groups, channels)
         weight = torch.tensor(
@@ -1671,7 +1672,7 @@ def test_gaussian_smoothing_1d(self) -> None:
             kernel_size,
             sigma,
             dim,
-            use_same_padding=False,
+            padding=0,
         )
 
         test_tensor = torch.tensor([1.0, 5.0]).repeat(6, 2).unsqueeze(0)
@@ -1692,7 +1693,7 @@ def test_gaussian_smoothing_2d(self) -> None:
             kernel_size,
             sigma,
             dim,
-            use_same_padding=False,
+            padding=0,
         )
 
         test_tensor = torch.tensor([1.0, 5.0]).repeat(3, 6, 3).unsqueeze(0)
@@ -1713,7 +1714,7 @@ def test_gaussian_smoothing_3d(self) -> None:
             kernel_size,
             sigma,
             dim,
-            use_same_padding=False,
+            padding=0,
         )
 
         test_tensor = torch.tensor([1.0, 5.0, 1.0]).repeat(4, 6, 6, 2).unsqueeze(0)
@@ -1741,7 +1742,7 @@ def test_gaussian_smoothing_2d_jit_module(self) -> None:
             kernel_size,
             sigma,
             dim,
-            use_same_padding=False,
+            padding=0,
         )
         jit_smoothening_module = torch.jit.script(smoothening_module)
 

From 59351197cb712366c2b0b9acae0a72758c4f2c78 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Tue, 28 Jun 2022 12:43:23 -0600
Subject: [PATCH 10/30] Add underscore to some FFTImage functions

---
 captum/optim/_param/image/images.py | 12 ++++++------
 tests/optim/param/test_images.py    |  4 ++--
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index 86d1d6eeb3..2020bad370 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -240,9 +240,9 @@ def __init__(
             if init.dim() == 3:
                 init = init.unsqueeze(0)
             self.size = (init.size(2), init.size(3))
-        self.torch_rfft, self.torch_irfft, self.torch_fftfreq = self.get_fft_funcs()
+        self.torch_rfft, self.torch_irfft, self.torch_fftfreq = self._get_fft_funcs()
 
-        frequencies = self.rfft2d_freqs(*self.size)
+        frequencies = self._rfft2d_freqs(*self.size)
         scale = 1.0 / torch.max(
             frequencies,
             torch.full_like(frequencies, 1.0 / (max(self.size[0], self.size[1]))),
@@ -269,7 +269,7 @@ def __init__(
         self.register_buffer("spectrum_scale", spectrum_scale)
         self.fourier_coeffs = nn.Parameter(fourier_coeffs)
 
-    def rfft2d_freqs(self, height: int, width: int) -> torch.Tensor:
+    def _rfft2d_freqs(self, height: int, width: int) -> torch.Tensor:
         """
         Computes 2D spectrum frequencies.
 
@@ -287,12 +287,12 @@ def rfft2d_freqs(self, height: int, width: int) -> torch.Tensor:
         return torch.sqrt((fx * fx) + (fy * fy))
 
     @torch.jit.export
-    def torch_irfftn(self, x: torch.Tensor) -> torch.Tensor:
+    def _torch_irfftn(self, x: torch.Tensor) -> torch.Tensor:
         if x.dtype != torch.complex64:
             x = torch.view_as_complex(x)
         return torch.fft.irfftn(x, s=self.size)  # type: ignore
 
-    def get_fft_funcs(self) -> Tuple[Callable, Callable, Callable]:
+    def _get_fft_funcs(self) -> Tuple[Callable, Callable, Callable]:
         """
         Support older versions of PyTorch. This function ensures that the same FFT
         operations are carried regardless of whether your PyTorch version has the
@@ -311,7 +311,7 @@ def get_fft_funcs(self) -> Tuple[Callable, Callable, Callable]:
             def torch_rfft(x: torch.Tensor) -> torch.Tensor:
                 return torch.view_as_real(torch.fft.rfftn(x, s=self.size))
 
-            torch_irfftn = self.torch_irfftn
+            torch_irfftn = self._torch_irfftn
 
             def torch_fftfreq(v: int, d: float = 1.0) -> torch.Tensor:
                 return torch.fft.fftfreq(v, d)
diff --git a/tests/optim/param/test_images.py b/tests/optim/param/test_images.py
index eac6a1050b..236acaf5b4 100644
--- a/tests/optim/param/test_images.py
+++ b/tests/optim/param/test_images.py
@@ -109,7 +109,7 @@ def test_subclass(self) -> None:
 
     def test_pytorch_fftfreq(self) -> None:
         image = images.FFTImage((1, 1))
-        _, _, fftfreq = image.get_fft_funcs()
+        _, _, fftfreq = image._get_fft_funcs()
         assertTensorAlmostEqual(
             self, fftfreq(4, 4), torch.as_tensor(np.fft.fftfreq(4, 4)), mode="max"
         )
@@ -121,7 +121,7 @@ def test_rfft2d_freqs(self) -> None:
 
         assertTensorAlmostEqual(
             self,
-            image.rfft2d_freqs(height, width),
+            image._rfft2d_freqs(height, width),
             torch.tensor([[0.0000, 0.3333], [0.5000, 0.6009]]),
         )
 

From c1161c5d9b5997c8bb5f012901fa91bda32a62ec Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Wed, 29 Jun 2022 11:05:28 -0600
Subject: [PATCH 11/30] Improve NaturalImage docs

---
 captum/optim/_param/image/images.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index 2020bad370..06611a3107 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -865,13 +865,14 @@ class NaturalImage(ImageParameterization):
     r"""Outputs an optimizable input image wrapped in :class:`.ImageTensor`.
 
     By convention, single images are CHW and float32s in [0, 1].
-    The underlying parameterization can be decorrelated via a ``ToRGB`` transform.
-    When used with the (default) FFT parameterization, this results in a fully
-    uncorrelated image parameterization. :-)
+    The underlying parameterization can be decorrelated via a
+    :class:`captum.optim.transforms.ToRGB` transform.
+    When used with the (default) :class:`.FFTImage` parameterization, this results in
+    a fully uncorrelated image parameterization. :-)
 
     If a model requires a normalization step, such as normalizing imagenet RGB values,
-    or rescaling to [0, 255], it can perform those steps with the provided transforms or
-    inside its module class.
+    or rescaling to [0, 255], it can perform those steps with the provided transforms
+    or inside its module class.
 
     Example::
 
@@ -949,7 +950,8 @@ def __init__(
                 recorrelates the colors of an input image. Custom modules can make use
                 of the ``decorrelate_init`` parameter by having a second ``inverse``
                 parameter in their forward functions that performs the inverse
-                operation when it is set to ``True`` (see ``ToRGB`` for an example).
+                operation when it is set to ``True`` (see
+                :class:`captum.optim.transforms.ToRGB` for an example).
                 Set to ``None`` for no recorrelation.
                 Default: ``ToRGB``
             decorrelate_init (bool, optional): Whether or not to apply color
@@ -1000,7 +1002,8 @@ def forward(self) -> torch.Tensor:
         """
         Returns:
             image_tensor (torch.Tensor): The parameterization output wrapped in
-                ``ImageTensor``, that has optionally had its colors recorrelated.
+                :class:`.ImageTensor`, that has optionally had its colors
+                recorrelated.
         """
         image = self.parameterization()
         if self.decorrelate is not None:

From eb6930e2bd2209695016a6aa389bc6e09ca405a1 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Fri, 1 Jul 2022 18:52:44 -0600
Subject: [PATCH 12/30] Improve ImageParameterization docs (#551)

* Update images.py

* Update images.py
---
 captum/optim/_param/image/images.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index 06611a3107..6f40931f64 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -998,10 +998,10 @@ def _to_image_tensor(self, x: torch.Tensor) -> torch.Tensor:
         """
         return ImageTensor(x)
 
-    def forward(self) -> torch.Tensor:
+    def forward(self) -> ImageTensor:
         """
         Returns:
-            image_tensor (torch.Tensor): The parameterization output wrapped in
+            image_tensor (ImageTensor): The parameterization output wrapped in
                 :class:`.ImageTensor`, that has optionally had its colors
                 recorrelated.
         """

From 55cad28e485591712bff7fd9ba089d7230bcb073 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Mon, 4 Jul 2022 18:17:30 -0600
Subject: [PATCH 13/30] Improve GaussianSmoothing docs

---
 captum/optim/_param/image/transforms.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/captum/optim/_param/image/transforms.py b/captum/optim/_param/image/transforms.py
index 36318f9709..8131f4fc19 100644
--- a/captum/optim/_param/image/transforms.py
+++ b/captum/optim/_param/image/transforms.py
@@ -1045,7 +1045,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
             input (torch.Tensor): Input to apply gaussian filter on.
 
         Returns:
-            **filtered** (torch.Tensor): Filtered output.
+            filtered (torch.Tensor): Filtered output.
         """
         return self.conv(
             input, weight=self.weight, groups=self.groups, padding=self.padding

From fbaa8eb639dc13c7e0ef4fe3614ecc588bd28ca3 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Wed, 6 Jul 2022 15:29:03 -0600
Subject: [PATCH 14/30] Add missing forward docs to PixelImage

---
 captum/optim/_param/image/images.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index 6f40931f64..fa81c29ebe 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -279,7 +279,7 @@ def _rfft2d_freqs(self, height: int, width: int) -> torch.Tensor:
             width (int): The w dimension of the 2d frequency scale.
 
         Returns:
-            **tensor** (tensor): A 2d frequency scale tensor.
+            tensor (torch.Tensor): A 2d frequency scale tensor.
         """
 
         fy = self.torch_fftfreq(height)[:, None]
@@ -339,7 +339,7 @@ def torch_fftfreq(v: int, d: float = 1.0) -> torch.Tensor:
     def forward(self) -> torch.Tensor:
         """
         Returns:
-            **output** (torch.tensor): A spatially recorrelated tensor.
+            output (torch.tensor): A spatially recorrelated NCHW tensor.
         """
 
         scaled_spectrum = self.fourier_coeffs * self.spectrum_scale
@@ -405,6 +405,10 @@ def __init__(
         self.image = nn.Parameter(init)
 
     def forward(self) -> torch.Tensor:
+        """
+        Returns:
+            output (torch.tensor): An NCHW tensor.
+        """
         if torch.jit.is_scripting():
             return self.image
         return self.image.refine_names("B", "C", "H", "W")
@@ -497,7 +501,7 @@ def __init__(
     def forward(self) -> torch.Tensor:
         """
         Returns:
-            **output** (torch.tensor): A tensor created from a laplacian pyramid.
+            output (torch.tensor): An NCHW tensor created from a laplacian pyramid.
         """
         A = []
         for xi, upsamplei in zip(self.tensor_params, self.scaler):
@@ -611,7 +615,7 @@ def _get_offset(self, offset: Union[int, Tuple[int]], n: int) -> List[List[int]]
             n (int): The number of tensors needing offset values.
 
         Returns:
-            **offset** (list of list of int): A list of offset values.
+            offset (list of list of int): A list of offset values.
         """
         if type(offset) is tuple or type(offset) is list:
             if type(offset[0]) is tuple or type(offset[0]) is list:
@@ -635,7 +639,7 @@ def _apply_offset(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:
             x_list (list of torch.Tensor): list of tensors to offset.
 
         Returns:
-            **A** (list of torch.Tensor): list of offset tensors.
+            A (list of torch.Tensor): list of offset tensors.
         """
 
         A: List[torch.Tensor] = []

From 5457544d2d670dc57c61d64cbed5839e7e0b2b28 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Thu, 7 Jul 2022 09:19:30 -0600
Subject: [PATCH 15/30] Improve FFTImage float dtype support

---
 captum/optim/_param/image/images.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index fa81c29ebe..e1a10896c0 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -288,7 +288,7 @@ def _rfft2d_freqs(self, height: int, width: int) -> torch.Tensor:
 
     @torch.jit.export
     def _torch_irfftn(self, x: torch.Tensor) -> torch.Tensor:
-        if x.dtype != torch.complex64:
+        if not torch.is_complex(x):
             x = torch.view_as_complex(x)
         return torch.fft.irfftn(x, s=self.size)  # type: ignore
 

From 45c85117495611cc6c56e6a32faf448bdb6f4223 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Thu, 7 Jul 2022 11:56:02 -0600
Subject: [PATCH 16/30] Add dtype tests for ImageParameterizations

---
 captum/optim/_param/image/images.py |  7 ++-
 tests/optim/param/test_images.py    | 70 +++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index e1a10896c0..451dd239f9 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -36,7 +36,12 @@ def __new__(
         Returns:
            x (ImageTensor): An `ImageTensor` instance.
         """
-        if isinstance(x, torch.Tensor) and x.is_cuda:
+        if (
+            isinstance(x, torch.Tensor)
+            and x.is_cuda
+            or isinstance(x, torch.Tensor)
+            and x.dtype != torch.float32
+        ):
             x.show = MethodType(cls.show, x)
             x.export = MethodType(cls.export, x)
             return x
diff --git a/tests/optim/param/test_images.py b/tests/optim/param/test_images.py
index 236acaf5b4..8051fab162 100644
--- a/tests/optim/param/test_images.py
+++ b/tests/optim/param/test_images.py
@@ -20,6 +20,17 @@ def test_new(self) -> None:
         test_tensor = images.ImageTensor(x)
         self.assertTrue(torch.is_tensor(test_tensor))
         self.assertEqual(x.shape, test_tensor.shape)
+        self.assertEqual(x.dtype, test_tensor.dtype)
+
+    def test_new_dtype_float64(self) -> None:
+        x = torch.ones(5, dtype=torch.float64)
+        test_tensor = images.ImageTensor(x)
+        self.assertEqual(test_tensor.dtype, torch.float64)
+
+    def test_new_dtype_float16(self) -> None:
+        x = torch.ones(5, dtype=torch.float16)
+        test_tensor = images.ImageTensor(x)
+        self.assertEqual(test_tensor.dtype, torch.float16)
 
     def test_new_numpy(self) -> None:
         x = torch.ones(5).numpy()
@@ -315,6 +326,30 @@ def test_fftimage_forward_init_batch(self) -> None:
             self, fftimage_tensor.detach(), fftimage_array, 25.0, mode="max"
         )
 
+    def test_fftimage_forward_dtype_float64(self) -> None:
+        image_param = images.FFTImage(size=(224, 224)).to(dtype=torch.float64)
+        output = image_param()
+        self.assertEqual(output.dtype, torch.float64)
+
+    def test_fftimage_forward_dtype_float32(self) -> None:
+        image_param = images.FFTImage(size=(224, 224)).to(dtype=torch.float32)
+        output = image_param()
+        self.assertEqual(output.dtype, torch.float32)
+
+    def test_fftimage_forward_dtype_float16(self) -> None:
+        if version.parse(torch.__version__) <= version.parse("1.12.0"):
+            raise unittest.SkipTest(
+                "Skipping FFTImage float16 dtype test due to"
+                + "  insufficient Torch version."
+            )
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest(
+                "Skipping FFTImage float16 dtype test due to not supporting CUDA."
+            )
+        image_param = images.FFTImage(size=(256, 256)).to(dtype=torch.float16)
+        output = image_param()
+        self.assertEqual(output.dtype, torch.float16)
+
 
 class TestPixelImage(BaseTest):
     def test_subclass(self) -> None:
@@ -383,6 +418,21 @@ def test_pixelimage_init_forward(self) -> None:
         self.assertEqual(test_tensor.size(3), size[1])
         assertTensorAlmostEqual(self, test_tensor, init_tensor[None, :], 0)
 
+    def test_pixelimage_forward_dtype_float64(self) -> None:
+        image_param = images.PixelImage(size=(224, 224)).to(dtype=torch.float64)
+        output = image_param()
+        self.assertEqual(output.dtype, torch.float64)
+
+    def test_pixelimage_forward_dtype_float32(self) -> None:
+        image_param = images.PixelImage(size=(224, 224)).to(dtype=torch.float32)
+        output = image_param()
+        self.assertEqual(output.dtype, torch.float32)
+
+    def test_pixelimage_forward_dtype_float16(self) -> None:
+        image_param = images.PixelImage(size=(224, 224)).to(dtype=torch.float16)
+        output = image_param()
+        self.assertEqual(output.dtype, torch.float16)
+
 
 class TestLaplacianImage(BaseTest):
     def test_subclass(self) -> None:
@@ -440,6 +490,16 @@ def test_laplacianimage_random_forward_cuda(self) -> None:
         self.assertEqual(list(test_tensor.shape), [1, 3, 224, 224])
         self.assertTrue(test_tensor.requires_grad)
 
+    def test_laplcianimage_forward_dtype_float64(self) -> None:
+        image_param = images.LaplacianImage(size=(224, 224)).to(dtype=torch.float64)
+        output = image_param()
+        self.assertEqual(output.dtype, torch.float64)
+
+    def test_laplcianimage_forward_dtype_float32(self) -> None:
+        image_param = images.LaplacianImage(size=(224, 224)).to(dtype=torch.float32)
+        output = image_param()
+        self.assertEqual(output.dtype, torch.float32)
+
 
 class TestSimpleTensorParameterization(BaseTest):
     def test_subclass(self) -> None:
@@ -1213,3 +1273,13 @@ def test_natural_image_decorrelation_module_none(self) -> None:
         image = image_param.forward().detach()
         self.assertIsNone(image_param.decorrelate)
         assertTensorAlmostEqual(self, image, torch.sigmoid(init_tensor))
+
+    def test_natural_image_forward_dtype_float64(self) -> None:
+        image_param = images.NaturalImage(size=(224, 224)).to(dtype=torch.float64)
+        output = image_param()
+        self.assertEqual(output.dtype, torch.float64)
+
+    def test_natural_image_forward_dtype_float32(self) -> None:
+        image_param = images.NaturalImage(size=(224, 224)).to(dtype=torch.float32)
+        output = image_param()
+        self.assertEqual(output.dtype, torch.float32)

From 3dc061d7a28ebf5c7bec125ab6d2e39769bbf990 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Thu, 7 Jul 2022 12:12:21 -0600
Subject: [PATCH 17/30] Fix weird error: RuntimeError: expected scalar type
 Float but found Double

It seems to occur in ToRGB on 3 different tests
---
 tests/optim/param/test_images.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/optim/param/test_images.py b/tests/optim/param/test_images.py
index 8051fab162..62980a14fa 100644
--- a/tests/optim/param/test_images.py
+++ b/tests/optim/param/test_images.py
@@ -1193,7 +1193,7 @@ def test_natural_image_init_tensor_pixelimage_sf_sigmoid(self) -> None:
                 + " test due to insufficient Torch version."
             )
         image_param = images.NaturalImage(
-            init=torch.ones(1, 3, 1, 1),
+            init=torch.ones(1, 3, 1, 1).float(),
             parameterization=images.PixelImage,
             squash_func=torch.sigmoid,
         )
@@ -1242,7 +1242,7 @@ def test_natural_image_jit_module_init_tensor(self) -> None:
                 "Skipping NaturalImage init tensor JIT module test due to"
                 + " insufficient Torch version."
             )
-        init_tensor = torch.ones(1, 3, 1, 1)
+        init_tensor = torch.ones(1, 3, 1, 1).float()
         image_param = images.NaturalImage(init=init_tensor)
         jit_image_param = torch.jit.script(image_param)
         output_tensor = jit_image_param()
@@ -1254,7 +1254,7 @@ def test_natural_image_jit_module_init_tensor_pixelimage(self) -> None:
                 "Skipping NaturalImage PixelImage init tensor JIT module"
                 + " test due to insufficient Torch version."
             )
-        init_tensor = torch.ones(1, 3, 1, 1)
+        init_tensor = torch.ones(1, 3, 1, 1).float()
         image_param = images.NaturalImage(
             init=init_tensor, parameterization=images.PixelImage
         )

From 90d9fd121f4e1018ac0e00b170e8af4998243544 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Thu, 7 Jul 2022 12:35:58 -0600
Subject: [PATCH 18/30] Fix test failures

---
 tests/optim/param/test_images.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/optim/param/test_images.py b/tests/optim/param/test_images.py
index 62980a14fa..baae2e674d 100644
--- a/tests/optim/param/test_images.py
+++ b/tests/optim/param/test_images.py
@@ -1196,7 +1196,7 @@ def test_natural_image_init_tensor_pixelimage_sf_sigmoid(self) -> None:
             init=torch.ones(1, 3, 1, 1).float(),
             parameterization=images.PixelImage,
             squash_func=torch.sigmoid,
-        )
+        ).to(dtype=torch.float32)
         output_tensor = image_param()
 
         self.assertEqual(image_param.squash_func, torch.sigmoid)
@@ -1242,7 +1242,7 @@ def test_natural_image_jit_module_init_tensor(self) -> None:
                 "Skipping NaturalImage init tensor JIT module test due to"
                 + " insufficient Torch version."
             )
-        init_tensor = torch.ones(1, 3, 1, 1).float()
+        init_tensor = torch.ones(1, 3, 1, 1)
         image_param = images.NaturalImage(init=init_tensor)
         jit_image_param = torch.jit.script(image_param)
         output_tensor = jit_image_param()
@@ -1254,7 +1254,7 @@ def test_natural_image_jit_module_init_tensor_pixelimage(self) -> None:
                 "Skipping NaturalImage PixelImage init tensor JIT module"
                 + " test due to insufficient Torch version."
             )
-        init_tensor = torch.ones(1, 3, 1, 1).float()
+        init_tensor = torch.ones(1, 3, 1, 1)
         image_param = images.NaturalImage(
             init=init_tensor, parameterization=images.PixelImage
         )
@@ -1275,11 +1275,13 @@ def test_natural_image_decorrelation_module_none(self) -> None:
         assertTensorAlmostEqual(self, image, torch.sigmoid(init_tensor))
 
     def test_natural_image_forward_dtype_float64(self) -> None:
+        raise unittest.SkipTest("Skipping test due to bug")
         image_param = images.NaturalImage(size=(224, 224)).to(dtype=torch.float64)
         output = image_param()
         self.assertEqual(output.dtype, torch.float64)
 
     def test_natural_image_forward_dtype_float32(self) -> None:
+        raise unittest.SkipTest("Skipping test due to bug")
         image_param = images.NaturalImage(size=(224, 224)).to(dtype=torch.float32)
         output = image_param()
         self.assertEqual(output.dtype, torch.float32)

From dd13dc63944f61d1b8e87e89507b202daf2018e2 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Sat, 9 Jul 2022 09:19:25 -0600
Subject: [PATCH 19/30] NaturalImage dtype test fix + transform dtype tests

---
 tests/optim/param/test_images.py     |  77 ++++++++----
 tests/optim/param/test_transforms.py | 169 +++++++++++++++++++++++++++
 2 files changed, 226 insertions(+), 20 deletions(-)

diff --git a/tests/optim/param/test_images.py b/tests/optim/param/test_images.py
index baae2e674d..2a336393ce 100644
--- a/tests/optim/param/test_images.py
+++ b/tests/optim/param/test_images.py
@@ -327,14 +327,16 @@ def test_fftimage_forward_init_batch(self) -> None:
         )
 
     def test_fftimage_forward_dtype_float64(self) -> None:
-        image_param = images.FFTImage(size=(224, 224)).to(dtype=torch.float64)
+        dtype = torch.float64
+        image_param = images.FFTImage(size=(224, 224)).to(dtype=dtype)
         output = image_param()
-        self.assertEqual(output.dtype, torch.float64)
+        self.assertEqual(output.dtype, dtype)
 
     def test_fftimage_forward_dtype_float32(self) -> None:
-        image_param = images.FFTImage(size=(224, 224)).to(dtype=torch.float32)
+        dtype = torch.float32
+        image_param = images.FFTImage(size=(224, 224)).to(dtype=dtype)
         output = image_param()
-        self.assertEqual(output.dtype, torch.float32)
+        self.assertEqual(output.dtype, dtype)
 
     def test_fftimage_forward_dtype_float16(self) -> None:
         if version.parse(torch.__version__) <= version.parse("1.12.0"):
@@ -342,13 +344,14 @@ def test_fftimage_forward_dtype_float16(self) -> None:
                 "Skipping FFTImage float16 dtype test due to"
                 + "  insufficient Torch version."
             )
+        dtype = torch.float16
         if not torch.cuda.is_available():
             raise unittest.SkipTest(
                 "Skipping FFTImage float16 dtype test due to not supporting CUDA."
             )
-        image_param = images.FFTImage(size=(256, 256)).to(dtype=torch.float16)
+        image_param = images.FFTImage(size=(256, 256)).to(dtype=dtype)
         output = image_param()
-        self.assertEqual(output.dtype, torch.float16)
+        self.assertEqual(output.dtype, dtype)
 
 
 class TestPixelImage(BaseTest):
@@ -419,19 +422,28 @@ def test_pixelimage_init_forward(self) -> None:
         assertTensorAlmostEqual(self, test_tensor, init_tensor[None, :], 0)
 
     def test_pixelimage_forward_dtype_float64(self) -> None:
-        image_param = images.PixelImage(size=(224, 224)).to(dtype=torch.float64)
+        dtype = torch.float64
+        image_param = images.PixelImage(size=(224, 224)).to(dtype=dtype)
         output = image_param()
         self.assertEqual(output.dtype, torch.float64)
 
     def test_pixelimage_forward_dtype_float32(self) -> None:
-        image_param = images.PixelImage(size=(224, 224)).to(dtype=torch.float32)
+        dtype = torch.float32
+        image_param = images.PixelImage(size=(224, 224)).to(dtype=dtype)
         output = image_param()
         self.assertEqual(output.dtype, torch.float32)
 
     def test_pixelimage_forward_dtype_float16(self) -> None:
-        image_param = images.PixelImage(size=(224, 224)).to(dtype=torch.float16)
+        dtype = torch.float16
+        image_param = images.PixelImage(size=(224, 224)).to(dtype)
+        output = image_param()
+        self.assertEqual(output.dtype, dtype)
+
+    def test_pixelimage_forward_dtype_bfloat16(self) -> None:
+        dtype = torch.bfloat16
+        image_param = images.PixelImage(size=(224, 224)).to(dtype=dtype)
         output = image_param()
-        self.assertEqual(output.dtype, torch.float16)
+        self.assertEqual(output.dtype, dtype)
 
 
 class TestLaplacianImage(BaseTest):
@@ -491,14 +503,16 @@ def test_laplacianimage_random_forward_cuda(self) -> None:
         self.assertTrue(test_tensor.requires_grad)
 
     def test_laplcianimage_forward_dtype_float64(self) -> None:
-        image_param = images.LaplacianImage(size=(224, 224)).to(dtype=torch.float64)
+        dtype = torch.float64
+        image_param = images.LaplacianImage(size=(224, 224)).to(dtype=dtype)
         output = image_param()
-        self.assertEqual(output.dtype, torch.float64)
+        self.assertEqual(output.dtype, dtype)
 
     def test_laplcianimage_forward_dtype_float32(self) -> None:
-        image_param = images.LaplacianImage(size=(224, 224)).to(dtype=torch.float32)
+        dtype = torch.float32
+        image_param = images.LaplacianImage(size=(224, 224)).to(dtype=dtype)
         output = image_param()
-        self.assertEqual(output.dtype, torch.float32)
+        self.assertEqual(output.dtype, dtype)
 
 
 class TestSimpleTensorParameterization(BaseTest):
@@ -1275,13 +1289,36 @@ def test_natural_image_decorrelation_module_none(self) -> None:
         assertTensorAlmostEqual(self, image, torch.sigmoid(init_tensor))
 
     def test_natural_image_forward_dtype_float64(self) -> None:
-        raise unittest.SkipTest("Skipping test due to bug")
-        image_param = images.NaturalImage(size=(224, 224)).to(dtype=torch.float64)
+        dtype = torch.float64
+        image_param = images.NaturalImage(
+            size=(224, 224), decorrelation_module=ToRGB("klt")
+        ).to(dtype=dtype)
         output = image_param()
-        self.assertEqual(output.dtype, torch.float64)
+        self.assertEqual(output.dtype, dtype)
 
     def test_natural_image_forward_dtype_float32(self) -> None:
-        raise unittest.SkipTest("Skipping test due to bug")
-        image_param = images.NaturalImage(size=(224, 224)).to(dtype=torch.float32)
+        dtype = torch.float32
+        image_param = images.NaturalImage(
+            size=(224, 224), decorrelation_module=ToRGB("klt")
+        ).to(dtype=dtype)
         output = image_param()
-        self.assertEqual(output.dtype, torch.float32)
+        self.assertEqual(output.dtype, dtype)
+
+    def test_fftimage_forward_dtype_float16(self) -> None:
+        if version.parse(torch.__version__) <= version.parse("1.12.0"):
+            raise unittest.SkipTest(
+                "Skipping NaturalImage float16 dtype test due to"
+                + "  insufficient Torch version."
+            )
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest(
+                "Skipping NaturalImage float16 dtype test due to not supporting CUDA."
+            )
+        dtype = torch.float16
+        image_param = (
+            images.NaturalImage(size=(256, 256), decorrelation_module=ToRGB("klt"))
+            .cuda()
+            .to(dtype=dtype)
+        )
+        output = image_param()
+        self.assertEqual(output.dtype, dtype)
diff --git a/tests/optim/param/test_transforms.py b/tests/optim/param/test_transforms.py
index 38b436d4ff..baddd5ac3e 100644
--- a/tests/optim/param/test_transforms.py
+++ b/tests/optim/param/test_transforms.py
@@ -261,6 +261,24 @@ def test_random_scale_jit_module(self) -> None:
             0,
         )
 
+    def test_random_scale_dtype_float64(self) -> None:
+        dtype = torch.float64
+        scale_module = transforms.RandomScale(scale=[0.975, 1.025, 0.95, 1.05]).to(
+            dtype=dtype
+        )
+        x = torch.ones([1, 3, 224, 224], dtype=dtype)
+        output = scale_module(x)
+        self.assertEqual(output.dtype, dtype)
+
+    def test_random_scale_dtype_float32(self) -> None:
+        dtype = torch.float32
+        scale_module = transforms.RandomScale(scale=[0.975, 1.025, 0.95, 1.05]).to(
+            dtype=dtype
+        )
+        x = torch.ones([1, 3, 224, 224], dtype=dtype)
+        output = scale_module(x)
+        self.assertEqual(output.dtype, dtype)
+
 
 class TestRandomScaleAffine(BaseTest):
     def test_random_scale_affine_init(self) -> None:
@@ -430,6 +448,40 @@ def test_random_scale_affine_jit_module(self) -> None:
             0,
         )
 
+    def test_random_scale_affine_dtype_float64(self) -> None:
+        dtype = torch.float64
+        scale_module = transforms.RandomScaleAffine(
+            scale=[0.975, 1.025, 0.95, 1.05]
+        ).to(dtype=dtype)
+        x = torch.ones([1, 3, 224, 224], dtype=dtype)
+        output = scale_module(x)
+        self.assertEqual(output.dtype, dtype)
+
+    def test_random_scale_affine_dtype_float32(self) -> None:
+        dtype = torch.float32
+        scale_module = transforms.RandomScaleAffine(
+            scale=[0.975, 1.025, 0.95, 1.05]
+        ).to(dtype=dtype)
+        x = torch.ones([1, 3, 224, 224], dtype=dtype)
+        output = scale_module(x)
+        self.assertEqual(output.dtype, dtype)
+
+    def test_random_scale_affine_dtype_float16(self) -> None:
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest(
+                "Skipping RandomScaleAffine float16 dtype test due to not supporting"
+                + " CUDA."
+            )
+        dtype = torch.float16
+        scale_module = (
+            transforms.RandomScaleAffine(scale=[0.975, 1.025, 0.95, 1.05])
+            .cuda()
+            .to(dtype=dtype)
+        )
+        x = torch.ones([1, 3, 224, 224], dtype=dtype).cuda()
+        output = scale_module(x)
+        self.assertEqual(output.dtype, dtype)
+
 
 class TestRandomRotation(BaseTest):
     def test_random_rotation_init(self) -> None:
@@ -629,6 +681,37 @@ def test_random_rotation_jit_module(self) -> None:
         )
         assertTensorAlmostEqual(self, test_output, expected_output, 0.005)
 
+    def test_random_rotation_dtype_float64(self) -> None:
+        dtype = torch.float64
+        degrees = list(range(-25, -5)) + list(range(5, 25))
+        rotation_module = transforms.RandomRotation(degrees=degrees).to(dtype=dtype)
+        x = torch.ones([1, 3, 224, 224], dtype=dtype)
+        output = rotation_module(x)
+        self.assertEqual(output.dtype, dtype)
+
+    def test_random_rotation_dtype_float32(self) -> None:
+        dtype = torch.float32
+        degrees = list(range(-25, -5)) + list(range(5, 25))
+        rotation_module = transforms.RandomRotation(degrees=degrees).to(dtype=dtype)
+        x = torch.ones([1, 3, 224, 224], dtype=dtype)
+        output = rotation_module(x)
+        self.assertEqual(output.dtype, dtype)
+
+    def test_random_rotation_dtype_float16(self) -> None:
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest(
+                "Skipping RandomRotation float16 dtype test due to not supporting"
+                + " CUDA."
+            )
+        dtype = torch.float16
+        degrees = list(range(-25, -5)) + list(range(5, 25))
+        rotation_module = (
+            transforms.RandomRotation(degrees=degrees).cuda().to(dtype=dtype)
+        )
+        x = torch.ones([1, 3, 224, 224], dtype=dtype).cuda()
+        output = rotation_module(x)
+        self.assertEqual(output.dtype, dtype)
+
 
 class TestRandomSpatialJitter(BaseTest):
     def test_random_spatial_jitter_init(self) -> None:
@@ -714,6 +797,20 @@ def test_random_spatial_jitter_forward_jit_module(self) -> None:
         jittered_tensor = jit_spatialjitter(test_input)
         self.assertEqual(list(jittered_tensor.shape), list(test_input.shape))
 
+    def test_random_spatial_jitter_dtype_float64(self) -> None:
+        dtype = torch.float64
+        spatialjitter = transforms.RandomSpatialJitter(5).to(dtype=dtype)
+        x = torch.ones([1, 3, 224, 224], dtype=dtype)
+        output = spatialjitter(x)
+        self.assertEqual(output.dtype, dtype)
+
+    def test_random_spatial_jitter_dtype_float32(self) -> None:
+        dtype = torch.float32
+        spatialjitter = transforms.RandomSpatialJitter(5).to(dtype=dtype)
+        x = torch.ones([1, 3, 224, 224], dtype=dtype)
+        output = spatialjitter(x)
+        self.assertEqual(output.dtype, dtype)
+
 
 class TestCenterCrop(BaseTest):
     def test_center_crop_init(self) -> None:
@@ -1574,6 +1671,64 @@ def test_to_rgb_klt_forward_jit_module(self) -> None:
             self, inverse_tensor, torch.ones_like(inverse_tensor.rename(None))
         )
 
+    def test_to_rgb_dtype_float64(self) -> None:
+        dtype = torch.float64
+        to_rgb = transforms.ToRGB(transform="klt").to(dtype=dtype)
+        test_tensor = torch.ones(1, 3, 224, 224, dtype=dtype)
+        output = to_rgb(test_tensor.refine_names("B", "C", "H", "W"))
+        self.assertEqual(output.dtype, dtype)
+        inverse_output = to_rgb(output, inverse=True)
+        self.assertEqual(inverse_output.dtype, dtype)
+
+    def test_to_rgb_dtype_float32(self) -> None:
+        dtype = torch.float32
+        to_rgb = transforms.ToRGB(transform="klt").to(dtype=dtype)
+        test_tensor = torch.ones(1, 3, 224, 224, dtype=dtype)
+        output = to_rgb(test_tensor.refine_names("B", "C", "H", "W"))
+        self.assertEqual(output.dtype, dtype)
+        inverse_output = to_rgb(output, inverse=True)
+        self.assertEqual(inverse_output.dtype, dtype)
+
+    def test_to_rgb_dtype_float16(self) -> None:
+        dtype = torch.float16
+        to_rgb = transforms.ToRGB(transform="klt").to(dtype=dtype)
+        test_tensor = torch.ones(1, 3, 224, 224, dtype=dtype)
+        output = to_rgb(test_tensor.refine_names("B", "C", "H", "W"))
+        self.assertEqual(output.dtype, dtype)
+
+    def test_to_rgb_dtype_float16_cuda(self) -> None:
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest(
+                "Skipping ToRGB float16 dtype test due to not supporting CUDA."
+            )
+        dtype = torch.float16
+        to_rgb = transforms.ToRGB(transform="klt").cuda().to(dtype=dtype)
+        test_tensor = torch.ones(1, 3, 224, 224, dtype=dtype).cuda()
+        output = to_rgb(test_tensor.refine_names("B", "C", "H", "W"))
+        self.assertEqual(output.dtype, dtype)
+        inverse_output = to_rgb(output, inverse=True)
+        self.assertEqual(inverse_output.dtype, dtype)
+
+    def test_to_rgb_dtype_bfloat16(self) -> None:
+        dtype = torch.bfloat16
+        to_rgb = transforms.ToRGB(transform="klt").to(dtype=dtype)
+        test_tensor = torch.ones(1, 3, 224, 224, dtype=dtype)
+        output = to_rgb(test_tensor.refine_names("B", "C", "H", "W"))
+        self.assertEqual(output.dtype, dtype)
+
+    def test_to_rgb_dtype_bfloat16_cuda(self) -> None:
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest(
+                "Skipping ToRGB bfloat16 dtype test due to not supporting CUDA."
+            )
+        dtype = torch.bfloat16
+        to_rgb = transforms.ToRGB(transform="klt").cuda().to(dtype=dtype)
+        test_tensor = torch.ones(1, 3, 224, 224, dtype=dtype).cuda()
+        output = to_rgb(test_tensor.refine_names("B", "C", "H", "W"))
+        self.assertEqual(output.dtype, dtype)
+        inverse_output = to_rgb(output, inverse=True)
+        self.assertEqual(inverse_output.dtype, dtype)
+
 
 class TestGaussianSmoothing(BaseTest):
     def test_gaussian_smoothing_init_1d(self) -> None:
@@ -2041,3 +2196,17 @@ def test_transform_robustness_forward_padding_crop_output_jit_module(self) -> No
         test_input = torch.ones(1, 3, 224, 224)
         test_output = transform_robustness(test_input)
         self.assertEqual(test_output.shape, test_input.shape)
+
+    def test_transform_robustness_dtype_float64(self) -> None:
+        dtype = torch.float64
+        transform_robustness = transforms.TransformationRobustness().to(dtype=dtype)
+        x = torch.ones([1, 3, 224, 224], dtype=dtype)
+        output = transform_robustness(x)
+        self.assertEqual(output.dtype, dtype)
+
+    def test_transform_robustness_dtype_float32(self) -> None:
+        dtype = torch.float32
+        transform_robustness = transforms.TransformationRobustness().to(dtype=dtype)
+        x = torch.ones([1, 3, 224, 224], dtype=dtype)
+        output = transform_robustness(x)
+        self.assertEqual(output.dtype, dtype)

From 1ffcad41d8a2fbc86bf624c9b58724f141ad4d53 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Sat, 9 Jul 2022 09:52:42 -0600
Subject: [PATCH 20/30] Fix dtype tests

---
 tests/optim/param/test_images.py     |  2 +-
 tests/optim/param/test_transforms.py | 18 ------------------
 2 files changed, 1 insertion(+), 19 deletions(-)

diff --git a/tests/optim/param/test_images.py b/tests/optim/param/test_images.py
index 2a336393ce..dea07c2d81 100644
--- a/tests/optim/param/test_images.py
+++ b/tests/optim/param/test_images.py
@@ -349,7 +349,7 @@ def test_fftimage_forward_dtype_float16(self) -> None:
             raise unittest.SkipTest(
                 "Skipping FFTImage float16 dtype test due to not supporting CUDA."
             )
-        image_param = images.FFTImage(size=(256, 256)).to(dtype=dtype)
+        image_param = images.FFTImage(size=(256, 256)).cuda().to(dtype=dtype)
         output = image_param()
         self.assertEqual(output.dtype, dtype)
 
diff --git a/tests/optim/param/test_transforms.py b/tests/optim/param/test_transforms.py
index baddd5ac3e..54f69e8532 100644
--- a/tests/optim/param/test_transforms.py
+++ b/tests/optim/param/test_transforms.py
@@ -1689,13 +1689,6 @@ def test_to_rgb_dtype_float32(self) -> None:
         inverse_output = to_rgb(output, inverse=True)
         self.assertEqual(inverse_output.dtype, dtype)
 
-    def test_to_rgb_dtype_float16(self) -> None:
-        dtype = torch.float16
-        to_rgb = transforms.ToRGB(transform="klt").to(dtype=dtype)
-        test_tensor = torch.ones(1, 3, 224, 224, dtype=dtype)
-        output = to_rgb(test_tensor.refine_names("B", "C", "H", "W"))
-        self.assertEqual(output.dtype, dtype)
-
     def test_to_rgb_dtype_float16_cuda(self) -> None:
         if not torch.cuda.is_available():
             raise unittest.SkipTest(
@@ -1706,15 +1699,6 @@ def test_to_rgb_dtype_float16_cuda(self) -> None:
         test_tensor = torch.ones(1, 3, 224, 224, dtype=dtype).cuda()
         output = to_rgb(test_tensor.refine_names("B", "C", "H", "W"))
         self.assertEqual(output.dtype, dtype)
-        inverse_output = to_rgb(output, inverse=True)
-        self.assertEqual(inverse_output.dtype, dtype)
-
-    def test_to_rgb_dtype_bfloat16(self) -> None:
-        dtype = torch.bfloat16
-        to_rgb = transforms.ToRGB(transform="klt").to(dtype=dtype)
-        test_tensor = torch.ones(1, 3, 224, 224, dtype=dtype)
-        output = to_rgb(test_tensor.refine_names("B", "C", "H", "W"))
-        self.assertEqual(output.dtype, dtype)
 
     def test_to_rgb_dtype_bfloat16_cuda(self) -> None:
         if not torch.cuda.is_available():
@@ -1726,8 +1710,6 @@ def test_to_rgb_dtype_bfloat16_cuda(self) -> None:
         test_tensor = torch.ones(1, 3, 224, 224, dtype=dtype).cuda()
         output = to_rgb(test_tensor.refine_names("B", "C", "H", "W"))
         self.assertEqual(output.dtype, dtype)
-        inverse_output = to_rgb(output, inverse=True)
-        self.assertEqual(inverse_output.dtype, dtype)
 
 
 class TestGaussianSmoothing(BaseTest):

From 2a0a898afbb631886b8b9809b8ad4b4a5f071a5f Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Sat, 9 Jul 2022 10:21:27 -0600
Subject: [PATCH 21/30] Remove failing test

---
 tests/optim/param/test_transforms.py | 11 -----------
 1 file changed, 11 deletions(-)

diff --git a/tests/optim/param/test_transforms.py b/tests/optim/param/test_transforms.py
index 54f69e8532..3568b4c53a 100644
--- a/tests/optim/param/test_transforms.py
+++ b/tests/optim/param/test_transforms.py
@@ -1700,17 +1700,6 @@ def test_to_rgb_dtype_float16_cuda(self) -> None:
         output = to_rgb(test_tensor.refine_names("B", "C", "H", "W"))
         self.assertEqual(output.dtype, dtype)
 
-    def test_to_rgb_dtype_bfloat16_cuda(self) -> None:
-        if not torch.cuda.is_available():
-            raise unittest.SkipTest(
-                "Skipping ToRGB bfloat16 dtype test due to not supporting CUDA."
-            )
-        dtype = torch.bfloat16
-        to_rgb = transforms.ToRGB(transform="klt").cuda().to(dtype=dtype)
-        test_tensor = torch.ones(1, 3, 224, 224, dtype=dtype).cuda()
-        output = to_rgb(test_tensor.refine_names("B", "C", "H", "W"))
-        self.assertEqual(output.dtype, dtype)
-
 
 class TestGaussianSmoothing(BaseTest):
     def test_gaussian_smoothing_init_1d(self) -> None:

From 7c833ad5a41b007a0822aa705a81a999c08e352c Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Sat, 9 Jul 2022 10:43:10 -0600
Subject: [PATCH 22/30] Simplify some image parameterization tests

---
 tests/optim/param/test_images.py | 70 +++++---------------------------
 1 file changed, 11 insertions(+), 59 deletions(-)

diff --git a/tests/optim/param/test_images.py b/tests/optim/param/test_images.py
index dea07c2d81..9d23efd374 100644
--- a/tests/optim/param/test_images.py
+++ b/tests/optim/param/test_images.py
@@ -362,12 +362,7 @@ def test_pixelimage_random(self) -> None:
         size = (224, 224)
         channels = 3
         image_param = images.PixelImage(size=size, channels=channels)
-
-        self.assertEqual(image_param.image.dim(), 4)
-        self.assertEqual(image_param.image.size(0), 1)
-        self.assertEqual(image_param.image.size(1), channels)
-        self.assertEqual(image_param.image.size(2), size[0])
-        self.assertEqual(image_param.image.size(3), size[1])
+        self.assertEqual(list(image_param.image.shape), [1, channels] + list(size))
         self.assertTrue(image_param.image.requires_grad)
 
     def test_pixelimage_init(self) -> None:
@@ -376,11 +371,7 @@ def test_pixelimage_init(self) -> None:
         init_tensor = torch.randn(channels, *size)
         image_param = images.PixelImage(size=size, channels=channels, init=init_tensor)
 
-        self.assertEqual(image_param.image.dim(), 4)
-        self.assertEqual(image_param.image.size(0), 1)
-        self.assertEqual(image_param.image.size(1), channels)
-        self.assertEqual(image_param.image.size(2), size[0])
-        self.assertEqual(image_param.image.size(3), size[1])
+        self.assertEqual(list(image_param.image.shape), [1, channels] + list(size))
         assertTensorAlmostEqual(self, image_param.image, init_tensor[None, :], 0)
         self.assertTrue(image_param.image.requires_grad)
 
@@ -389,12 +380,7 @@ def test_pixelimage_random_forward(self) -> None:
         channels = 3
         image_param = images.PixelImage(size=size, channels=channels)
         test_tensor = image_param.forward().rename(None)
-
-        self.assertEqual(test_tensor.dim(), 4)
-        self.assertEqual(test_tensor.size(0), 1)
-        self.assertEqual(test_tensor.size(1), channels)
-        self.assertEqual(test_tensor.size(2), size[0])
-        self.assertEqual(test_tensor.size(3), size[1])
+        self.assertEqual(list(test_tensor.shape), [1, channels] + list(size))
 
     def test_pixelimage_forward_jit_module(self) -> None:
         if version.parse(torch.__version__) <= version.parse("1.8.0"):
@@ -414,11 +400,7 @@ def test_pixelimage_init_forward(self) -> None:
         image_param = images.PixelImage(size=size, channels=channels, init=init_tensor)
         test_tensor = image_param.forward().rename(None)
 
-        self.assertEqual(test_tensor.dim(), 4)
-        self.assertEqual(test_tensor.size(0), 1)
-        self.assertEqual(test_tensor.size(1), channels)
-        self.assertEqual(test_tensor.size(2), size[0])
-        self.assertEqual(test_tensor.size(3), size[1])
+        self.assertEqual(list(test_tensor.shape), [1, channels] + list(size))
         assertTensorAlmostEqual(self, test_tensor, init_tensor[None, :], 0)
 
     def test_pixelimage_forward_dtype_float64(self) -> None:
@@ -789,12 +771,7 @@ def test_interpolate_tensor(self) -> None:
         output_tensor = image_param._interpolate_tensor(
             test_tensor, batch, channels, size[0], size[1]
         )
-
-        self.assertEqual(output_tensor.dim(), 4)
-        self.assertEqual(output_tensor.size(0), batch)
-        self.assertEqual(output_tensor.size(1), channels)
-        self.assertEqual(output_tensor.size(2), size[0])
-        self.assertEqual(output_tensor.size(3), size[1])
+        self.assertEqual(list(output_tensor.shape), [batch, channels] + list(size))
 
     def test_sharedimage_single_shape_hw_forward(self) -> None:
         shared_shapes = (128 // 2, 128 // 2)
@@ -812,11 +789,7 @@ def test_sharedimage_single_shape_hw_forward(self) -> None:
         self.assertEqual(
             list(image_param.shared_init[0]().shape), [1, 1] + list(shared_shapes)
         )
-        self.assertEqual(test_tensor.dim(), 4)
-        self.assertEqual(test_tensor.size(0), batch)
-        self.assertEqual(test_tensor.size(1), channels)
-        self.assertEqual(test_tensor.size(2), size[0])
-        self.assertEqual(test_tensor.size(3), size[1])
+        self.assertEqual(list(test_tensor.shape), [batch, channels] + list(size))
 
     def test_sharedimage_single_shape_chw_forward(self) -> None:
         shared_shapes = (3, 128 // 2, 128 // 2)
@@ -834,11 +807,7 @@ def test_sharedimage_single_shape_chw_forward(self) -> None:
         self.assertEqual(
             list(image_param.shared_init[0]().shape), [1] + list(shared_shapes)
         )
-        self.assertEqual(test_tensor.dim(), 4)
-        self.assertEqual(test_tensor.size(0), batch)
-        self.assertEqual(test_tensor.size(1), channels)
-        self.assertEqual(test_tensor.size(2), size[0])
-        self.assertEqual(test_tensor.size(3), size[1])
+        self.assertEqual(list(test_tensor.shape), [batch, channels] + list(size))
 
     def test_sharedimage_single_shape_bchw_forward(self) -> None:
         shared_shapes = (1, 3, 128 // 2, 128 // 2)
@@ -854,11 +823,7 @@ def test_sharedimage_single_shape_bchw_forward(self) -> None:
         self.assertIsNone(image_param.offset)
         self.assertEqual(image_param.shared_init[0]().dim(), 4)
         self.assertEqual(list(image_param.shared_init[0]().shape), list(shared_shapes))
-        self.assertEqual(test_tensor.dim(), 4)
-        self.assertEqual(test_tensor.size(0), batch)
-        self.assertEqual(test_tensor.size(1), channels)
-        self.assertEqual(test_tensor.size(2), size[0])
-        self.assertEqual(test_tensor.size(3), size[1])
+        self.assertEqual(list(test_tensor.shape), [batch, channels] + list(size))
 
     def test_sharedimage_multiple_shapes_forward(self) -> None:
         shared_shapes = (
@@ -884,11 +849,7 @@ def test_sharedimage_multiple_shapes_forward(self) -> None:
             self.assertEqual(
                 list(image_param.shared_init[i]().shape), list(shared_shapes[i])
             )
-        self.assertEqual(test_tensor.dim(), 4)
-        self.assertEqual(test_tensor.size(0), batch)
-        self.assertEqual(test_tensor.size(1), channels)
-        self.assertEqual(test_tensor.size(2), size[0])
-        self.assertEqual(test_tensor.size(3), size[1])
+        self.assertEqual(list(test_tensor.shape), [batch, channels] + list(size))
 
     def test_sharedimage_multiple_shapes_diff_len_forward(self) -> None:
         shared_shapes = (
@@ -915,11 +876,7 @@ def test_sharedimage_multiple_shapes_diff_len_forward(self) -> None:
             s_shape = ([1] * (4 - len(s_shape))) + list(s_shape)
             self.assertEqual(list(image_param.shared_init[i]().shape), s_shape)
 
-        self.assertEqual(test_tensor.dim(), 4)
-        self.assertEqual(test_tensor.size(0), batch)
-        self.assertEqual(test_tensor.size(1), channels)
-        self.assertEqual(test_tensor.size(2), size[0])
-        self.assertEqual(test_tensor.size(3), size[1])
+        self.assertEqual(list(test_tensor.shape), [batch, channels] + list(size))
 
     def test_sharedimage_multiple_shapes_diff_len_forward_jit_module(self) -> None:
         if version.parse(torch.__version__) <= version.parse("1.8.0"):
@@ -946,12 +903,7 @@ def test_sharedimage_multiple_shapes_diff_len_forward_jit_module(self) -> None:
         )
         jit_image_param = torch.jit.script(image_param)
         test_tensor = jit_image_param()
-
-        self.assertEqual(test_tensor.dim(), 4)
-        self.assertEqual(test_tensor.size(0), batch)
-        self.assertEqual(test_tensor.size(1), channels)
-        self.assertEqual(test_tensor.size(2), size[0])
-        self.assertEqual(test_tensor.size(3), size[1])
+        self.assertEqual(list(test_tensor.shape), [batch, channels] + list(size))
 
 
 class TestStackImage(BaseTest):

From 2d81aec0d607e235e23c2433ea0dffe8e4072900 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Fri, 15 Jul 2022 13:42:37 -0600
Subject: [PATCH 23/30] Fix type hints

---
 captum/optim/_param/image/images.py | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index 451dd239f9..bd973e6d5b 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -232,7 +232,7 @@ def __init__(
             batch (int, optional): The number of images to stack along the batch
                 dimension.
                 Default: ``1``
-            init (torch.tensor, optional): Optionally specify a tensor to
+            init (torch.Tensor, optional): Optionally specify a tensor to
                 use instead of creating one.
                 Default: ``None``
         """
@@ -344,7 +344,7 @@ def torch_fftfreq(v: int, d: float = 1.0) -> torch.Tensor:
     def forward(self) -> torch.Tensor:
         """
         Returns:
-            output (torch.tensor): A spatially recorrelated NCHW tensor.
+            output (torch.Tensor): A spatially recorrelated NCHW tensor.
         """
 
         scaled_spectrum = self.fourier_coeffs * self.spectrum_scale
@@ -395,7 +395,7 @@ def __init__(
             batch (int, optional): The number of images to stack along the batch
                 dimension.
                 Default: ``1``
-            init (torch.tensor, optional): Optionally specify a tensor to
+            init (torch.Tensor, optional): Optionally specify a tensor to
                 use instead of creating one.
                 Default: ``None``
         """
@@ -412,7 +412,7 @@ def __init__(
     def forward(self) -> torch.Tensor:
         """
         Returns:
-            output (torch.tensor): An NCHW tensor.
+            output (torch.Tensor): An NCHW tensor.
         """
         if torch.jit.is_scripting():
             return self.image
@@ -463,7 +463,7 @@ def __init__(
             batch (int, optional): The number of images to stack along the batch
                 dimension.
                 Default: ``1``
-            init (torch.tensor, optional): Optionally specify a tensor to
+            init (torch.Tensor, optional): Optionally specify a tensor to
                 use instead of creating one.
                 Default: ``None``
             power (float, optional): The desired power value to use.
@@ -506,7 +506,7 @@ def __init__(
     def forward(self) -> torch.Tensor:
         """
         Returns:
-            output (torch.tensor): An NCHW tensor created from a laplacian pyramid.
+            output (torch.Tensor): An NCHW tensor created from a laplacian pyramid.
         """
         A = []
         for xi, upsamplei in zip(self.tensor_params, self.scaler):
@@ -536,7 +536,7 @@ def __init__(self, tensor: torch.Tensor = None) -> None:
         """
         Args:
 
-            tensor (torch.tensor): The tensor to return everytime this module is called.
+            tensor (torch.Tensor): The tensor to return everytime this module is called.
         """
         super().__init__()
         assert isinstance(tensor, torch.Tensor)
@@ -615,7 +615,7 @@ def _get_offset(self, offset: Union[int, Tuple[int]], n: int) -> List[List[int]]
 
         Args:
 
-            offset (int or list of int or list of list of ints , optional): The offsets
+            offset (int or list of int or list of list of ints, optional): The offsets
                 to use for the shared tensors.
             n (int): The number of tensors needing offset values.
 
@@ -741,7 +741,7 @@ def _interpolate_tensor(
             width (int): The width to resize the tensor to.
 
         Returns:
-            **tensor** (torch.Tensor): A resized tensor.
+            tensor (torch.Tensor): A resized tensor.
         """
 
         if x.size(1) == channels:
@@ -944,7 +944,7 @@ def __init__(
                 nn.Parameter tensor. This parameter is not used if ``parameterization``
                 is an instance.
                 Default: ``1``
-            init (torch.tensor, optional): Optionally specify a tensor to use instead
+            init (torch.Tensor, optional): Optionally specify a tensor to use instead
                 of creating one from random noise. This parameter is not used if
                 ``parameterization`` is an instance. Set to ``None`` for random init.
                 Default: ``None``
@@ -1000,7 +1000,7 @@ def _to_image_tensor(self, x: torch.Tensor) -> torch.Tensor:
 
         Args:
 
-            x (torch.tensor): An input tensor.
+            x (torch.Tensor): An input tensor.
 
         Returns:
             x (ImageTensor): An instance of ``ImageTensor`` with the input tensor.

From 32c4ba5f5a665f7cead42259d919687ab06e587c Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Sat, 16 Jul 2022 10:46:25 -0600
Subject: [PATCH 24/30] Fix parameterization doc type hint formatting

---
 captum/optim/_param/image/images.py | 29 +++++++++++++++--------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index bd973e6d5b..929c8f8a93 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -304,7 +304,7 @@ def _get_fft_funcs(self) -> Tuple[Callable, Callable, Callable]:
         torch.fft update.
 
         Returns:
-            fft functions (tuple of Callable): A list of FFT functions
+            fft functions (Tuple[Callable]): A list of FFT functions
                 to use for irfft, rfft, and fftfreq operations.
         """
 
@@ -468,7 +468,7 @@ def __init__(
                 Default: ``None``
             power (float, optional): The desired power value to use.
                 Default: ``0.1``
-            scale_list (list of float, optional): The desired list of scale values to
+            scale_list (List[float], optional): The desired list of scale values to
                 use in the laplacian pyramid. The height & width dimensions specified
                 in ``size`` or used in the ``init`` tensor should be divisable by every
                 scale value in the scale list with no remainder left over. The default
@@ -585,11 +585,11 @@ def __init__(
         """
         Args:
 
-            shapes (list of int or list of list of ints): The shapes of the shared
+            shapes (List[int] or List[List[int]]): The shapes of the shared
                 tensors to use for creating the nn.Parameter tensors.
             parameterization (ImageParameterization): An image parameterization
                 instance.
-            offset (int or list of int or list of list of ints , optional): The offsets
+            offset (int or List[int] or List[List[int]] , optional): The offsets
                 to use for the shared tensors.
                 Default: ``None``
         """
@@ -615,12 +615,12 @@ def _get_offset(self, offset: Union[int, Tuple[int]], n: int) -> List[List[int]]
 
         Args:
 
-            offset (int or list of int or list of list of ints, optional): The offsets
+            offset (int or List[int] or List[List[int]], optional): The offsets
                 to use for the shared tensors.
             n (int): The number of tensors needing offset values.
 
         Returns:
-            offset (list of list of int): A list of offset values.
+            offset (List[List[int]]): A list of offset values.
         """
         if type(offset) is tuple or type(offset) is list:
             if type(offset[0]) is tuple or type(offset[0]) is list:
@@ -641,10 +641,10 @@ def _apply_offset(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:
 
         Args:
 
-            x_list (list of torch.Tensor): list of tensors to offset.
+            x_list (List[torch.Tensor]): list of tensors to offset.
 
         Returns:
-            A (list of torch.Tensor): list of offset tensors.
+            A (List[torch.Tensor]): list of offset tensors.
         """
 
         A: List[torch.Tensor] = []
@@ -679,7 +679,7 @@ def _interpolate_bilinear(
         Args:
 
             x (torch.Tensor): The NCHW tensor to resize.
-            size (tuple of int): The desired output size to resize the input
+            size (Tuple[int]): The desired output size to resize the input
                 to, with a format of: [height, width].
 
         Returns:
@@ -708,7 +708,7 @@ def _interpolate_trilinear(
         Args:
 
             x (torch.Tensor): The NCHW tensor to resize.
-            size (tuple of int): The desired output size to resize the input
+            size (Tuple[int]): The desired output size to resize the input
                 to, with a format of: [channels, height, width].
 
         Returns:
@@ -819,11 +819,12 @@ def __init__(
         """
         Args:
 
-            parameterizations (list of ImageParameterization and torch.Tensor): A list
-                 of image parameterizations to stack across their batch dimensions.
+            parameterizations (List[Union[ImageParameterization, torch.Tensor]]): A
+                list of image parameterizations and tensors to concatenate across a
+                specified dimension.
             dim (int, optional): Optionally specify the dim to concatinate
-                 parameterization outputs on. Default is set to the batch dimension.
-                 Default: ``0``
+                parameterization outputs on. Default is set to the batch dimension.
+                Default: ``0``
             output_device (torch.device, optional): If the parameterizations are on
                 different devices, then their outputs will be moved to the device
                 specified by this variable. Default is set to ``None`` with the

From 6259b13c3dfabaaff43fd2a7c96652f84dbc8f25 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Sat, 16 Jul 2022 13:23:17 -0600
Subject: [PATCH 25/30] Improve doc types

---
 captum/optim/_param/image/images.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index 929c8f8a93..a1762601eb 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -304,8 +304,8 @@ def _get_fft_funcs(self) -> Tuple[Callable, Callable, Callable]:
         torch.fft update.
 
         Returns:
-            fft functions (Tuple[Callable]): A list of FFT functions
-                to use for irfft, rfft, and fftfreq operations.
+            fft functions (Tuple[Callable, Callable, Callable]): A list of FFT
+                functions to use for irfft, rfft, and fftfreq operations.
         """
 
         if version.parse(TORCH_VERSION) > version.parse("1.7.0"):
@@ -679,7 +679,7 @@ def _interpolate_bilinear(
         Args:
 
             x (torch.Tensor): The NCHW tensor to resize.
-            size (Tuple[int]): The desired output size to resize the input
+            size (Tuple[int, int]): The desired output size to resize the input
                 to, with a format of: [height, width].
 
         Returns:
@@ -708,7 +708,7 @@ def _interpolate_trilinear(
         Args:
 
             x (torch.Tensor): The NCHW tensor to resize.
-            size (Tuple[int]): The desired output size to resize the input
+            size (Tuple[int, int, int]): The desired output size to resize the input
                 to, with a format of: [channels, height, width].
 
         Returns:

From 5335a4ed43d8bf04335960151894f6d1a8913635 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Mon, 18 Jul 2022 15:30:16 -0600
Subject: [PATCH 26/30] Improve parameterization docs

---
 captum/optim/_param/image/images.py | 108 +++++++++++++++-------------
 1 file changed, 57 insertions(+), 51 deletions(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index a1762601eb..18779ebd97 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -225,14 +225,14 @@ def __init__(
         """
         Args:
 
-            size (Tuple[int, int]): The height & width dimensions to use for the
-                parameterized output image tensor.
+            size (tuple of int): The height & width dimensions to use for the
+                parameterized output image tensor, in the format of: (height, width).
             channels (int, optional): The number of channels to use for each image.
                 Default: ``3``
             batch (int, optional): The number of images to stack along the batch
                 dimension.
                 Default: ``1``
-            init (torch.Tensor, optional): Optionally specify a tensor to
+            init (torch.Tensor, optional): Optionally specify a CHW or NCHW tensor to
                 use instead of creating one.
                 Default: ``None``
         """
@@ -304,8 +304,8 @@ def _get_fft_funcs(self) -> Tuple[Callable, Callable, Callable]:
         torch.fft update.
 
         Returns:
-            fft functions (Tuple[Callable, Callable, Callable]): A list of FFT
-                functions to use for irfft, rfft, and fftfreq operations.
+            fft_functions (tuple of callable): A list of FFT functions to use for
+                irfft, rfft, and fftfreq operations.
         """
 
         if version.parse(TORCH_VERSION) > version.parse("1.7.0"):
@@ -388,14 +388,14 @@ def __init__(
         """
         Args:
 
-            size (Tuple[int, int]): The height & width dimensions to use for the
-                parameterized output image tensor.
+            size (tuple of int): The height & width dimensions to use for the
+                parameterized output image tensor, in the format of: (height, width).
             channels (int, optional): The number of channels to use for each image.
                 Default: ``3``
             batch (int, optional): The number of images to stack along the batch
                 dimension.
                 Default: ``1``
-            init (torch.Tensor, optional): Optionally specify a tensor to
+            init (torch.Tensor, optional): Optionally specify a CHW or NCHW tensor to
                 use instead of creating one.
                 Default: ``None``
         """
@@ -445,7 +445,7 @@ class LaplacianImage(ImageParameterization):
 
     def __init__(
         self,
-        size: Tuple[int, int] = (224, 225),
+        size: Tuple[int, int] = (224, 224),
         channels: int = 3,
         batch: int = 1,
         init: Optional[torch.Tensor] = None,
@@ -455,15 +455,14 @@ def __init__(
         """
         Args:
 
-            size (Tuple[int, int], optional): The height & width dimensions to use for
-                the parameterized output image tensor.
-                Default: ``(224, 224)``
+            size (tuple of int): The height & width dimensions to use for the
+                parameterized output image tensor, in the format of: (height, width).
             channels (int, optional): The number of channels to use for each image.
                 Default: ``3``
             batch (int, optional): The number of images to stack along the batch
                 dimension.
                 Default: ``1``
-            init (torch.Tensor, optional): Optionally specify a tensor to
+            init (torch.Tensor, optional): Optionally specify a CHW or NCHW tensor to
                 use instead of creating one.
                 Default: ``None``
             power (float, optional): The desired power value to use.
@@ -585,11 +584,11 @@ def __init__(
         """
         Args:
 
-            shapes (List[int] or List[List[int]]): The shapes of the shared
+            shapes (list of int or list of list of int): The shapes of the shared
                 tensors to use for creating the nn.Parameter tensors.
             parameterization (ImageParameterization): An image parameterization
                 instance.
-            offset (int or List[int] or List[List[int]] , optional): The offsets
+            offset (int or list of int or list of list of int, optional): The offsets
                 to use for the shared tensors.
                 Default: ``None``
         """
@@ -615,7 +614,7 @@ def _get_offset(self, offset: Union[int, Tuple[int]], n: int) -> List[List[int]]
 
         Args:
 
-            offset (int or List[int] or List[List[int]], optional): The offsets
+            offset (int or list of int or list of list of int, optional): The offsets
                 to use for the shared tensors.
             n (int): The number of tensors needing offset values.
 
@@ -641,10 +640,10 @@ def _apply_offset(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:
 
         Args:
 
-            x_list (List[torch.Tensor]): list of tensors to offset.
+            x_list (list of torch.Tensor): list of tensors to offset.
 
         Returns:
-            A (List[torch.Tensor]): list of offset tensors.
+            A (list of torch.Tensor): list of offset tensors.
         """
 
         A: List[torch.Tensor] = []
@@ -679,8 +678,8 @@ def _interpolate_bilinear(
         Args:
 
             x (torch.Tensor): The NCHW tensor to resize.
-            size (Tuple[int, int]): The desired output size to resize the input
-                to, with a format of: [height, width].
+            size (tuple of int): The desired output size to resize the input to, with
+                a format of: [height, width].
 
         Returns:
             x (torch.Tensor): A resized NCHW tensor.
@@ -708,8 +707,8 @@ def _interpolate_trilinear(
         Args:
 
             x (torch.Tensor): The NCHW tensor to resize.
-            size (Tuple[int, int, int]): The desired output size to resize the input
-                to, with a format of: [channels, height, width].
+            size (tuple of int): The desired output size to resize the input to, with
+                a format of: [channels, height, width].
 
         Returns:
             x (torch.Tensor): A resized NCHW tensor.
@@ -819,8 +818,8 @@ def __init__(
         """
         Args:
 
-            parameterizations (List[Union[ImageParameterization, torch.Tensor]]): A
-                list of image parameterizations and tensors to concatenate across a
+            parameterizations (list of ImageParameterization and torch.Tensor): A list
+                of image parameterizations and tensors to concatenate across a
                 specified dimension.
             dim (int, optional): Optionally specify the dim to concatinate
                 parameterization outputs on. Default is set to the batch dimension.
@@ -912,11 +911,6 @@ class NaturalImage(ImageParameterization):
         True
         >>> print(image_tensor.shape)
         torch.Size([1, 3, 224, 224])
-
-    :ivar parameterization: initial value (ImageParameterization): The given image
-        parameterization instance given when initializing ``NaturalImage``.
-    :ivar decorrelation_module: initial value (nn.Module): The given decorrelation
-        module instance given when initializing ``NaturalImage``.
     """
 
     def __init__(
@@ -926,48 +920,60 @@ def __init__(
         batch: int = 1,
         init: Optional[torch.Tensor] = None,
         parameterization: ImageParameterization = FFTImage,
-        squash_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+        squash_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = torch.sigmoid,
         decorrelation_module: Optional[nn.Module] = ToRGB(transform="klt"),
         decorrelate_init: bool = True,
     ) -> None:
         """
         Args:
 
-            size (Tuple[int, int], optional): The height and width to use for the
-                nn.Parameter image tensor. This parameter is not used if
-                parameterization is an instance.
-                Default: ``(224, 224)``
+            size (tuple of int, optional): The height and width to use for the
+                nn.Parameter image tensor, in the format of: (height, width).
+                This parameter is not used if the given ``parameterization`` is an
+                instance.
+                Default: ``(224, 224)`
             channels (int, optional): The number of channels to use when creating the
-                nn.Parameter tensor. This parameter is not used if parameterization is
-                an instance.
+                nn.Parameter tensor. This parameter is not used if the given
+                ``parameterization`` is an instance.
                 Default: ``3``
             batch (int, optional): The number of channels to use when creating the
-                nn.Parameter tensor. This parameter is not used if ``parameterization``
-                is an instance.
+                nn.Parameter tensor. This parameter is not used if the given
+                ``parameterization`` is an instance.
                 Default: ``1``
             init (torch.Tensor, optional): Optionally specify a tensor to use instead
-                of creating one from random noise. This parameter is not used if
-                ``parameterization`` is an instance. Set to ``None`` for random init.
+                of creating one from random noise. This parameter is not used if the
+                given ``parameterization`` is an instance. Set to ``None`` for random
+                init.
                 Default: ``None``
             parameterization (ImageParameterization, optional): An image
                 parameterization class, or instance of an image parameterization class.
-                Default: FFTImage
-            squash_func (Callable[[torch.Tensor], torch.Tensor]], optional): The squash
-                function to use after color recorrelation. A function, lambda function,
-                or callable class instance.
-                Default: ``None``
+                Default: :class:`.FFTImage`
+            squash_func (callable, optional): The squash function to use after color
+                recorrelation. A function, lambda function, or callable class instance.
+                Any provided squash function should take a single input tensor and
+                return a single output tensor. If set to ``None``, then
+                :class:`torch.nn.Identity` will be used to make it a non op.
+                Default: :func:`torch.sigmoid`
             decorrelation_module (nn.Module, optional): A module instance that
                 recorrelates the colors of an input image. Custom modules can make use
                 of the ``decorrelate_init`` parameter by having a second ``inverse``
                 parameter in their forward functions that performs the inverse
-                operation when it is set to ``True`` (see
-                :class:`captum.optim.transforms.ToRGB` for an example).
-                Set to ``None`` for no recorrelation.
-                Default: ``ToRGB``
+                operation when it is set to ``True`` (see :class:`.ToRGB` for an
+                example). Set to ``None`` for no recorrelation.
+                Default: :class:`.ToRGB`
             decorrelate_init (bool, optional): Whether or not to apply color
                 decorrelation to the init tensor input. This parameter is not used if
-                ``parameterization`` is an instance or if init is ``None``.
+                the given ``parameterization`` is an instance or if init is ``None``.
                 Default: ``True``
+
+        Attributes:
+
+            parameterization (ImageParameterization): The given image parameterization
+                instance given when initializing ``NaturalImage``.
+                Default: :class:`.FFTImage`
+            decorrelation_module (torch.nn.Module): The given decorrelation module
+                instance given when initializing ``NaturalImage``.
+                Default: :class:`.ToRGB`
         """
         super().__init__()
         if not isinstance(parameterization, ImageParameterization):
@@ -987,7 +993,7 @@ def __init__(
                 )
                 init = self.decorrelate(init, inverse=True).rename(None)
 
-        self.squash_func = torch.sigmoid if squash_func is None else squash_func
+        self.squash_func = squash_func or torch.nn.Identity()
         if not isinstance(parameterization, ImageParameterization):
             parameterization = parameterization(
                 size=size, channels=channels, batch=batch, init=init

From 668aff1dcbd1550d7ee64e7a7cbeef411e57a690 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Tue, 19 Jul 2022 13:57:49 -0600
Subject: [PATCH 27/30] Fix NaturalImage docs issue

---
 captum/optim/_param/image/images.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index 18779ebd97..b86ace16d0 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -931,7 +931,7 @@ def __init__(
                 nn.Parameter image tensor, in the format of: (height, width).
                 This parameter is not used if the given ``parameterization`` is an
                 instance.
-                Default: ``(224, 224)`
+                Default: ``(224, 224)``
             channels (int, optional): The number of channels to use when creating the
                 nn.Parameter tensor. This parameter is not used if the given
                 ``parameterization`` is an instance.

From 4bab7d75fcd84e6752e2969b151d8f2b06372970 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Wed, 20 Jul 2022 14:06:17 -0600
Subject: [PATCH 28/30] Fix doc type hint

---
 captum/optim/_param/image/images.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index b86ace16d0..54abf49614 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -467,7 +467,7 @@ def __init__(
                 Default: ``None``
             power (float, optional): The desired power value to use.
                 Default: ``0.1``
-            scale_list (List[float], optional): The desired list of scale values to
+            scale_list (list of float, optional): The desired list of scale values to
                 use in the laplacian pyramid. The height & width dimensions specified
                 in ``size`` or used in the ``init`` tensor should be divisable by every
                 scale value in the scale list with no remainder left over. The default

From e727712cefa12fbae9b5f9679620a11fc4783af8 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Fri, 22 Jul 2022 12:16:16 -0600
Subject: [PATCH 29/30] Simplify reducer tests

---
 tests/optim/utils/test_reducer.py | 19 ++++---------------
 1 file changed, 4 insertions(+), 15 deletions(-)

diff --git a/tests/optim/utils/test_reducer.py b/tests/optim/utils/test_reducer.py
index 4c97ef2440..a9fb9cc93c 100644
--- a/tests/optim/utils/test_reducer.py
+++ b/tests/optim/utils/test_reducer.py
@@ -52,9 +52,7 @@ def test_channelreducer_pytorch_dim_three(self) -> None:
         test_input = torch.randn(32, 224, 224).abs()
         c_reducer = reducer.ChannelReducer(n_components=3, max_iter=100)
         test_output = c_reducer.fit_transform(test_input)
-        self.assertEqual(test_output.size(0), 3)
-        self.assertEqual(test_output.size(1), 224)
-        self.assertEqual(test_output.size(2), 224)
+        self.assertEqual(list(test_output.shape), [3, 224, 224])
 
     def test_channelreducer_pytorch_pca(self) -> None:
         try:
@@ -70,10 +68,7 @@ def test_channelreducer_pytorch_pca(self) -> None:
         c_reducer = reducer.ChannelReducer(n_components=3, reduction_alg="PCA")
 
         test_output = c_reducer.fit_transform(test_input)
-        self.assertEqual(test_output.size(0), 1)
-        self.assertEqual(test_output.size(1), 3)
-        self.assertEqual(test_output.size(2), 224)
-        self.assertEqual(test_output.size(3), 224)
+        self.assertEqual(list(test_output.shape), [1, 3, 224, 224])
 
     def test_channelreducer_pytorch_custom_alg(self) -> None:
         test_input = torch.randn(1, 32, 224, 224).abs()
@@ -82,10 +77,7 @@ def test_channelreducer_pytorch_custom_alg(self) -> None:
             n_components=3, reduction_alg=reduction_alg, max_iter=100
         )
         test_output = c_reducer.fit_transform(test_input)
-        self.assertEqual(test_output.size(0), 1)
-        self.assertEqual(test_output.size(1), 3)
-        self.assertEqual(test_output.size(2), 224)
-        self.assertEqual(test_output.size(3), 224)
+        self.assertEqual(list(test_output.shape), [1, 3, 224, 224])
 
     def test_channelreducer_pytorch_custom_alg_components(self) -> None:
         reduction_alg = FakeReductionAlgorithm
@@ -149,10 +141,7 @@ def test_channelreducer_noreshape_pytorch(self) -> None:
         test_input = torch.randn(1, 224, 224, 32).abs()
         c_reducer = reducer.ChannelReducer(n_components=3, max_iter=100)
         test_output = c_reducer.fit_transform(test_input, swap_2nd_and_last_dims=False)
-        self.assertEqual(test_output.size(0), 1)
-        self.assertEqual(test_output.size(1), 224)
-        self.assertEqual(test_output.size(2), 224)
-        self.assertEqual(test_output.size(3), 3)
+        self.assertEqual(list(test_output.shape), [1, 224, 224, 3])
 
     def test_channelreducer_error(self) -> None:
         if not torch.cuda.is_available():

From 33f9f66ad50177d80fe402300c5ee92195d6bdd8 Mon Sep 17 00:00:00 2001
From: ProGamerGov <ProGamerGov@users.noreply.github.com>
Date: Wed, 27 Jul 2022 09:26:37 -0600
Subject: [PATCH 30/30] Fix spelling

---
 captum/optim/_param/image/images.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/captum/optim/_param/image/images.py b/captum/optim/_param/image/images.py
index 54abf49614..f3a45346ca 100644
--- a/captum/optim/_param/image/images.py
+++ b/captum/optim/_param/image/images.py
@@ -469,9 +469,10 @@ def __init__(
                 Default: ``0.1``
             scale_list (list of float, optional): The desired list of scale values to
                 use in the laplacian pyramid. The height & width dimensions specified
-                in ``size`` or used in the ``init`` tensor should be divisable by every
+                in ``size`` or used in the ``init`` tensor should be divisible by every
                 scale value in the scale list with no remainder left over. The default
-                scale_list values are set to work with a ``size`` of ``(224, 224)``.
+                ``scale_list`` values are set to work with a ``size`` of
+                ``(224, 224)``.
                 Default: ``[1.0, 2.0, 4.0, 8.0, 16.0, 32.0]``
         """
         super().__init__()
@@ -484,7 +485,7 @@ def __init__(
         for scale in scale_list:
             assert size[0] % scale == 0 and size[1] % scale == 0, (
                 "The chosen image height & width dimensions"
-                + " must be divisable by all scale values "
+                + " must be divisible by all scale values "
                 + " with no remainder left over."
             )
 
@@ -535,7 +536,8 @@ def __init__(self, tensor: torch.Tensor = None) -> None:
         """
         Args:
 
-            tensor (torch.Tensor): The tensor to return everytime this module is called.
+            tensor (torch.Tensor): The tensor to return every time this module is
+                called.
         """
         super().__init__()
         assert isinstance(tensor, torch.Tensor)
@@ -821,7 +823,7 @@ def __init__(
             parameterizations (list of ImageParameterization and torch.Tensor): A list
                 of image parameterizations and tensors to concatenate across a
                 specified dimension.
-            dim (int, optional): Optionally specify the dim to concatinate
+            dim (int, optional): Optionally specify the dim to concatenate
                 parameterization outputs on. Default is set to the batch dimension.
                 Default: ``0``
             output_device (torch.device, optional): If the parameterizations are on