chore: import upstream snapshot with attribution

tensorflow/python/ops/linalg/sparse/BUILD

@@ -0,0 +1,93 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+
+# Description: Sparse CSR support for TensorFlow.
+load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+tf_gen_op_wrapper_py(
+    name = "gen_sparse_csr_matrix_ops",
+    out = "gen_sparse_csr_matrix_ops.py",
+    api_def_srcs = ["//tensorflow/core/api_def:base_api_def"],
+    extra_py_deps = [
+        "//tensorflow/python:pywrap_tfe",
+        "//tensorflow/python/util:dispatch",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+    py_lib_rule = py_library,
+    visibility = ["//visibility:private"],
+    deps = ["//tensorflow/core:sparse_csr_matrix_ops_op_lib"],
+)
+
+py_library(
+    name = "sparse_py",
+    srcs = ["sparse.py"],
+    strict_deps = True,
+    deps = [
+        ":conjugate_gradient",
+        ":sparse_csr_matrix_grad",
+        ":sparse_csr_matrix_ops",
+    ],
+)
+
+py_library(
+    name = "sparse_csr_matrix_grad",
+    srcs = ["sparse_csr_matrix_grad.py"],
+    strict_deps = True,
+    deps = [
+        ":sparse_csr_matrix_ops",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:sparse_ops",
+    ],
+)
+
+py_library(
+    name = "__init__",
+    srcs = ["__init__.py"],
+    strict_deps = True,
+)
+
+pytype_strict_library(
+    name = "conjugate_gradient",
+    srcs = ["conjugate_gradient.py"],
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:while_loop",
+        "//tensorflow/python/ops/linalg:linalg_impl",
+        "//tensorflow/python/util:dispatch",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "sparse_csr_matrix_ops",
+    srcs = ["sparse_csr_matrix_ops.py"],
+    strict_deps = True,
+    deps = [
+        ":gen_sparse_csr_matrix_ops",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:cpp_shape_inference_proto_py",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+    ],
+)

tensorflow/python/ops/linalg/sparse/conjugate_gradient.py

@@ -0,0 +1,138 @@
+# 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.
+# ==============================================================================
+"""Preconditioned Conjugate Gradient."""
+
+import collections
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import while_loop
+from tensorflow.python.ops.linalg import linalg_impl as linalg
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export('linalg.experimental.conjugate_gradient')
+@dispatch.add_dispatch_support
+def conjugate_gradient(operator,
+                       rhs,
+                       preconditioner=None,
+                       x=None,
+                       tol=1e-5,
+                       max_iter=20,
+                       name='conjugate_gradient'):
+  r"""Conjugate gradient solver.
+
+  Solves a linear system of equations `A*x = rhs` for self-adjoint, positive
+  definite matrix `A` and right-hand side vector `rhs`, using an iterative,
+  matrix-free algorithm where the action of the matrix A is represented by
+  `operator`. The iteration terminates when either the number of iterations
+  exceeds `max_iter` or when the residual norm has been reduced to `tol`
+  times its initial value, i.e. \\(||rhs - A x_k|| <= tol ||rhs||\\).
+
+  Args:
+    operator: A `LinearOperator` that is self-adjoint and positive definite.
+    rhs: A possibly batched vector of shape `[..., N]` containing the right-hand
+      size vector.
+    preconditioner: A `LinearOperator` that approximates the inverse of `A`.
+      An efficient preconditioner could dramatically improve the rate of
+      convergence. If `preconditioner` represents matrix `M`(`M` approximates
+      `A^{-1}`), the algorithm uses `preconditioner.apply(x)` to estimate
+      `A^{-1}x`. For this to be useful, the cost of applying `M` should be
+      much lower than computing `A^{-1}` directly.
+    x: A possibly batched vector of shape `[..., N]` containing the initial
+      guess for the solution.
+    tol: A float scalar convergence tolerance.
+    max_iter: An integer giving the maximum number of iterations.
+    name: A name scope for the operation.
+
+  Returns:
+    output: A namedtuple representing the final state with fields:
+      - i: A scalar `int32` `Tensor`. Number of iterations executed.
+      - x: A rank-1 `Tensor` of shape `[..., N]` containing the computed
+          solution.
+      - r: A rank-1 `Tensor` of shape `[.., M]` containing the residual vector.
+      - p: A rank-1 `Tensor` of shape `[..., N]`. `A`-conjugate basis vector.
+      - gamma: \\(r \dot M \dot r\\), equivalent to  \\(||r||_2^2\\) when
+        `preconditioner=None`.
+  """
+  if not (operator.is_self_adjoint and operator.is_positive_definite):
+    raise ValueError('Expected a self-adjoint, positive definite operator.')
+
+  cg_state = collections.namedtuple('CGState', ['i', 'x', 'r', 'p', 'gamma'])
+
+  def stopping_criterion(i, state):
+    return math_ops.logical_and(
+        i < max_iter,
+        math_ops.reduce_any(linalg.norm(state.r, axis=-1) > tol))
+
+  def dot(x, y):
+    return array_ops.squeeze(
+        math_ops.matvec(
+            x[..., array_ops.newaxis],
+            y, adjoint_a=True), axis=-1)
+
+  def cg_step(i, state):  # pylint: disable=missing-docstring
+    z = math_ops.matvec(operator, state.p)
+    alpha = state.gamma / dot(state.p, z)
+    x = state.x + alpha[..., array_ops.newaxis] * state.p
+    r = state.r - alpha[..., array_ops.newaxis] * z
+    if preconditioner is None:
+      q = r
+    else:
+      q = preconditioner.matvec(r)
+    gamma = dot(r, q)
+    beta = gamma / state.gamma
+    p = q + beta[..., array_ops.newaxis] * state.p
+    return i + 1, cg_state(i + 1, x, r, p, gamma)
+
+  # We now broadcast initial shapes so that we have fixed shapes per iteration.
+
+  with ops.name_scope(name):
+    broadcast_shape = array_ops.broadcast_dynamic_shape(
+        array_ops.shape(rhs)[:-1],
+        operator.batch_shape_tensor())
+    if preconditioner is not None:
+      broadcast_shape = array_ops.broadcast_dynamic_shape(
+          broadcast_shape,
+          preconditioner.batch_shape_tensor()
+      )
+    broadcast_rhs_shape = array_ops.concat([
+        broadcast_shape, [array_ops.shape(rhs)[-1]]], axis=-1)
+    r0 = array_ops.broadcast_to(rhs, broadcast_rhs_shape)
+    tol *= linalg.norm(r0, axis=-1)
+
+    if x is None:
+      x = array_ops.zeros(
+          broadcast_rhs_shape, dtype=rhs.dtype.base_dtype)
+    else:
+      r0 = rhs - math_ops.matvec(operator, x)
+    if preconditioner is None:
+      p0 = r0
+    else:
+      p0 = math_ops.matvec(preconditioner, r0)
+    gamma0 = dot(r0, p0)
+    i = constant_op.constant(0, dtype=dtypes.int32)
+    state = cg_state(i=i, x=x, r=r0, p=p0, gamma=gamma0)
+    _, state = while_loop.while_loop(stopping_criterion, cg_step, [i, state])
+    return cg_state(
+        state.i,
+        x=state.x,
+        r=state.r,
+        p=state.p,
+        gamma=state.gamma)

tensorflow/python/ops/linalg/sparse/sparse.py

@@ -0,0 +1,26 @@
+# 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.
+# ==============================================================================
+"""Public API for tf.linalg.sparse namespace."""
+
+# go/tf-wildcard-import
+# pylint: disable=wildcard-import
+from tensorflow.python.ops.linalg.sparse.conjugate_gradient import conjugate_gradient
+from tensorflow.python.ops.linalg.sparse.sparse_csr_matrix_grad import *
+from tensorflow.python.ops.linalg.sparse.sparse_csr_matrix_ops import *
+# pylint: enable=wildcard-import
+
+__all__ = [
+    'conjugate_gradient'
+]

tensorflow/python/ops/linalg/sparse/sparse_csr_matrix_grad.py

@@ -0,0 +1,364 @@
+# 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 Gradients."""
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops.linalg.sparse import sparse_csr_matrix_ops
+
+
+@ops.RegisterGradient("DenseToCSRSparseMatrix")
+def _DenseToCSRSparseMatrixGrad(op: ops.Operation, grad):
+  """Gradient for dense_to_csr_sparse_matrix op."""
+  grad_values = (
+      sparse_csr_matrix_ops.csr_sparse_matrix_to_dense(
+          grad, type=op.get_attr("T")))
+  # inputs to fw op were: params, indices.
+  return (grad_values, None)
+
+
+@ops.RegisterGradient("CSRSparseMatrixToDense")
+def _CSRSparseMatrixToDenseGrad(op: ops.Operation, grad):
+  """Gradient for csr_sparse_matrix_to_dense op."""
+  coo_sparse_tensor = sparse_csr_matrix_ops.csr_sparse_matrix_to_sparse_tensor(
+      op.inputs[0], type=grad.dtype)
+  return sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
+      indices=coo_sparse_tensor.indices,
+      values=array_ops.gather_nd(grad, coo_sparse_tensor.indices),
+      dense_shape=grad.shape)
+
+
+@ops.RegisterGradient("SparseTensorToCSRSparseMatrix")
+def _SparseTensorToCSRSparseMatrixGrad(op: ops.Operation, grad):
+  """Gradient for sparse_tensor_to_csr_sparse_matrix op."""
+  grad_values = sparse_csr_matrix_ops.csr_sparse_matrix_to_sparse_tensor(
+      grad, type=op.get_attr("T")).values
+  return (None, grad_values, None)
+
+
+@ops.RegisterGradient("CSRSparseMatrixToSparseTensor")
+def _CSRSparseMatrixToSparseTensorGrad(op: ops.Operation, *grads):
+  """Gradient for csr_sparse_matrix_to_sparse_tensor op."""
+  return sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
+      indices=op.outputs[0], values=grads[1], dense_shape=op.outputs[2])
+
+
+ops.NotDifferentiable("SparseMatrixNNZ")
+
+ops.NotDifferentiable("SparseMatrixZeros")
+
+
+def _PruneSparseTensor(unpruned, pruned_pattern):
+  """Helper function to prune COO sparse tensor.
+
+  Given two sparse tensors 'unpruned' and 'pruned_pattern', generates another
+  sparse tensor with indices and values fron 'unpruned' only if its indices also
+  occur in pruned_pattern.
+
+  Args:
+    unpruned: COO matrix with unpruned indices
+    pruned_pattern: COO matrix with pruned pattern.
+
+  TODO(tabakg): This is far from optimal. Consider a C++ implementation.
+
+  Returns:
+    Indices, values, and dense_shape of the pruned matrix.
+  """
+  pruned_indices = sparse_ops.sparse_reshape(
+      pruned_pattern, shape=(-1,)).indices[..., 0]
+  unpruned_indices = sparse_ops.sparse_reshape(
+      unpruned, shape=(-1,)).indices[..., 0]
+  best_match = array_ops.searchsorted(unpruned_indices, pruned_indices)
+  keep_indices = array_ops.gather(
+      best_match,
+      array_ops.where(
+          math_ops.equal(
+              array_ops.gather(unpruned_indices, best_match), pruned_indices)))
+  return (array_ops.gather_nd(unpruned.indices, keep_indices),
+          array_ops.gather_nd(unpruned.values,
+                              keep_indices), pruned_pattern.dense_shape)
+
+
+def _PruneCSRMatrix(unpruned, pruned_pattern):
+  """TODO(tabakg): Consider re-writing in C++."""
+  _, dtype = sparse_csr_matrix_ops.dense_shape_and_type(pruned_pattern)
+  coo_unpruned = sparse_tensor.SparseTensor(
+      *sparse_csr_matrix_ops.csr_sparse_matrix_to_sparse_tensor(
+          unpruned, type=dtype))
+  coo_pruned_pattern = sparse_tensor.SparseTensor(
+      *sparse_csr_matrix_ops.csr_sparse_matrix_to_sparse_tensor(
+          pruned_pattern, type=dtype))
+  return sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
+      *_PruneSparseTensor(coo_unpruned, coo_pruned_pattern))
+
+
+@ops.RegisterGradient("SparseMatrixAdd")
+def _SparseMatrixAddGrad(op: ops.Operation, grad):
+  """Gradient for sparse_matrix_add op."""
+  # input to sparse_matrix_add is (a, b, alpha, beta)
+  # with a, b CSR and alpha beta scalars.
+  # output is: alpha * a + beta * b
+
+  # d(a*A + b*B)/dA . grad = a * grad
+
+  # May have gotten the transposes wrong below.
+  # d(a*A + b*B)/da . grad = tr(A' . grad)
+
+  # For now, only implement gradients w.r.t. A and B.
+  # TODO(ebrevdo): Implement reduce_sum for SparseMatrix so that we
+  # can implement gradients w.r.t. a and b.
+  (a_csr, b_csr, alpha, beta) = op.inputs
+  return (sparse_csr_matrix_ops.sparse_matrix_mul(
+      _PruneCSRMatrix(grad, a_csr), alpha),
+          sparse_csr_matrix_ops.sparse_matrix_mul(
+              _PruneCSRMatrix(grad, b_csr), beta), None, None)
+
+
+def _PrunedDenseMatrixMultiplication(a,
+                                     b,
+                                     indices,
+                                     transpose_a=False,
+                                     adjoint_a=False,
+                                     transpose_b=False,
+                                     adjoint_b=False):
+  """Multiplies two dense matrices at selected indices.
+
+  The two inputs `a` and `b` must have matching rank (2 or 3). If using rank 3,
+  the first rank is used for the batch number. The last two dimensions should
+  also be compatible for matrix multiplication.
+
+  TODO(tabakg): Consider C++ implementation. There is also a more efficient way
+  to handle transposes here.
+
+  Args:
+    a: The left dense matrix (or batched matrices).
+    b: The right dense matrix (or batched matrices).
+    indices: The selected output indices where values should be produced. Other
+      indices will be pruned (not computed in the first place). Indices are
+      specified as a tensor of shape (length, rank), where length is the number
+      of entries and rank is the rank of the dense inputs (2 or 3).
+    transpose_a: Whether to transpose a.
+    adjoint_a: Whether to take the conjugate transpose of a.
+    transpose_b: Whether to transpose b.
+    adjoint_b: Whether to take the conjugate transpose of b.
+
+  Returns:
+    A CSR matrix.
+  """
+  transpose_a = transpose_a or adjoint_a
+  transpose_b = transpose_b or adjoint_b
+
+  a = math_ops.conj(a) if adjoint_a else a
+  b = math_ops.conj(b) if adjoint_b else b
+
+  rank = len(a.shape)
+  dense_shape = (a.shape[-1] if transpose_a else a.shape[-2],
+                 b.shape[-2] if transpose_b else b.shape[-1])
+  if rank == 2:
+    rows = indices[:, 0]
+    cols = indices[:, 1]
+    transpose = array_ops.transpose
+    gather_op = array_ops.gather
+  elif rank == 3:
+    dense_shape = (a.shape[0],) + dense_shape
+    rows = indices[:, :2]
+    cols = array_ops_stack.stack([indices[:, 0], indices[:, 2]], axis=1)
+    transpose = lambda x: array_ops.transpose(x, perm=[0, 2, 1])
+    gather_op = array_ops.gather_nd
+
+  a_rows = gather_op(transpose(a) if transpose_a else a, indices=rows)
+  b_cols = gather_op(b if transpose_b else transpose(b), indices=cols)
+  values = math_ops.reduce_sum(a_rows * b_cols, axis=1)
+
+  return sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
+      indices=indices, values=values, dense_shape=dense_shape)
+
+
+@ops.RegisterGradient("SparseMatrixTranspose")
+def _SparseMatrixTransposeGrad(op: ops.Operation, grad):
+  """Gradient for sparse_matrix_transpose op."""
+  return sparse_csr_matrix_ops.sparse_matrix_transpose(
+      grad, type=op.get_attr("type"), conjugate=op.get_attr("conjugate"))
+
+
+@ops.RegisterGradient("SparseMatrixSoftmax")
+def _SparseMatrixSoftmaxGrad(op: ops.Operation, grad_softmax):
+  """Gradient for sparse_matrix_softmax op."""
+  softmax = op.outputs[0]
+  return sparse_csr_matrix_ops.sparse_matrix_softmax_grad(
+      softmax, grad_softmax, type=op.get_attr("type"))
+
+
+@ops.RegisterGradient("SparseMatrixMatMul")
+def _SparseMatrixMatMulGrad(op: ops.Operation, grad):
+  """Gradient for sparse_matrix_mat_mul op."""
+  # input to sparse_matrix_mat_mul is (A, B) with CSR A and dense B.
+  # Output is dense:
+  #   C = opA(A) . opB(B) if transpose_output = false
+  #   C = (opA(A) . opB(B))' = opB(B)' . opA(A)' if transpose_output = true.
+  # where opA = transpose if transpose_a = True else identity
+  # and   opB = transpose if transpose_b = True else identity
+
+  t_a = op.get_attr("transpose_a")
+  t_b = op.get_attr("transpose_b")
+  adj_a = op.get_attr("adjoint_a")
+  adj_b = op.get_attr("adjoint_b")
+  transpose_output = op.get_attr("transpose_output")
+  conjugate_output = op.get_attr("conjugate_output")
+  a = op.inputs[0]  # sparse matrix
+  b = op.inputs[1]  # dense matrix
+  conj = math_ops.conj
+  sparse_matmul = sparse_csr_matrix_ops.sparse_matrix_mat_mul
+
+  def matmul(x, y, **kwargs):  # pylint: disable=invalid-name
+    return _PrunedDenseMatrixMultiplication(
+        x,
+        y,
+        indices=sparse_csr_matrix_ops.csr_sparse_matrix_to_sparse_tensor(
+            a, type=x.dtype).indices,
+        **kwargs)
+
+  if conjugate_output:
+    grad = conj(grad)
+  if not transpose_output:
+    # C = opA(A) . opB(B)
+    if not adj_a and not adj_b:
+      a = conj(a)
+      b = conj(b)
+      if not t_a:
+        grad_a = matmul(grad, b, transpose_b=not t_b)
+      else:
+        grad_a = matmul(b, grad, transpose_a=t_b, transpose_b=True)
+      grad_b = sparse_matmul(a, grad, transpose_a=not t_a, transpose_output=t_b)
+    elif not t_a and not t_b:
+      if not adj_a:
+        grad_a = matmul(grad, b, adjoint_b=not adj_b)
+      else:
+        grad_a = matmul(b, grad, adjoint_a=adj_b, adjoint_b=True)
+      grad_b = sparse_matmul(
+          a,
+          grad,
+          adjoint_a=not adj_a,
+          transpose_output=adj_b,
+          conjugate_output=adj_b)
+    elif adj_a and t_b:
+      grad_a = matmul(b, grad, transpose_a=True, adjoint_b=True)
+      grad_b = sparse_matmul(a, grad, transpose_output=True)
+    elif t_a and adj_b:
+      grad_a = matmul(b, grad, transpose_a=True, transpose_b=True)
+      grad_b = sparse_matmul(
+          conj(a), grad, transpose_output=True, conjugate_output=True)
+  else:
+    # C = (opA(A) . opB(B))' =  opB(B)' . opA(A)'
+    if not adj_a and not adj_b:
+      a = conj(a)
+      b = conj(b)
+      if not t_a:
+        grad_a = matmul(grad, b, transpose_a=True, transpose_b=not t_b)
+      else:
+        grad_a = matmul(b, grad, transpose_a=t_b)
+      grad_b = sparse_matmul(
+          a, grad, transpose_a=not t_a, transpose_b=True, transpose_output=t_b)
+    elif not t_a and not t_b:
+      if not adj_a:
+        grad_a = matmul(grad, b, transpose_a=True, adjoint_b=not adj_b)
+      else:
+        grad_a = matmul(b, conj(grad), adjoint_a=adj_b)
+      grad_b = sparse_matmul(
+          a,
+          grad,
+          adjoint_a=not adj_a,
+          transpose_b=True,
+          transpose_output=adj_b,
+          conjugate_output=adj_b)
+    elif adj_a and t_b:
+      grad_a = matmul(b, conj(grad), transpose_a=True)
+      grad_b = sparse_matmul(a, grad, transpose_b=True, transpose_output=True)
+    elif t_a and adj_b:
+      grad_a = matmul(b, grad, transpose_a=True)
+      grad_b = sparse_matmul(a, grad, adjoint_b=True, transpose_output=True)
+
+  return (grad_a, grad_b)
+
+
+@ops.RegisterGradient("SparseMatrixSparseMatMul")
+def _SparseMatrixSparseMatMulGrad(op: ops.Operation, grad):
+  """Gradient for sparse_matrix_sparse_mat_mul op."""
+  t_a = op.get_attr("transpose_a")
+  t_b = op.get_attr("transpose_b")
+  adj_a = op.get_attr("adjoint_a")
+  adj_b = op.get_attr("adjoint_b")
+  dtype = op.get_attr("type")
+
+  # input to sparse_matrix_sparse_mat_mul is (A, B) with CSR A and B.
+  # Output is CSR:
+  #   C = opA(A) . opB(B)
+  # where opA = transpose if transpose_a = True else identity
+  # and   opB = transpose if transpose_b = True else identity
+  a = op.inputs[0]
+  b = op.inputs[1]
+  conj = math_ops.conj
+  matmul = sparse_csr_matrix_ops.sparse_matrix_sparse_mat_mul
+  if not t_a and not t_b:
+    if not adj_a:
+      if not adj_b:
+        grad_a = matmul(grad, b, adjoint_b=True, type=dtype)
+        grad_b = matmul(a, grad, adjoint_a=True, type=dtype)
+      else:
+        grad_a = matmul(grad, b, type=dtype)
+        grad_b = matmul(grad, a, adjoint_a=True, type=dtype)
+    else:
+      if not adj_b:
+        grad_a = matmul(b, grad, adjoint_b=True, type=dtype)
+        grad_b = matmul(a, grad, type=dtype)
+      else:
+        grad_a = matmul(b, grad, adjoint_a=True, adjoint_b=True, type=dtype)
+        grad_b = matmul(grad, a, adjoint_a=True, adjoint_b=True, type=dtype)
+  elif not adj_a and not adj_b:
+    if not t_a and t_b:
+      grad_a = matmul(grad, conj(b), type=dtype)
+      grad_b = matmul(grad, conj(a), transpose_a=True, type=dtype)
+    elif t_a and not t_b:
+      grad_a = matmul(conj(b), grad, transpose_b=True, type=dtype)
+      grad_b = matmul(conj(a), grad, type=dtype)
+    else:
+      grad_a = matmul(b, grad, adjoint_a=True, transpose_b=True, type=dtype)
+      grad_b = matmul(grad, a, transpose_a=True, adjoint_b=True, type=dtype)
+  elif adj_a and t_b:
+    grad_a = matmul(b, grad, transpose_a=True, adjoint_b=True, type=dtype)
+    grad_b = matmul(grad, a, transpose_a=True, transpose_b=True, type=dtype)
+  elif t_a and adj_b:
+    grad_a = matmul(b, grad, transpose_a=True, transpose_b=True, type=dtype)
+    grad_b = matmul(grad, a, adjoint_a=True, transpose_b=True, type=dtype)
+
+  # TODO(tabakg): There should be a C++ function for sparse-sparse
+  # multiplication with pre-determined indices, instead of pruning after the
+  # multiplication.
+  return (_PruneCSRMatrix(grad_a, a), _PruneCSRMatrix(grad_b, b))
+
+
+@ops.RegisterGradient("SparseMatrixMul")
+def _SparseMatrixMulGrad(op: ops.Operation, grad):
+  """Gradient for sparse_matrix_mul op."""
+  # input to sparse_matrix_mul is (A, B) with CSR A and dense B.
+  # Output is CSR:
+  #   C = A .* B
+  del op
+  del grad
+  raise NotImplementedError

