add recursive convert

Author: SkafteNicki

tests/test_torchplot.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import string
 from collections import namedtuple
 from inspect import getmembers, isfunction
 
@@ -50,7 +49,10 @@ def test_members(member):
 @pytest.mark.parametrize("test_case", _cases)
 def test_cpu(test_case):
     """ test that it works on cpu """
+    # passed as args
     assert tp.plot(test_case.x, test_case.y, ".")
+    # passed as kwargs
+    assert tp.scatter(x=test_case.x, y=test_case.y)
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires cuda")
@@ -60,5 +62,9 @@ def test_gpu(test_case):
     assert tp.plot(
         test_case.x.cuda() if isinstance(test_case.x, torch.Tensor) else test_case.x,
         test_case.y.cuda() if isinstance(test_case.y, torch.Tensor) else test_case.y,
-        ".",
+    )
+    # passed as kwargs
+    assert tp.scatter(
+        x=test_case.x.cuda() if isinstance(test_case.x, torch.Tensor) else test_case.x,
+        y=test_case.y.cuda() if isinstance(test_case.y, torch.Tensor) else test_case.y,
     )

torchplot/core.py

@@ -1,25 +1,63 @@
 #!/usr/bin/env python
 from inspect import getdoc, getmembers, isfunction
+from typing import Any, Callable, Mapping, Sequence, Union
 
 import matplotlib.pyplot as plt
 import torch
 
 
+# Taken from
+# https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/utilities/apply_func.py
+def apply_to_collection(data: Any, dtype: Union[type, tuple], function: Callable, *args, **kwargs) -> Any:
+    """
+    Recursively applies a function to all elements of a certain dtype.
+    Args:
+        data: the collection to apply the function to
+        dtype: the given function will be applied to all elements of this dtype
+        function: the function to apply
+        *args: positional arguments (will be forwarded to calls of ``function``)
+        **kwargs: keyword arguments (will be forwarded to calls of ``function``)
+    Returns:
+        the resulting collection
+    """
+    elem_type = type(data)
+
+    # Breaking condition
+    if isinstance(data, dtype):
+        return function(data, *args, **kwargs)
+
+    # Recursively apply to collection items
+    if isinstance(data, Mapping):
+        return elem_type({k: apply_to_collection(v, dtype, function, *args, **kwargs) for k, v in data.items()})
+
+    if isinstance(data, tuple) and hasattr(data, "_fields"):  # named tuple
+        return elem_type(*(apply_to_collection(d, dtype, function, *args, **kwargs) for d in data))
+
+    if isinstance(data, Sequence) and not isinstance(data, str):
+        return elem_type([apply_to_collection(d, dtype, function, *args, **kwargs) for d in data])
+
+    # data is neither of dtype, nor a collection
+    return data
+
+
 # Function to convert a list of arguments containing torch tensors, into
 # a corresponding list of arguments containing numpy arrays
 def _torch2np(*args, **kwargs):
+    """
+    Convert a list of arguments containing torch tensors into a list of
+    arguments containing numpy arrays
+    """
+
     def convert(arg):
-        return arg.detach().cpu().numpy() if isinstance(arg, torch.Tensor) else arg
+        return arg.detach().cpu().numpy()
 
     # first unnamed arguments
-    outargs = [convert(arg) for arg in args]
+    outargs = apply_to_collection(args, torch.Tensor, convert)
 
     # then keyword arguments
-    outkwargs = dict()
-    for key, value in kwargs.items():
-        outkwargs[key] = convert(value)
+    outkwargs = apply_to_collection(kwargs, torch.Tensor, convert)
 
-    return outargs, kwargs
+    return outargs, outkwargs
 
 
 # Iterate over all members of 'plt' in order to duplicate them