Gradient for SparseMatMul.
(op, grad)
| 1598 | |
| 1599 | @ops.RegisterGradient("SparseMatMul") |
| 1600 | def _SparseMatMulGrad(op, grad): |
| 1601 | """Gradient for SparseMatMul.""" |
| 1602 | |
| 1603 | t_a = op.get_attr("transpose_a") |
| 1604 | t_b = op.get_attr("transpose_b") |
| 1605 | is_sparse = object_identity.ObjectIdentityDictionary() |
| 1606 | is_sparse[op.inputs[0]] = op.get_attr("a_is_sparse") |
| 1607 | is_sparse[op.inputs[1]] = op.get_attr("b_is_sparse") |
| 1608 | # Use heuristic to figure out if grad might be sparse |
| 1609 | is_sparse[grad] = not context.executing_eagerly() and ( |
| 1610 | grad.op.type == "ReluGrad") |
| 1611 | |
| 1612 | def _SparseMatMul(t1, t2, out_dtype, transpose_a=False, transpose_b=False): |
| 1613 | """Helper function to create SparseMatMul op.""" |
| 1614 | |
| 1615 | assert t1 in is_sparse and t2 in is_sparse |
| 1616 | t1_sparse = is_sparse[t1] |
| 1617 | t2_sparse = is_sparse[t2] |
| 1618 | if transpose_b: |
| 1619 | t2 = array_ops.transpose(t2) |
| 1620 | transpose_b = False |
| 1621 | prod = math_ops.matmul( |
| 1622 | t1, |
| 1623 | t2, |
| 1624 | transpose_a=transpose_a, |
| 1625 | transpose_b=transpose_b, |
| 1626 | a_is_sparse=t1_sparse, |
| 1627 | b_is_sparse=t2_sparse) |
| 1628 | if prod.dtype != out_dtype: |
| 1629 | prod = math_ops.cast(prod, out_dtype) |
| 1630 | return prod |
| 1631 | |
| 1632 | dtype_a = op.inputs[0].dtype |
| 1633 | dtype_b = op.inputs[1].dtype |
| 1634 | if not t_a and not t_b: |
| 1635 | return (_SparseMatMul(grad, op.inputs[1], dtype_a, transpose_b=True), |
| 1636 | _SparseMatMul(op.inputs[0], grad, dtype_b, transpose_a=True)) |
| 1637 | elif not t_a and t_b: |
| 1638 | return (_SparseMatMul(grad, op.inputs[1], dtype_a), |
| 1639 | _SparseMatMul(grad, op.inputs[0], dtype_b, transpose_a=True)) |
| 1640 | elif t_a and not t_b: |
| 1641 | return (_SparseMatMul(op.inputs[1], grad, dtype_a, transpose_b=True), |
| 1642 | _SparseMatMul(op.inputs[0], grad, dtype_b)) |
| 1643 | elif t_a and t_b: |
| 1644 | return (_SparseMatMul( |
| 1645 | op.inputs[1], grad, dtype_a, transpose_a=True, transpose_b=True), |
| 1646 | _SparseMatMul( |
| 1647 | grad, op.inputs[0], dtype_b, transpose_a=True, |
| 1648 | transpose_b=True)) |
| 1649 | |
| 1650 | |
| 1651 | @ops.RegisterGradient("Floor") |
nothing calls this directly
no test coverage detected