Add sparse concat operator support

keras/backend/mxnet_backend.py

@@ -17,6 +17,8 @@
 _REENTRY = False
 NAME_SCOPE_STACK = []
 
+py_all = all
+
 
 class name_scope(object):
     def __init__(self, name):
@@ -2030,6 +2032,10 @@ def concatenate(tensors, axis=-1):
 
     # Returns
         A tensor.
+
+    Note:
+    - MXNet supports sparse concat only for dim=0
+    - https://mxnet.apache.org/api/python/symbol/sparse.html#mxnet.symbol.sparse.concat
     """
     if axis < 0:
         rank = ndim(tensors[0])
@@ -2038,8 +2044,12 @@ def concatenate(tensors, axis=-1):
         else:
             axis = 0
 
-    tensors = [t.symbol for t in tensors]
-    return KerasSymbol(mx.sym.concat(*tensors, dim=axis))
+    symbols = [t.symbol for t in tensors]
+
+    if axis == 0 and py_all([is_sparse(t) for t in tensors]):
+        return KerasSymbol(mx.sym.sparse.concat(*symbols, dim=axis))
+
+    return KerasSymbol(mx.sym.concat(*symbols, dim=axis))
 
 
 @keras_mxnet_symbol

tests/keras/backend/mxnet_sparse_test.py

@@ -104,6 +104,64 @@ def test_sparse_dot(self):
         assert k_s.shape == k_d.shape
         assert_allclose(k_s, k_d, atol=1e-05)
 
+    def test_sparse_concat(self):
+        x_sparse_1 = self.generate_test_sparse_matrix()
+        x_sparse_2 = self.generate_test_sparse_matrix()
+
+        assert K.is_sparse(K.variable(x_sparse_1))
+        assert K.is_sparse(K.variable(x_sparse_2))
+        x_dense_1 = x_sparse_1.toarray()
+        x_dense_2 = x_sparse_2.toarray()
+
+        k_s = K.concatenate(tensors=[K.variable(x_sparse_1), K.variable(x_sparse_2)], axis=0)
+        assert K.is_sparse(k_s)
+
+        k_s_d = K.eval(k_s)
+
+        # mx.sym.sparse.concat only supported for axis=0
+        k_d = K.eval(K.concatenate(tensors=[K.variable(x_dense_1), K.variable(x_dense_2)], axis=0))
+
+        assert k_s_d.shape == k_d.shape
+        assert_allclose(k_s_d, k_d, atol=1e-05)
+
+    def test_sparse_concat_partial_dense(self):
+        x_sparse_1 = self.generate_test_sparse_matrix()
+        x_sparse_2 = self.generate_test_sparse_matrix()
+
+        assert K.is_sparse(K.variable(x_sparse_1))
+        x_dense_1 = x_sparse_1.toarray()
+        x_dense_2 = x_sparse_2.toarray()
+
+        k_s = K.concatenate(tensors=[K.variable(x_sparse_1), K.variable(x_dense_2)], axis=0)
+        assert not(K.is_sparse(k_s))
+
+        k_s_d = K.eval(k_s)
+
+        # mx.sym.sparse.concat only supported for axis=0
+        k_d = K.eval(K.concatenate(tensors=[K.variable(x_dense_1), K.variable(x_dense_2)], axis=0))
+
+        assert k_s_d.shape == k_d.shape
+        assert_allclose(k_s_d, k_d, atol=1e-05)
+
+    def test_sparse_concat_axis_non_zero(self):
+        x_sparse_1 = self.generate_test_sparse_matrix()
+        x_sparse_2 = self.generate_test_sparse_matrix()
+
+        assert K.is_sparse(K.variable(x_sparse_1))
+        x_dense_1 = x_sparse_1.toarray()
+        x_dense_2 = x_sparse_2.toarray()
+
+        k_s = K.concatenate(tensors=[K.variable(x_sparse_1), K.variable(x_dense_2)])
+        assert not (K.is_sparse(k_s))
+
+        k_s_d = K.eval(k_s)
+
+        # mx.sym.sparse.concat only supported for axis=0
+        k_d = K.eval(K.concatenate(tensors=[K.variable(x_dense_1), K.variable(x_dense_2)]))
+
+        assert k_s_d.shape == k_d.shape
+        assert_allclose(k_s_d, k_d, atol=1e-05)
+
 
 if __name__ == '__main__':
     pytest.main([__file__])