[pycharm] PY-38294 Add Sparse tensors support

GitOrigin-RevId: 0b9b046ae1da37b5fd34be800fbdd6f241f7514d

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/outputs/statistic/NotebookOutputKeyType.kt

@@ -13,6 +13,7 @@ enum class NotebookOutputKeyType {
   NUMPY_ARRAY,
   EAGER_TENSOR,
   RESOURCE_VARIABLE,
+  SPARSE_TENSOR,
   TORCH_TENSOR,
   PANDAS_DATA_FRAME,
   PANDAS_SERIES,

python/helpers/pydev/_pydevd_bundle/pydevd_tables.py

@@ -75,6 +75,7 @@ def __get_table_provider(output):
     elif type_qualified_name in ['numpy.ndarray',
                                  'tensorflow.python.framework.ops.EagerTensor',
                                  'tensorflow.python.ops.resource_variable_ops.ResourceVariable',
+                                 'tensorflow.python.framework.sparse_tensor.SparseTensor',
                                  'torch.Tensor',
                                  'builtins.dict']:
         import _pydevd_bundle.tables.pydevd_numpy as table_provider

python/helpers/pydev/_pydevd_bundle/pydevd_thrift.py

@@ -412,6 +412,11 @@ def tensor_to_thrift_struct(tensor, name, roffset, coffset, rows, cols, format):
     return array_to_thrift_struct(tensor.numpy(), name, roffset, coffset, rows, cols, format)
 
 
+def sparse_tensor_to_thrift_struct(tensor, name, roffset, coffset, rows, cols, format):
+    import tensorflow as tf
+    return tensor_to_thrift_struct(tf.sparse.to_dense(tf.sparse.reorder(tensor)), name, roffset, coffset, rows, cols, format)
+
+
 def array_to_meta_thrift_struct(array, name, format):
     type = array.dtype.kind
     slice = name
@@ -609,6 +614,7 @@ TYPE_TO_THRIFT_STRUCT_CONVERTERS = {
     "ndarray": array_to_thrift_struct,
     "EagerTensor": tensor_to_thrift_struct,
     "ResourceVariable": tensor_to_thrift_struct,
+    "SparseTensor": sparse_tensor_to_thrift_struct,
     "Tensor": tensor_to_thrift_struct,
     "DataFrame": dataframe_to_thrift_struct,
     "Series": dataframe_to_thrift_struct,

python/helpers/pydev/_pydevd_bundle/pydevd_vars.py

@@ -586,6 +586,11 @@ def tensor_to_xml(tensor, name, roffset, coffset, rows, cols, format):
     return array_to_xml(tensor.numpy(), name, roffset, coffset, rows, cols, format)
 
 
+def sparse_tensor_to_xml(tensor, name, roffset, coffset, rows, cols, format):
+    import tensorflow as tf
+    return tensor_to_xml(tf.sparse.to_dense(tf.sparse.reorder(tensor)), name, roffset, coffset, rows, cols, format)
+
+
 class ExceedingArrayDimensionsException(Exception):
     pass
 
@@ -811,6 +816,7 @@ TYPE_TO_XML_CONVERTERS = {
     "GeoSeries": dataframe_to_xml,
     "EagerTensor": tensor_to_xml,
     "ResourceVariable": tensor_to_xml,
+    "SparseTensor": sparse_tensor_to_xml,
     "Tensor": tensor_to_xml
 }
 

python/helpers/pydev/_pydevd_bundle/tables/pydevd_numpy.py

@@ -220,7 +220,12 @@ def _create_table(command, start_index=None, end_index=None):
         np_array = command['data']
         sort_keys = command['sort_keys']
     else:
-        np_array = command
+        try:
+            import tensorflow as tf
+            if isinstance(command, tf.SparseTensor):
+                command = tf.sparse.to_dense(tf.sparse.reorder(command))
+        finally:
+            np_array = command
 
     if is_pd:
         sorting_arr = _sort_df(pd.DataFrame(np_array), sort_keys)

python/pydevSrc/src/com/jetbrains/python/debugger/PyDebugValue.java

@@ -37,6 +37,7 @@ public class PyDebugValue extends XNamedValue {
     "ndarray", ARRAY,
     "EagerTensor", ARRAY,
     "ResourceVariable", ARRAY,
+    "SparseTensor", ARRAY,
     "Tensor", ARRAY,
     DATA_FRAME, DATA_FRAME,
     SERIES, SERIES,

python/src/com/jetbrains/python/debugger/array/ArrayViewStrategy.java

@@ -27,6 +27,10 @@ public class ArrayViewStrategy extends DataViewStrategy {
     return new ArrayViewStrategy("ResourceVariable");
   }
 
+  public static @NotNull ArrayViewStrategy createInstanceForSparseTensor() {
+    return new ArrayViewStrategy("SparseTensor");
+  }
+
   public static @NotNull ArrayViewStrategy createInstanceForTensor() {
     return new ArrayViewStrategy("Tensor");
   }

python/src/com/jetbrains/python/debugger/containerview/DataViewStrategy.java

@@ -21,6 +21,7 @@ public abstract class DataViewStrategy {
       ArrayViewStrategy.createInstanceForNumpyArray(),
       ArrayViewStrategy.createInstanceForEagerTensor(),
       ArrayViewStrategy.createInstanceForResourceVariable(),
+      ArrayViewStrategy.createInstanceForSparseTensor(),
       ArrayViewStrategy.createInstanceForTensor(),
       DataFrameViewStrategy.createInstanceForDataFrame(),
       DataFrameViewStrategy.createInstanceForGeoDataFrame(),

python/src/com/jetbrains/python/debugger/containerview/PyViewNumericContainerAction.java

@@ -64,7 +64,11 @@ public class PyViewNumericContainerAction extends XDebuggerTreeActionBase {
     }
 
     String nodeType = debugValue.getType();
-    if ("ndarray".equals(nodeType) || "EagerTensor".equals(nodeType) || "ResourceVariable".equals(nodeType) || "Tensor".equals(nodeType)) {
+    if ("ndarray".equals(nodeType) ||
+        "EagerTensor".equals(nodeType) ||
+        "ResourceVariable".equals(nodeType) ||
+        "SparseTensor".equals(nodeType) ||
+        "Tensor".equals(nodeType)) {
       e.getPresentation().setText(PyBundle.message("debugger.numeric.view.as.array"));
       e.getPresentation().setVisible(true);
     }