optim-wip: Fix objectives.py, images.py, RedirectedReLU etc.

captum/optim/__init__.py

@@ -1,11 +1,31 @@
 """optim submodule."""
 
-from captum.optim._core import loss  # noqa: F401
-from captum.optim._core import objectives  # noqa: F401
-from captum.optim._core.objectives import InputOptimization  # noqa: F401
-from captum.optim._param.image import images  # noqa: F401
-from captum.optim._param.image import transform  # noqa: F401
+from captum.optim import models
+from captum.optim._core import loss, optimization  # noqa: F401
+from captum.optim._core.optimization import InputOptimization  # noqa: F401
+from captum.optim._param.image import images, transforms  # noqa: F401
 from captum.optim._param.image.images import ImageTensor  # noqa: F401
-from captum.optim._utils import circuits, models, reducer  # noqa: F401
-from captum.optim._utils.image.common import nchannels_to_rgb  # noqa: F401
-from captum.optim._utils.image.common import weights_to_heatmap_2d  # noqa: F401
+from captum.optim._utils import circuits, reducer  # noqa: F401
+from captum.optim._utils.image.common import (  # noqa: F401
+    nchannels_to_rgb,
+    save_tensor_as_image,
+    show,
+    weights_to_heatmap_2d,
+)
+from captum.optim._utils.reducer import ChannelReducer, posneg  # noqa: F401
+
+__all__ = [
+    "InputOptimization",
+    "ImageTensor",
+    "loss",
+    "optimization",
+    "images",
+    "transforms",
+    "circuits",
+    "models",
+    "reducer",
+    "nchannels_to_rgb",
+    "save_tensor_as_image",
+    "show",
+    "weights_to_heatmap_2d",
+]