Add torch array support for the variable explorer

spyder/plugins/variableexplorer/widgets/arrayeditor.py

@@ -46,6 +46,11 @@
 from spyder.utils.qthelpers import keybinding, safe_disconnect
 from spyder.utils.stylesheet import AppStyle, MAC
 
+# try to import torch.
+try:    
+    import torch
+except ImportError:
+    torch = None
 
 # =============================================================================
 # ---- Constants
@@ -904,7 +909,7 @@ def do_nothing():
 
         # Set title
         if title:
-            title = str(title) + " - " + _("NumPy object array")
+            title = str(title) + " - " + _("Object array")
         else:
             title = _("Array editor")
         if readonly:
@@ -922,13 +927,23 @@ def set_data_and_check(self, data, readonly=False):
         Setup ArrayEditor:
         return False if data is not supported, True otherwise
         """
+
         if not isinstance(data, (np.ndarray, np.ma.MaskedArray)):
-            return False
-
-        self.data = data
-        readonly = readonly or not self.data.flags.writeable
+            if torch is not None and isinstance(data, torch.Tensor):
+                pass
+            else:
+                return False
         is_masked_array = isinstance(data, np.ma.MaskedArray)
+        is_torch_tensor = isinstance(data, torch.Tensor)
 
+        if is_torch_tensor:
+            readonly = True
+            data = data.detach().cpu().numpy() #data is converted to numpy array for display, but the original tensor is not modified when accepting changes, so we can keep it read-only
+        else:
+            readonly = readonly or not data.flags.writeable
+
+        self.data = data
+
         # Reset data for 3d arrays
         self.dim_indexes = [{}, {}, {}]
         self.last_dim = 0

spyder/plugins/variableexplorer/widgets/collectionsdelegate.py

@@ -57,6 +57,11 @@
 from spyder.plugins.variableexplorer.widgets.texteditor import TextEditor
 from spyder.utils.icon_manager import ima
 
+# try to import torch.
+try:    
+    import torch
+except ImportError:
+    torch = None
 
 LARGE_COLLECTION = 1e5
 LARGE_ARRAY = 5e6
@@ -374,6 +379,25 @@ def createEditor(self, parent, option, index, object_explorer=False):
                 # editor.returnPressed.connect(self.commitAndCloseEditor)
                 self.sig_editor_shown.emit()
                 return editor
+
+        # ArrayEditor for a torch tensors
+        elif torch is not None and isinstance(value, torch.Tensor) and not object_explorer:
+
+            #             # We need to leave this import here for tests to pass.
+            from .arrayeditor import ArrayEditor
+            editor = ArrayEditor(
+                parent=parent,
+                data_function=self.make_data_function(index)
+            )
+
+            # set torch tensors as read-only for now, since we convert them to numpy arrays.
+            if not editor.setup_and_check(value, title=key, readonly=True):
+                self.sig_editor_shown.emit()
+                return
+            self.create_dialog(editor, dict(model=index.model(), editor=editor,
+                                            key=key, readonly=True))
+            return None
+
         # ObjectExplorer for an arbitrary Python object
         else:
             # Don't show the object explorer for short bytes because it's not

spyder/plugins/variableexplorer/widgets/tests/test_arrayeditor.py

@@ -233,6 +233,15 @@ def test_arrayeditor_with_3d_array(qtbot):
     arr[0,0,1]=2
     arr[0,0,2]=3
     assert_array_equal(arr, launch_arrayeditor(arr, "3D array"))
+
+def test_arrayeditor_with_torch_array(qtbot):
+    """Test that the array editor can handle PyTorch tensors."""
+    try:
+        import torch
+    except ImportError:
+        pytest.skip("PyTorch is not installed")
+    arr = torch.rand((5, 5, 5))
+    assert_array_equal(arr, launch_arrayeditor(arr, "torch array"))
 
 
 def test_arrayeditor_with_empty_3d_array(qtbot):