tensorflow/python/ops/linalg/sparse/sparse_csr_matrix_ops.py

@@ -0,0 +1,376 @@
+# 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 Operations."""
+
+import abc
+import collections
+
+# pylint: disable=g-direct-tensorflow-import, wildcard-import
+from tensorflow.python.eager import context
+from tensorflow.python.framework import cpp_shape_inference_pb2
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops.linalg.sparse import gen_sparse_csr_matrix_ops as sm_ops
+from tensorflow.python.ops.linalg.sparse.gen_sparse_csr_matrix_ops import *
+
+
+__all__ = [
+    "SparseMatrix",
+    "CSRSparseMatrix",
+    "matmul",
+    "dense_shape_and_type",
+]
+# pylint: disable=invalid-name
+__all__ += [_x for _x in dir(sm_ops) if not _x.startswith("_")]
+
+
+class DenseShapeAndType(
+    collections.namedtuple("DenseShapeAndType", ("shape", "dtype"))):
+  pass
+
+
+def _get_handle_data(tensor):
+  return resource_variable_ops.get_eager_safe_handle_data(tensor)
+
+
+def _create_handle_data_proto(shape_proto, dtype_enum):
+  """Create handle data based on shape and dtype protos."""
+  variant_shape_and_type_data = \
+    cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData()
+  variant_shape_and_type_data.is_set = True
+  # NOTE(ebrevdo): shape_and_type lacks append() in some versions of protobuf.
+  variant_shape_and_type_data.shape_and_type.extend([
+      cpp_shape_inference_pb2.CppShapeInferenceResult.HandleShapeAndType(
+          shape=shape_proto, dtype=dtype_enum)
+  ])
+  return variant_shape_and_type_data
+
+
+def _make_handle_data(tensor):
+  """Create handle data based on tensor shape and dtype."""
+  return _create_handle_data_proto(tensor.shape.as_proto(),
+                                   tensor.dtype.as_datatype_enum)
+
+
+def get_shape_and_type(matrix):
+  """Return matrix's shape and type if available."""
+  handle_data = getattr(matrix, "_handle_data", None)
+  if handle_data is None:
+    return None
+  if len(handle_data.shape_and_type) != 1:
+    raise ValueError(
+        "shape_and_type array in _handle_data must have length one, but saw: %d"
+        % len(handle_data.shape_and_type))
+  return handle_data.shape_and_type[0]
+
+
+def dense_shape_and_type(matrix):
+  """Get dense shape and dtype of the tf.Tensor containing the matrix.
+
+  Args:
+    matrix: A `tf.Tensor` of type `tf.variant` storing a sparse matrix.
+
+  Returns:
+    An instance of `ShapeAndType` with properties `shape` (a `tf.TensorShape`)
+    and `dtype` (a `tf.DType`).
+
+  Raises:
+    TypeError: if `matrix` is not a tensor or its dtype is not variant.
+    ValueError: if `matrix` lacks static handle data containing the dense
+      shape and dtype.
+  """
+  if not isinstance(matrix, tensor_lib.Tensor):
+    raise TypeError("matrix should be a tensor, but saw: %s" % (matrix,))
+  if matrix.dtype != dtypes.variant:
+    raise TypeError(
+        "expected matrix to be type tf.variant, but saw: %s" % (matrix.dtype,))
+  handle_data = _get_handle_data(matrix)
+  if not handle_data or not handle_data.is_set:
+    raise ValueError("matrix has missing handle data: %s" % (matrix,))
+  if len(handle_data.shape_and_type) != 1:
+    raise ValueError("len(matrix.handle_data.shape_and_type) != 1: '%s'" %
+                     (handle_data.shape_and_type,))
+  return DenseShapeAndType(
+      tensor_shape.TensorShape(handle_data.shape_and_type[0].shape),
+      dtypes.DType(handle_data.shape_and_type[0].dtype))
+
+
+def matmul_shape_inference(a, b, c, transpose_a, transpose_b, adjoint_a,
+                           adjoint_b):
+  """Helper function for matmul to set the result matrix's handle data."""
+  c_handle = getattr(c, "_handle_data", None)
+  a_shape_and_type = get_shape_and_type(a)
+  b_shape_and_type = get_shape_and_type(b)
+  if (c_handle is None and a_shape_and_type is not None and
+      b_shape_and_type is not None):
+
+    transpose_a = transpose_a or adjoint_a
+    transpose_b = transpose_b or adjoint_b
+
+    a_shape = a_shape_and_type.shape
+    b_shape = b_shape_and_type.shape
+    rank = len(a_shape.dim)
+
+    # Creates the output shape.
+    c_rows = a_shape.dim[rank - (1 if transpose_a else 2)].size
+    c_cols = b_shape.dim[rank - (2 if transpose_b else 1)].size
+    c_shape = tensor_shape.TensorShape(a_shape)
+    c_shape = tensor_shape.TensorShape(c_shape[:rank - 2] + [c_rows, c_cols])
+    c_handle = _create_handle_data_proto(c_shape.as_proto(),
+                                         a_shape_and_type.dtype)
+  return c_handle
+
+
+def matmul(a,
+           b,
+           transpose_a=False,
+           transpose_b=False,
+           adjoint_a=False,
+           adjoint_b=False,
+           name=None):
+  """Perform a sparse matrix matmul between `a` and `b`.
+
+  Performs a contraction between `a` and `b` along the two innermost dimensions.
+  If both `a` and `b` are instances of `SparseMatrix`, returns a new instance
+  of `SparseMatrix` (same type as `a`).  If one is not an instance of
+  `SparseMatrix`, returns a dense `Tensor`:
+
+  ```
+  c = opA(a) . opB(b)
+  ```
+  where `opA` (resp. `opB`) is the transpose or hermitian transpose depending
+  on the values of `transpose_a` (resp. `transpose_b`) and `adjoint_a`
+  (resp. `adjoint_b`).
+
+  Args:
+    a: `Tensor` or `SparseMatrix`, having rank `2` or `3`.
+    b: `Tensor` or `SparseMatrix`, having rank `2` or `3`.
+    transpose_a: Python `bool`.
+    transpose_b: Python `bool`.
+    adjoint_a: Python `bool`.
+    adjoint_b: Python `bool`.
+    name: Optional name to use when creating ops.
+
+  Returns:
+    A `SparseMatrix` if both `a` and `b` are instances of `SparseMatrix`,
+    otherwise a dense `Tensor`.
+  """
+  if not isinstance(a, SparseMatrix) and not isinstance(b, SparseMatrix):
+    return math_ops.matmul(
+        a,
+        b,
+        transpose_a=transpose_a,
+        transpose_b=transpose_b,
+        adjoint_a=adjoint_a,
+        adjoint_b=adjoint_b,
+        name=name)
+
+  # pylint: disable=protected-access
+  a_matrix = a._matrix if isinstance(a, SparseMatrix) else a
+  b_matrix = b._matrix if isinstance(b, SparseMatrix) else b
+  with ops.name_scope(name, "SparseMatrixMatMul", [a_matrix, b_matrix]):
+    if isinstance(a, SparseMatrix) and isinstance(b, SparseMatrix):
+      if not (isinstance(a, type(b)) or isinstance(b, type(a))):
+        raise TypeError("SparseMatrix types don't inherit from each other: "
+                        "%s and %s" % (type(a), type(b)))
+      c = sm_ops.sparse_matrix_sparse_mat_mul(
+          a_matrix,
+          b_matrix,
+          transpose_a=transpose_a,
+          transpose_b=transpose_b,
+          adjoint_a=adjoint_a,
+          adjoint_b=adjoint_b,
+          type=a.dtype)
+
+      # In eager mode, shape inference functions are not called, and the output
+      # shape is not set. We have to infer the output shape here.
+      # TODO(penporn): Set this from the C++ kernel instead.
+      c_handle = matmul_shape_inference(a_matrix, b_matrix, c, transpose_a,
+                                        transpose_b, adjoint_a, adjoint_b)
+      return a._from_matrix(c, handle_data=c_handle)
+
+    elif isinstance(a, SparseMatrix):
+      return sm_ops.sparse_matrix_mat_mul(
+          a_matrix,
+          b,
+          transpose_a=transpose_a,
+          transpose_b=transpose_b,
+          adjoint_a=adjoint_a,
+          adjoint_b=adjoint_b)
+    else:
+      # opA(A) . opB(B) = t(nopB(B) . nopA(A))
+      if not adjoint_a and not adjoint_b:
+        return sm_ops.sparse_matrix_mat_mul(
+            b_matrix,
+            a,
+            transpose_a=not transpose_b,
+            transpose_b=not transpose_a,
+            transpose_output=True)
+      elif not transpose_a and not transpose_b:
+        return sm_ops.sparse_matrix_mat_mul(
+            b_matrix,
+            a,
+            adjoint_a=not adjoint_b,
+            adjoint_b=not adjoint_a,
+            transpose_output=True,
+            conjugate_output=True)
+      else:
+        return sm_ops.sparse_matrix_mat_mul(
+            b_matrix,
+            math_ops.conj(a),
+            transpose_output=True,
+            conjugate_output=adjoint_b)
+
+
+class SparseMatrix(metaclass=abc.ABCMeta):
+  """Abstract class for sparse matrix types."""
+
+  @abc.abstractmethod
+  def __init__(self):
+    self._eager_mode = context.executing_eagerly()
+
+  @abc.abstractproperty
+  def _matrix(self):
+    pass
+
+  @abc.abstractmethod
+  def _from_matrix(self, matrix, handle_data=None):
+    pass
+
+  @abc.abstractmethod
+  def to_dense(self):
+    pass
+
+  @abc.abstractmethod
+  def to_sparse_tensor(self):
+    pass
+
+  @property
+  def graph(self):
+    return self._matrix.graph
+
+  @property
+  def shape(self):
+    return dense_shape_and_type(self._matrix).shape
+
+  @property
+  def dtype(self):
+    return dense_shape_and_type(self._matrix).dtype
+
+  @property
+  def eager_handle_data(self):
+    """Return the matrix's handle data iff in eager mode."""
+    return _get_handle_data(self._matrix) if self._eager_mode else None
+
+  def conj(self):
+    return self._from_matrix(
+        math_ops.conj(self._matrix), self.eager_handle_data)
+
+  def hermitian_transpose(self):
+    """Return the hermitian transpose of the matrix."""
+    return self._from_matrix(
+        sm_ops.sparse_matrix_transpose(
+            self._matrix, conjugate=True, type=self.dtype),
+        self.eager_handle_data)
+
+  def nnz(self):
+    """Number of stored values, including explicit zeros."""
+    return sm_ops.sparse_matrix_nnz(self._matrix)
+
+  nonzero = nnz
+
+  def sorted_indices(self):
+    # TODO(ebrevdo): A more efficient implementation?
+    return self.to_sparse_tensor().indices
+
+  def transpose(self):
+    return self._from_matrix(
+        sm_ops.sparse_matrix_transpose(self._matrix, type=self.dtype),
+        self.eager_handle_data)
+
+
+class CSRSparseMatrix(SparseMatrix):
+  """(Optionally batched) CSR Sparse Matrix."""
+
+  def __init__(self, value, indices=None, name=None):
+    """Construct a CSRSparseMatrix from a dense matrix or SparseTensor.
+
+    Args:
+      value: A dense `2D` or `3D` Tensor or `SparseTensor`.
+      indices: The nonzero indices of `value`
+        (if `value` is not a `SparseTensor`).
+      name: Optional op name.
+
+    Raises:
+      ValueError: if `value` is a `SparseTensor` and `indices` is not `None`.
+    """
+    del name  # Unused.
+    super(CSRSparseMatrix, self).__init__()
+    if isinstance(value, sparse_tensor.SparseTensor):
+      if indices is not None:
+        raise ValueError("indices must be None if value is a SparseTensor.")
+      self._dtype = value.dtype
+      self._csr_matrix = sm_ops.sparse_tensor_to_csr_sparse_matrix(
+          indices=value.indices,
+          values=value.values,
+          dense_shape=value.dense_shape)
+    else:
+      value = ops.convert_to_tensor(value)
+      self._dtype = value.dtype
+      if indices is not None:
+        indices = ops.convert_to_tensor(indices, dtype=dtypes.int64)
+      else:
+        indices = array_ops.stop_gradient(array_ops.where(value))
+      self._csr_matrix = sm_ops.dense_to_csr_sparse_matrix(value, indices)
+
+    # Eager mode doesn't call shape inference functions, so we have to set the
+    # shape and dtype handle data directly.
+    if self._eager_mode:
+      # pylint: disable=protected-access
+      self._csr_matrix._handle_data = _make_handle_data(value)
+      # pylint: enable=protected-access
+
+  @property
+  def _matrix(self):
+    return self._csr_matrix
+
+  def _from_matrix(self, matrix, handle_data=None):
+    assert (
+        isinstance(matrix, tensor_lib.Tensor) and matrix.dtype == dtypes.variant
+    )
+    ret = type(self).__new__(type(self))
+    # pylint: disable=protected-access
+    ret._dtype = self._dtype
+    if self._eager_mode:
+      if matrix._handle_data is None:
+        matrix._handle_data = handle_data
+      assert matrix._handle_data is not None
+    ret._csr_matrix = matrix
+    # pylint: enable=protected-access
+    return ret
+
+  def to_dense(self):
+    return sm_ops.csr_sparse_matrix_to_dense(self._matrix, type=self.dtype)
+
+  def to_sparse_tensor(self):
+    r = sm_ops.csr_sparse_matrix_to_sparse_tensor(self._matrix, type=self.dtype)
+    return sparse_tensor.SparseTensor(
+        indices=r.indices, values=r.values, dense_shape=r.dense_shape)