Add sparse support to ops.ones_like and ops.zeros_like. (#21181)

`ops.zeros_like` is in particular useful for creating a mask of the populated values in the sparse tensor.

Author: hertschuh

keras/src/backend/jax/numpy.py

@@ -895,10 +895,12 @@ def not_equal(x1, x2):
     return jnp.not_equal(x1, x2)
 
 
+@sparse.elementwise_unary(linear=False)
 def ones_like(x, dtype=None):
     return jnp.ones_like(x, dtype=dtype)
 
 
+@sparse.elementwise_unary(linear=True)
 def zeros_like(x, dtype=None):
     return jnp.zeros_like(x, dtype=dtype)
 

keras/src/backend/tensorflow/numpy.py

@@ -1798,10 +1798,12 @@ def not_equal(x1, x2):
     return tf.not_equal(x1, x2)
 
 
+@sparse.elementwise_unary
 def ones_like(x, dtype=None):
     return tf.ones_like(x, dtype=dtype)
 
 
+@sparse.elementwise_unary
 def zeros_like(x, dtype=None):
     return tf.zeros_like(x, dtype=dtype)
 

keras/src/ops/numpy.py

@@ -4391,7 +4391,8 @@ def call(self, x, dtype=None):
     def compute_output_spec(self, x, dtype=None):
         if dtype is None:
             dtype = x.dtype
-        return KerasTensor(x.shape, dtype=dtype)
+        sparse = getattr(x, "sparse", False)
+        return KerasTensor(x.shape, dtype=dtype, sparse=sparse)
 
 
 @keras_export(["keras.ops.ones_like", "keras.ops.numpy.ones_like"])
@@ -4417,7 +4418,8 @@ def call(self, x, dtype=None):
     def compute_output_spec(self, x, dtype=None):
         if dtype is None:
             dtype = x.dtype
-        return KerasTensor(x.shape, dtype=dtype)
+        sparse = getattr(x, "sparse", False)
+        return KerasTensor(x.shape, dtype=dtype, sparse=sparse)
 
 
 @keras_export(

keras/src/ops/numpy_test.py

@@ -5097,6 +5097,7 @@ class SparseTest(testing.TestCase):
         "imag",
         "log1p",
         "negative",
+        "ones_like",
         "real",
         "round",
         "sign",
@@ -5106,6 +5107,7 @@ class SparseTest(testing.TestCase):
         "square",
         "tan",
         "tanh",
+        "zeros_like",
     ]
     ELEMENTWISE_UNARY_OPS_TESTS = [
         {
@@ -5287,10 +5289,11 @@ def test_elementwise_unary_sparse_correctness(
             x = np.array([[1, 0.5, -0.7], [0.9, 0.2, -1]])
         x = create_sparse_tensor(x)
         x_np = backend.convert_to_numpy(x)
+        expected = np_op(x_np) * backend.convert_to_numpy(knp.ones_like(x))
 
-        self.assertAllClose(op_function(x), np_op(x_np))
+        self.assertAllClose(op_function(x), expected)
         self.assertSameSparseness(op_function(x), x)
-        self.assertAllClose(op_class()(x), np_op(x_np))
+        self.assertAllClose(op_class()(x), expected)
         self.assertSameSparseness(op_class()(x), x)
 
     @parameterized.named_parameters(ELEMENTWISE_UNARY_OPS_TESTS)
@@ -5303,10 +5306,11 @@ def test_elementwise_unary_indexed_slices_correctness(
             x = np.array([[1, 0.5, -0.7], [0.9, 0.2, -1]])
         x = create_indexed_slices(x)
         x_np = backend.convert_to_numpy(x)
+        expected = np_op(x_np) * backend.convert_to_numpy(knp.ones_like(x))
 
-        self.assertAllClose(op_function(x), np_op(x_np))
+        self.assertAllClose(op_function(x), expected)
         self.assertSameSparseness(op_function(x), x)
-        self.assertAllClose(op_class()(x), np_op(x_np))
+        self.assertAllClose(op_class()(x), expected)
         self.assertSameSparseness(op_class()(x), x)
 
     @parameterized.named_parameters(OTHER_UNARY_OPS_TESTS)