chore: import upstream snapshot with attribution

tensorflow/python/kernel_tests/linalg/sparse/BUILD

@@ -0,0 +1,165 @@
+# Tests of TensorFlow sparse linear algebra kernels using the Python API.
+
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+cuda_py_strict_test(
+    name = "conjugate_gradient_test",
+    size = "medium",
+    srcs = ["conjugate_gradient_test.py"],
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops/linalg",
+        "//tensorflow/python/ops/linalg/sparse:conjugate_gradient",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "csr_sparse_matrix_test",
+    size = "medium",
+    srcs = ["csr_sparse_matrix_test.py"],
+    main = "csr_sparse_matrix_test.py",
+    tags = [
+        "no_cuda_asan",  # TODO(b/190824595)
+    ],
+    deps = [
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops/linalg/sparse:sparse_csr_matrix_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "csr_sparse_matrix_ops_test",
+    size = "medium",
+    timeout = "long",
+    srcs = ["csr_sparse_matrix_ops_test.py"],
+    main = "csr_sparse_matrix_ops_test.py",
+    shard_count = 10,
+    tags = [
+        "no_gpu",  # b/203655060 (cuda 11.5 specific)
+    ],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:linalg_ops",
+        "//tensorflow/python/ops:map_fn",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:sparse_ops",
+        "//tensorflow/python/ops:state_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops/linalg/sparse:sparse_csr_matrix_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:tf_logging",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "csr_sparse_matrix_grad_test",
+    size = "medium",
+    srcs = ["csr_sparse_matrix_grad_test.py"],
+    main = "csr_sparse_matrix_grad_test.py",
+    shard_count = 3,
+    deps = [
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:gradients_impl",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops/linalg/sparse:sparse_csr_matrix_grad",
+        "//tensorflow/python/ops/linalg/sparse:sparse_csr_matrix_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:tf_logging",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "csr_sparse_matrix_dense_mat_mul_grad_test",
+    size = "medium",
+    srcs = ["csr_sparse_matrix_dense_mat_mul_grad_test.py"],
+    main = "csr_sparse_matrix_dense_mat_mul_grad_test.py",
+    shard_count = 50,
+    tags = [
+        "no_cuda_asan",  # TODO(b/190824595)
+    ],
+    deps = [
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:gradient_checker",
+        "//tensorflow/python/ops/linalg/sparse:sparse_csr_matrix_grad",
+        "//tensorflow/python/ops/linalg/sparse:sparse_csr_matrix_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:tf_logging",
+    ],
+)
+
+tf_py_strict_test(
+    name = "csr_sparse_matrix_dense_mat_mul_onednn_grad_test",
+    size = "medium",
+    srcs = ["csr_sparse_matrix_dense_mat_mul_grad_test.py"],
+    env = {
+        "TF_ENABLE_ONEDNN_SPMM": "1",
+    },
+    main = "csr_sparse_matrix_dense_mat_mul_grad_test.py",
+    shard_count = 50,
+    tags = [
+        "no_cuda_asan",  # TODO(b/190824595)
+    ],
+    deps = [
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:gradient_checker",
+        "//tensorflow/python/ops/linalg/sparse:sparse_csr_matrix_grad",
+        "//tensorflow/python/ops/linalg/sparse:sparse_csr_matrix_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:tf_logging",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "csr_sparse_matrix_sparse_mat_mul_grad_test",
+    size = "medium",
+    srcs = ["csr_sparse_matrix_sparse_mat_mul_grad_test.py"],
+    main = "csr_sparse_matrix_sparse_mat_mul_grad_test.py",
+    shard_count = 10,
+    deps = [
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:gradient_checker",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops/linalg/sparse:sparse_csr_matrix_grad",
+        "//tensorflow/python/ops/linalg/sparse:sparse_csr_matrix_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:tf_logging",
+    ],
+)

tensorflow/python/kernel_tests/linalg/sparse/conjugate_gradient_test.py

@@ -0,0 +1,110 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+# ==============================================================================
+"""Tests for preconditioned conjugate gradient."""
+
+import itertools
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.linalg import linalg
+from tensorflow.python.ops.linalg.sparse import conjugate_gradient
+from tensorflow.python.platform import test
+
+
+class ConjugateGradientTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      itertools.product([np.float32, np.float64], [1, 4, 10], [True, False])
+  )
+  def test_conjugate_gradient(self, dtype, size, use_static_shape):
+    shape = [size, size]
+    np.random.seed(1)
+    a_np = (
+        np.random.uniform(low=-1.0, high=1.0, size=np.prod(shape))
+        .reshape(shape)
+        .astype(dtype)
+    )
+    # Make a self-adjoint, positive definite.
+    a_np = np.dot(a_np.T, a_np)
+    # jacobi preconditioner
+    jacobi_np = np.zeros_like(a_np)
+    jacobi_np[range(a_np.shape[0]), range(a_np.shape[1])] = (
+        1.0 / a_np.diagonal()
+    )
+    rhs_np = np.random.uniform(low=-1.0, high=1.0, size=shape[0]).astype(dtype)
+    x_np = np.zeros_like(rhs_np)
+    tol = 1e-6 if dtype == np.float64 else 1e-3
+    max_iter = 20
+    if use_static_shape:
+      a = constant_op.constant(a_np)
+      rhs = constant_op.constant(rhs_np)
+      x = constant_op.constant(x_np)
+      jacobi = constant_op.constant(jacobi_np)
+    else:
+      a = array_ops.placeholder_with_default(a_np, shape=None)
+      rhs = array_ops.placeholder_with_default(rhs_np, shape=None)
+      x = array_ops.placeholder_with_default(x_np, shape=None)
+      jacobi = array_ops.placeholder_with_default(jacobi_np, shape=None)
+    operator = linalg.LinearOperatorFullMatrix(
+        a, is_positive_definite=True, is_self_adjoint=True
+    )
+    preconditioners = [
+        None,
+        # Preconditioner that does nothing beyond change shape.
+        linalg.LinearOperatorIdentity(
+            a_np.shape[-1],
+            dtype=a_np.dtype,
+            is_positive_definite=True,
+            is_self_adjoint=True,
+        ),
+        # Jacobi preconditioner.
+        linalg.LinearOperatorFullMatrix(
+            jacobi, is_positive_definite=True, is_self_adjoint=True
+        ),
+    ]
+    cg_results = []
+    for preconditioner in preconditioners:
+      cg_graph = conjugate_gradient.conjugate_gradient(
+          operator,
+          rhs,
+          preconditioner=preconditioner,
+          x=x,
+          tol=tol,
+          max_iter=max_iter,
+      )
+      cg_val = self.evaluate(cg_graph)
+      norm_r0 = np.linalg.norm(rhs_np)
+      norm_r = np.linalg.norm(cg_val.r)
+      self.assertLessEqual(norm_r, tol * norm_r0)
+      # Validate that we get an equally small residual norm with numpy
+      # using the computed solution.
+      r_np = rhs_np - np.dot(a_np, cg_val.x)
+      norm_r_np = np.linalg.norm(r_np)
+      self.assertLessEqual(norm_r_np, tol * norm_r0)
+      cg_results.append(cg_val)
+    # Validate that we get same results using identity_preconditioner
+    # and None
+    self.assertEqual(cg_results[0].i, cg_results[1].i)
+    self.assertAlmostEqual(cg_results[0].gamma, cg_results[1].gamma)
+    self.assertAllClose(cg_results[0].r, cg_results[1].r, rtol=tol)
+    self.assertAllClose(cg_results[0].x, cg_results[1].x, rtol=tol)
+    self.assertAllClose(cg_results[0].p, cg_results[1].p, rtol=tol)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/kernel_tests/linalg/sparse/csr_sparse_matrix_dense_mat_mul_grad_test.py

@@ -0,0 +1,145 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+# ==============================================================================
+"""CSR sparse matrix tests."""
+
+import itertools
+
+import numpy as np
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradient_checker
+from tensorflow.python.ops.linalg.sparse import sparse_csr_matrix_grad  # pylint: disable=unused-import
+from tensorflow.python.ops.linalg.sparse import sparse_csr_matrix_ops
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging
+
+
+def _add_test(test, op_name, testcase_name, fn):  # pylint: disable=redefined-outer-name
+  if fn is None:
+    return
+  test_name = "_".join(["test", op_name, testcase_name])
+  if hasattr(test, test_name):
+    raise RuntimeError("Test %s defined more than once" % test_name)
+  setattr(test, test_name, fn)
+
+
+class CSRSparseMatrixDenseMatMulGradTest(test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super(CSRSparseMatrixDenseMatMulGradTest, cls).setUpClass()
+    cls._gpu_available = test_util.is_gpu_available()
+
+  # TODO(penporn): Make these tests runnable on eager mode.
+  # (tf.gradients and gradient_checker only run in graph mode.)
+  @test_util.run_deprecated_v1
+  def _testLargeBatchSparseMatrixMatMulGrad(
+      self,
+      datatype,
+      transpose_a,
+      transpose_b,
+      adjoint_a,
+      adjoint_b,
+      transpose_output,
+      conjugate_output,
+      batched_inputs,
+  ):
+    if batched_inputs:
+      a_shape = (3, 5, 11)
+      b_shape = (3, 11, 13)
+      transpose = lambda x: np.transpose(x, (0, 2, 1))
+    else:
+      a_shape = (5, 11)
+      b_shape = (11, 13)
+      transpose = np.transpose
+
+    sparsify = lambda m: m * (m > 0)
+    a_mats_val = sparsify(
+        np.random.randn(*a_shape) +
+        1.j * np.random.randn(*a_shape)).astype(datatype)
+    if transpose_a or adjoint_a:
+      a_mats_val = transpose(a_mats_val)
+    if adjoint_a:
+      a_mats_val = np.conj(a_mats_val)
+    b_mats_val = (np.random.randn(*b_shape) +
+                  1.j * np.random.randn(*b_shape)).astype(datatype)
+    if transpose_b or adjoint_b:
+      b_mats_val = transpose(b_mats_val)
+    if adjoint_b:
+      b_mats_val = np.conj(b_mats_val)
+    with self.test_session():
+      a_mats = ops.convert_to_tensor(a_mats_val, dtype=datatype)
+      b_mats = ops.convert_to_tensor(b_mats_val, dtype=datatype)
+      locs = array_ops.where(abs(a_mats_val) > 0)
+      a_sm = sparse_csr_matrix_ops.dense_to_csr_sparse_matrix(a_mats, locs)
+      c_mats = sparse_csr_matrix_ops.sparse_matrix_mat_mul(
+          a_sm,
+          b_mats,
+          transpose_a=transpose_a,
+          transpose_b=transpose_b,
+          adjoint_a=adjoint_a,
+          adjoint_b=adjoint_b,
+          transpose_output=transpose_output,
+          conjugate_output=conjugate_output)
+      for [ten, val, nn] in [[a_mats, a_mats_val, "a"],
+                             [b_mats, b_mats_val, "b"]]:
+        tf_logging.info("Testing gradients for %s" % nn)
+        theoretical, numerical = gradient_checker.compute_gradient(
+            ten,
+            ten.get_shape().as_list(),
+            c_mats,
+            c_mats.get_shape().as_list(),
+            x_init_value=val,
+            delta=1e-3)
+        self.assertAllClose(theoretical, numerical, atol=1e-3, rtol=1e-3)
+
+
+# These tests are refactored from sparse_csr_matrix_grad_test to keep its size
+# "medium".
+dtypes_to_test = [np.float32, np.complex64]
+for dtype in dtypes_to_test:
+  for (t_a, t_b, adj_a, adj_b, t_out,
+       conj_out, batched) in itertools.product(*(([False, True],) * 7)):
+
+    def create_mat_mul_test_fn(dtype_, t_a_, t_b_, adj_a_, adj_b_, t_out_,
+                               conj_out_, batched_):
+      # Skip invalid cases.
+      if (t_a_ and adj_a_) or (t_b_ and adj_b_):
+        return
+      # Skip cases where we conjugate real matrices.
+      if dtype_ == np.float32 and (adj_a_ or adj_b_ or conj_out_):
+        return
+
+      def test_fn(self):
+        self._testLargeBatchSparseMatrixMatMulGrad(dtype_, t_a_, t_b_, adj_a_,
+                                                   adj_b_, t_out_, conj_out_,
+                                                   batched_)
+
+      return test_fn
+
+    name = (
+        "_testLargeBatchSparseMatrixMatMulGrad_dtype_%s_t_a_%s_t_b_%s_adj_a_%s_"
+        "adj_b_%s_t_out_%s_conj_out_%s_batched_%s" %
+        (dtype.__name__, t_a, t_b, adj_a, adj_b, t_out, conj_out, batched))
+
+    _add_test(
+        CSRSparseMatrixDenseMatMulGradTest, "CSRSparseMatrixGradTest", name,
+        create_mat_mul_test_fn(dtype, t_a, t_b, adj_a, adj_b, t_out, conj_out,
+                               batched))
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/kernel_tests/linalg/sparse/csr_sparse_matrix_grad_test.py

@@ -0,0 +1,145 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+# ==============================================================================
+"""CSR sparse matrix tests."""
+
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.linalg.sparse import sparse_csr_matrix_grad  # pylint: disable=unused-import
+from tensorflow.python.ops.linalg.sparse import sparse_csr_matrix_ops
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging
+
+
+def dense_to_csr_sparse_matrix(dense):
+  dense_t = ops.convert_to_tensor(dense)
+  locs = array_ops.where(math_ops.abs(dense_t) > 0)
+  return sparse_csr_matrix_ops.dense_to_csr_sparse_matrix(dense_t, locs)
+
+
+def _add_test(test, op_name, testcase_name, fn):  # pylint: disable=redefined-outer-name
+  if fn is None:
+    return
+  test_name = "_".join(["test", op_name, testcase_name])
+  if hasattr(test, test_name):
+    raise RuntimeError("Test %s defined more than once" % test_name)
+  setattr(test, test_name, fn)
+
+
+class CSRSparseMatrixGradTest(test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super(CSRSparseMatrixGradTest, cls).setUpClass()
+    cls._gpu_available = test_util.is_gpu_available()
+
+  # TODO(penporn): Make these tests runnable on eager mode.
+  # (tf.gradients and gradient_checker only run in graph mode.)
+  @test_util.run_deprecated_v1
+  def testLargeBatchConversionGrad(self):
+    if not self._gpu_available:
+      return
+
+    sparsify = lambda m: m * (m > 0)
+    for dense_shape in ([53, 65, 127], [127, 65]):
+      mats_val = sparsify(np.random.randn(*dense_shape))
+      with self.test_session() as sess:
+        mats = math_ops.cast(mats_val, dtype=dtypes.float32)
+        sparse_mats = dense_to_csr_sparse_matrix(mats)
+        dense_mats = sparse_csr_matrix_ops.csr_sparse_matrix_to_dense(
+            sparse_mats, dtypes.float32)
+        grad_vals = np.random.randn(*dense_shape).astype(np.float32)
+        grad_out = gradients_impl.gradients([dense_mats], [mats],
+                                            [grad_vals])[0]
+        self.assertEqual(grad_out.dtype, dtypes.float32)
+        self.assertEqual(grad_out.shape, dense_shape)
+        grad_out_value = sess.run(grad_out)
+        tf_logging.info("testLargeBatchConversionGrad: Testing shape %s" %
+                        dense_shape)
+        nonzero_indices = abs(mats_val) > 0.0
+        self.assertAllEqual(grad_out_value[nonzero_indices],
+                            grad_vals[nonzero_indices])
+        self.assertTrue(
+            np.all(grad_out_value[np.logical_not(nonzero_indices)] == 0.0))
+
+  @test_util.run_deprecated_v1
+  def testLargeBatchSparseConversionGrad(self):
+    sparsify = lambda m: m * (m > 0)
+    for dense_shape in ([53, 65, 127], [127, 65]):
+      mats_val = sparsify(np.random.randn(*dense_shape))
+
+      with self.session(use_gpu=True) as sess:
+        indices = array_ops.where_v2(
+            math_ops.not_equal(mats_val, array_ops.zeros_like(mats_val)))
+        values = math_ops.cast(
+            array_ops.gather_nd(mats_val, indices), dtype=dtypes.float32)
+
+        grad_vals = np.random.randn(*sess.run(values).shape).astype(np.float32)
+        csr_matrix = sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
+            indices, values, dense_shape)
+        new_coo_tensor = (
+            sparse_csr_matrix_ops.csr_sparse_matrix_to_sparse_tensor(
+                csr_matrix, type=dtypes.float32))
+        grad_out = gradients_impl.gradients([new_coo_tensor.values], [values],
+                                            [grad_vals])[0]
+        self.assertEqual(grad_out.dtype, dtypes.float32)
+        grad_out_vals = sess.run(grad_out)
+        self.assertAllClose(grad_vals, grad_out_vals)
+
+  @test_util.run_deprecated_v1
+  def testLargeBatchSparseMatrixAddGrad(self):
+    if not self._gpu_available:
+      return
+
+    sparsify = lambda m: m * (m > 0)
+    for dense_shape in ([53, 65, 127], [127, 65]):
+      a_mats_val = sparsify(np.random.randn(*dense_shape))
+      b_mats_val = sparsify(np.random.randn(*dense_shape))
+      alpha = np.float32(0.5)
+      beta = np.float32(-1.5)
+      grad_vals = np.random.randn(*dense_shape).astype(np.float32)
+      expected_a_grad = alpha * grad_vals
+      expected_b_grad = beta * grad_vals
+      expected_a_grad[abs(a_mats_val) == 0.0] = 0.0
+      expected_b_grad[abs(b_mats_val) == 0.0] = 0.0
+      with self.test_session() as sess:
+        a_mats = math_ops.cast(a_mats_val, dtype=dtypes.float32)
+        b_mats = math_ops.cast(b_mats_val, dtype=dtypes.float32)
+        a_sm = dense_to_csr_sparse_matrix(a_mats)
+        b_sm = dense_to_csr_sparse_matrix(b_mats)
+        c_sm = sparse_csr_matrix_ops.sparse_matrix_add(
+            a_sm, b_sm, alpha=alpha, beta=beta)
+        c_dense = sparse_csr_matrix_ops.csr_sparse_matrix_to_dense(
+            c_sm, dtypes.float32)
+        a_grad, b_grad = gradients_impl.gradients([c_dense], [a_mats, b_mats],
+                                                  [grad_vals])
+        self.assertEqual(a_grad.dtype, dtypes.float32)
+        self.assertEqual(b_grad.dtype, dtypes.float32)
+        self.assertEqual(a_grad.shape, dense_shape)
+        self.assertEqual(b_grad.shape, dense_shape)
+        a_grad_value, b_grad_value = sess.run((a_grad, b_grad))
+        tf_logging.info("testLargeBatchConversionGrad: Testing shape %s" %
+                        dense_shape)
+        self.assertAllEqual(expected_a_grad, a_grad_value)
+        self.assertAllEqual(expected_b_grad, b_grad_value)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/kernel_tests/linalg/sparse/csr_sparse_matrix_sparse_mat_mul_grad_test.py

@@ -0,0 +1,132 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+# ==============================================================================
+"""CSR sparse matrix tests."""
+
+import itertools
+
+import numpy as np
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradient_checker
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.linalg.sparse import sparse_csr_matrix_grad  # pylint: disable=unused-import
+from tensorflow.python.ops.linalg.sparse import sparse_csr_matrix_ops
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging
+
+
+def dense_and_sparse_from_vals(vals, datatype):
+  locs = array_ops.where(math_ops.abs(vals) > 0)
+  dense_t = ops.convert_to_tensor(vals, dtype=datatype)
+  return (dense_t,
+          sparse_csr_matrix_ops.dense_to_csr_sparse_matrix(dense_t, locs))
+
+
+def _add_test(test, op_name, testcase_name, fn):  # pylint: disable=redefined-outer-name
+  if fn is None:
+    return
+  test_name = "_".join(["test", op_name, testcase_name])
+  if hasattr(test, test_name):
+    raise RuntimeError("Test %s defined more than once" % test_name)
+  setattr(test, test_name, fn)
+
+
+class CSRSparseMatrixGradTest(test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super(CSRSparseMatrixGradTest, cls).setUpClass()
+    cls._gpu_available = test_util.is_gpu_available()
+
+  # TODO(penporn): Make these tests runnable on eager mode.
+  # (tf.gradients and gradient_checker only run in graph mode.)
+  @test_util.run_deprecated_v1
+  def _testLargeBatchSparseMatrixSparseMatMulGrad(self, datatype, transpose_a,
+                                                  transpose_b, adjoint_a,
+                                                  adjoint_b):
+    if not self._gpu_available:
+      return
+
+    sparsify = lambda m: m * (m > 0)
+    a_mats_val = sparsify(
+        np.random.randn(3, 5, 11) +
+        1.j * np.random.randn(3, 5, 11)).astype(datatype)
+    if transpose_a or adjoint_a:
+      a_mats_val = np.transpose(a_mats_val, (0, 2, 1))
+    if adjoint_a:
+      a_mats_val = np.conj(a_mats_val)
+    b_mats_val = sparsify(
+        np.random.randn(3, 11, 13) +
+        1.j * np.random.randn(3, 11, 13)).astype(datatype)
+    if transpose_b or adjoint_b:
+      b_mats_val = np.transpose(b_mats_val, (0, 2, 1))
+    if adjoint_b:
+      b_mats_val = np.conj(b_mats_val)
+    with self.test_session():
+      a_mats, a_sm = dense_and_sparse_from_vals(a_mats_val, datatype)
+      b_mats, b_sm = dense_and_sparse_from_vals(b_mats_val, datatype)
+      c_sm = sparse_csr_matrix_ops.sparse_matrix_sparse_mat_mul(
+          a_sm,
+          b_sm,
+          transpose_a=transpose_a,
+          transpose_b=transpose_b,
+          adjoint_a=adjoint_a,
+          adjoint_b=adjoint_b,
+          type=datatype)
+      c_dense = sparse_csr_matrix_ops.csr_sparse_matrix_to_dense(
+          c_sm, type=datatype)
+      for ten, val, nn in [[a_mats, a_mats_val, "a"], [b_mats, b_mats_val,
+                                                       "b"]]:
+        tf_logging.info("Testing gradients for %s" % nn)
+        theoretical, numerical = gradient_checker.compute_gradient(
+            ten,
+            ten.get_shape().as_list(),
+            c_dense,
+            c_dense.get_shape().as_list(),
+            x_init_value=val,
+            delta=1e-3)
+        self.assertAllClose(theoretical, numerical, atol=1e-3, rtol=1e-3)
+
+
+# These tests are refactored from sparse_csr_matrix_grad_test to keep its size
+# "medium".
+for dtype in (np.float32, np.complex64):
+  for (t_a, t_b, adj_a, adj_b) in itertools.product(*(([False, True],) * 4)):
+
+    def create_sparse_mat_mul_test_fn(dtype_, t_a_, t_b_, adj_a_, adj_b_):
+      # Skip invalid cases.
+      if (t_a_ and adj_a_) or (t_b_ and adj_b_):
+        return
+      # Skip cases where we conjugate real matrices.
+      if dtype_ == np.float32 and (adj_a_ or adj_b_):
+        return
+
+      def test_fn(self):
+        self._testLargeBatchSparseMatrixSparseMatMulGrad(
+            dtype_, t_a_, t_b_, adj_a_, adj_b_)
+
+      return test_fn
+
+    name = (
+        "_testLargeBatchSparseMatrixSparseMatMulGrad_dtype_%s_t_a_%s_t_b_%s_"
+        "adj_a_%s_adj_b_%s" % (dtype.__name__, t_a, t_b, adj_a, adj_b))
+
+    _add_test(CSRSparseMatrixGradTest, "CSRSparseMatrixSparseGradTest", name,
+              create_sparse_mat_mul_test_fn(dtype, t_a, t_b, adj_a, adj_b))
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/kernel_tests/linalg/sparse/csr_sparse_matrix_test.py

@@ -0,0 +1,265 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+# ==============================================================================
+"""CSR sparse matrix tests."""
+
+import itertools
+
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.linalg.sparse import sparse_csr_matrix_ops
+from tensorflow.python.platform import test
+
+
+class CSRSparseMatrixTest(test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):  # pylint: disable=g-missing-super-call
+    cls._gpu_available = test_util.is_gpu_available()
+
+  @test_util.run_in_graph_and_eager_modes
+  def testConstructorFromSparseTensor(self):
+    if not self._gpu_available:
+      return
+
+    a_indices = np.array([[0, 0], [2, 3], [2, 4], [3, 0]])
+    a_values = [1.0, 5.0, -1.0, -2.0]
+    a_dense_shape = [5, 6]
+
+    a_st = sparse_tensor.SparseTensor(a_indices, a_values, a_dense_shape)
+    a_st = math_ops.cast(a_st, dtypes.float32)
+    a_sm = sparse_csr_matrix_ops.CSRSparseMatrix(a_st)
+    self.assertEqual(a_sm.shape, a_dense_shape)
+
+    a_st_rt = a_sm.to_sparse_tensor()
+    a_st_rt = self.evaluate(a_st_rt)
+
+    self.assertAllEqual(a_indices, a_st_rt.indices)
+    self.assertAllClose(a_values, a_st_rt.values)
+    self.assertAllEqual(a_dense_shape, a_st_rt.dense_shape)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testConstructorFromDenseTensorNoIndices(self):
+    if not self._gpu_available:
+      return
+
+    sparsify = lambda m: m * (m > 0)
+    dense_shape = [5, 7, 13]
+    a_mats = sparsify(np.random.randn(*dense_shape)).astype(np.float32)
+
+    a_sm = sparse_csr_matrix_ops.CSRSparseMatrix(a_mats)
+    self.assertEqual(a_sm.shape, a_mats.shape)
+
+    a_sm_rt = a_sm.to_dense()
+    a_sm_nnz = a_sm.nnz()
+    a_sm_nnz, a_sm_rt = self.evaluate([a_sm_nnz, a_sm_rt])
+
+    # Count number of nonzero entries for each batch using bincount.
+    nz = np.bincount(a_mats.nonzero()[0], minlength=a_mats.shape[0])
+    self.assertAllEqual(nz, a_sm_nnz)
+    self.assertAllClose(a_mats, a_sm_rt)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testConstructorFromDenseTensorWithIndices(self):
+    if not self._gpu_available:
+      return
+
+    dense_shape = [5, 7, 13]
+    a_mats = np.random.randn(*dense_shape).astype(np.float32)
+    indices = np.array([[0, 0, 0],
+                        [1, 0, 0]], dtype=np.int64)
+
+    a_sm = sparse_csr_matrix_ops.CSRSparseMatrix(a_mats, indices=indices)
+    self.assertEqual(a_sm.shape, a_mats.shape)
+
+    a_sm_st = a_sm.to_sparse_tensor()
+    a_sm_st = self.evaluate(a_sm_st)
+
+    # Count number of nonzero entries for each batch using bincount.
+    self.assertAllEqual(indices, a_sm_st.indices)
+    self.assertAllEqual(dense_shape, a_sm.shape)
+    self.assertAllEqual(dense_shape, a_sm_st.dense_shape)
+    self.assertAllClose([a_mats[tuple(x)] for x in indices], a_sm_st.values)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testConj(self):
+    if not self._gpu_available:
+      return
+
+    sparsify = lambda m: m * (m.real > 0)
+    dense_shape = [5, 7, 13]
+    a_mats = sparsify(
+        (np.random.randn(*dense_shape) + 1.j * np.random.randn(*dense_shape))
+        .astype(np.complex64))
+    a_sm = sparse_csr_matrix_ops.CSRSparseMatrix(a_mats)
+    a_sm_conj = a_sm.conj()
+    self.assertIsInstance(a_sm_conj, sparse_csr_matrix_ops.CSRSparseMatrix)
+    a_sm_conj_dense = a_sm_conj.to_dense()
+    a_sm_conj_dense = self.evaluate(a_sm_conj_dense)
+    self.assertAllClose(a_mats.conj(), a_sm_conj_dense)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testTranspose(self):
+    if not self._gpu_available:
+      return
+
+    for conjugate in False, True:
+      sparsify = lambda m: m * (m > 0)
+      dense_shape = [5, 7, 13]
+      a_mats = sparsify((np.random.randn(*dense_shape) +
+                         1.j * np.random.randn(*dense_shape))).astype(
+                             np.complex64)
+      expected = np.transpose(a_mats, (0, 2, 1))
+      if conjugate:
+        expected = np.conj(expected)
+      a_sm = sparse_csr_matrix_ops.CSRSparseMatrix(a_mats)
+      if conjugate:
+        a_sm_t = a_sm.hermitian_transpose()
+      else:
+        a_sm_t = a_sm.transpose()
+      self.assertIsInstance(a_sm_t, sparse_csr_matrix_ops.CSRSparseMatrix)
+      a_sm_t_dense = a_sm_t.to_dense()
+      a_sm_t_dense = self.evaluate(a_sm_t_dense)
+      self.assertAllClose(expected, a_sm_t_dense)
+
+
+class SparseMatrixMatmulTest(test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):  # pylint: disable=g-missing-super-call
+    cls._gpu_available = test_util.is_gpu_available()
+
+  def _testSparseSparse(self, transpose_a, transpose_b, adjoint_a, adjoint_b):
+    if not self._gpu_available:
+      return
+    sparsify = lambda m: m * (m > 0)
+    dense_shape_a = [5, 13, 7] if transpose_a or adjoint_a else [5, 7, 13]
+    dense_shape_b = [5, 15, 13] if transpose_b or adjoint_b else [5, 13, 15]
+    dtypes_to_test = [np.float32, np.complex64]
+    for dtype in dtypes_to_test:
+      a_mats = sparsify((np.random.randn(*dense_shape_a) +
+                         1.j * np.random.randn(*dense_shape_a))).astype(dtype)
+      b_mats = sparsify((np.random.randn(*dense_shape_b) +
+                         1.j * np.random.randn(*dense_shape_b))).astype(dtype)
+      a_sm = sparse_csr_matrix_ops.CSRSparseMatrix(a_mats)
+      b_sm = sparse_csr_matrix_ops.CSRSparseMatrix(b_mats)
+      c_dense = test_util.matmul_without_tf32(
+          a_mats,
+          b_mats,
+          transpose_a=transpose_a,
+          transpose_b=transpose_b,
+          adjoint_a=adjoint_a,
+          adjoint_b=adjoint_b)
+      c_sm = sparse_csr_matrix_ops.matmul(
+          a_sm,
+          b_sm,
+          transpose_a=transpose_a,
+          transpose_b=transpose_b,
+          adjoint_a=adjoint_a,
+          adjoint_b=adjoint_b)
+      self.assertIsInstance(c_sm, sparse_csr_matrix_ops.CSRSparseMatrix)
+      c_sm_dense = c_sm.to_dense()
+      c_dense, c_sm_dense = self.evaluate([c_dense, c_sm_dense])
+      self.assertAllClose(c_dense, c_sm_dense)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testSparseSparse(self):
+    for (t_a, t_b, adj_a, adj_b) in itertools.product(*(([False, True],) * 4)):
+      if (t_a and adj_a) or (t_b and adj_b):
+        continue
+      self._testSparseSparse(t_a, t_b, adj_a, adj_b)
+
+  def _testSparseDense(self, transpose_a, transpose_b, adjoint_a, adjoint_b):
+    if not self._gpu_available:
+      return
+
+    sparsify = lambda m: m * (m > 0)
+    dense_shape_a = [5, 13, 7] if transpose_a or adjoint_a else [5, 7, 13]
+    dense_shape_b = [5, 15, 13] if transpose_b or adjoint_b else [5, 13, 15]
+    dtypes_to_test = [np.float32, np.complex64]
+    for dtype in dtypes_to_test:
+      a_mats = sparsify((np.random.randn(*dense_shape_a) +
+                         1.j * np.random.randn(*dense_shape_a))).astype(dtype)
+      b_mats = (np.random.randn(*dense_shape_b) +
+                1.j * np.random.randn(*dense_shape_b)).astype(dtype)
+      a_sm = sparse_csr_matrix_ops.CSRSparseMatrix(a_mats)
+      c_dense = test_util.matmul_without_tf32(
+          a_mats,
+          b_mats,
+          transpose_a=transpose_a,
+          transpose_b=transpose_b,
+          adjoint_a=adjoint_a,
+          adjoint_b=adjoint_b)
+      c_sm_dense = sparse_csr_matrix_ops.matmul(
+          a_sm,
+          b_mats,
+          transpose_a=transpose_a,
+          transpose_b=transpose_b,
+          adjoint_a=adjoint_a,
+          adjoint_b=adjoint_b)
+      c_dense, c_sm_dense = self.evaluate([c_dense, c_sm_dense])
+      self.assertAllClose(c_dense, c_sm_dense)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testSparseDense(self):
+    for (t_a, t_b, adj_a, adj_b) in itertools.product(*(([False, True],) * 4)):
+      if (t_a and adj_a) or (t_b and adj_b):
+        continue
+      self._testSparseDense(t_a, t_b, adj_a, adj_b)
+
+  def _testDenseSparse(self, transpose_a, transpose_b, adjoint_a, adjoint_b):
+    if not self._gpu_available:
+      return
+
+    sparsify = lambda m: m * (m > 0)
+    dense_shape_a = [5, 13, 7] if transpose_a or adjoint_a else [5, 7, 13]
+    dense_shape_b = [5, 15, 13] if transpose_b or adjoint_b else [5, 13, 15]
+    dtypes_to_test = [np.float32, np.complex64]
+    for dtype in dtypes_to_test:
+      a_mats = (np.random.randn(*dense_shape_a) +
+                1.j * np.random.randn(*dense_shape_a)).astype(dtype)
+      b_mats = sparsify((np.random.randn(*dense_shape_b) +
+                         1.j * np.random.randn(*dense_shape_b))).astype(dtype)
+      b_sm = sparse_csr_matrix_ops.CSRSparseMatrix(b_mats)
+      c_dense = test_util.matmul_without_tf32(
+          a_mats,
+          b_mats,
+          transpose_a=transpose_a,
+          transpose_b=transpose_b,
+          adjoint_a=adjoint_a,
+          adjoint_b=adjoint_b)
+      c_sm_dense = sparse_csr_matrix_ops.matmul(
+          a_mats,
+          b_sm,
+          transpose_a=transpose_a,
+          transpose_b=transpose_b,
+          adjoint_a=adjoint_a,
+          adjoint_b=adjoint_b)
+      c_dense, c_sm_dense = self.evaluate([c_dense, c_sm_dense])
+      self.assertAllClose(c_dense, c_sm_dense)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testDenseSparse(self):
+    for (t_a, t_b, adj_a, adj_b) in itertools.product(*(([False, True],) * 4)):
+      if (t_a and adj_a) or (t_b and adj_b):
+        continue
+      self._testDenseSparse(t_a, t_b, adj_a, adj_b)
+
+
+if __name__ == "__main__":
+  test.main()