FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Implement gradients for SparseAdd. · codeplaysoftware/tensorflow@9033377 · GitHub

  1. codeplaysoftware
  2. tensorflow

Repository navigation

Commit 9033377

Browse files
concretevitamin
authored and
tensorflower-gardener
committed
Implement gradients for SparseAdd.
Change: 118216819
1 parent 7913f95 commit 9033377

10 files changed

Lines changed: 351 additions & 11 deletions

File tree

‎tensorflow/core/kernels/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1018,6 +1018,7 @@ tf_kernel_libraries(
10181018
tf_kernel_libraries(
10191019
name = "sparse",
10201020
prefixes = [
1021+
"sparse_add_grad_op",
10211022
"sparse_add_op",
10221023
"sparse_concat_op",
10231024
"sparse_reorder_op",

‎tensorflow/core/kernels/sparse_add_grad_op.cc‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
/* Copyright 2016 Google Inc. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#include "tensorflow/core/framework/op_kernel.h"
17+
#include "tensorflow/core/framework/register_types.h"
18+
#include "tensorflow/core/framework/tensor.h"
19+
#include "tensorflow/core/framework/tensor_util.h"
20+
#include "tensorflow/core/framework/types.h"
21+
#include "tensorflow/core/util/sparse/sparse_tensor.h"
22+
23+
namespace tensorflow {
24+
25+
template <typename T>
26+
class SparseAddGradOp : public OpKernel {
27+
public:
28+
explicit SparseAddGradOp(OpKernelConstruction *ctx) : OpKernel(ctx) {}
29+
30+
void Compute(OpKernelContext *ctx) override {
31+
// Gradient for op: SparseAdd(a, b) == sum.
32+
const Tensor *backprop_val_grad, *a_indices, *b_indices, *sum_indices;
33+
OP_REQUIRES_OK(ctx, ctx->input("backprop_val_grad", &backprop_val_grad));
34+
OP_REQUIRES_OK(ctx, ctx->input("a_indices", &a_indices));
35+
OP_REQUIRES_OK(ctx, ctx->input("b_indices", &b_indices));
36+
OP_REQUIRES_OK(ctx, ctx->input("sum_indices", &sum_indices));
37+
38+
OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(a_indices->shape()) &&
39+
TensorShapeUtils::IsMatrix(b_indices->shape()) &&
40+
TensorShapeUtils::IsMatrix(sum_indices->shape()),
41+
errors::InvalidArgument(
42+
"Input indices should be matrices but received shapes: ",
43+
a_indices->shape().DebugString(), " and ",
44+
b_indices->shape().DebugString(), " and ",
45+
sum_indices->shape().DebugString()));
46+
OP_REQUIRES(
47+
ctx, TensorShapeUtils::IsVector(backprop_val_grad->shape()),
48+
errors::InvalidArgument(
49+
"Input backprop_val_grad should be a vector but received shape: ",
50+
backprop_val_grad->shape().DebugString()));
51+
OP_REQUIRES(
52+
ctx, a_indices->dim_size(1) == b_indices->dim_size(1) &&
53+
b_indices->dim_size(1) == sum_indices->dim_size(1),
54+
errors::InvalidArgument("The densified operands should have the same "
55+
"ndims; for A, B, sum got: ",
56+
a_indices->dim_size(1), b_indices->dim_size(1),
57+
sum_indices->dim_size(1)));
58+
OP_REQUIRES(
59+
ctx, backprop_val_grad->NumElements() == sum_indices->dim_size(0),
60+
errors::InvalidArgument("# elements of backprop_val_grad and # rows of "
61+
"sum_indices should match (#nnz of sum): got ",
62+
backprop_val_grad->NumElements(), " and ",
63+
sum_indices->dim_size(0)));
64+
65+
const int num_dims = a_indices->dim_size(1);
66+
const int64 a_nnz = a_indices->dim_size(0);
67+
const int64 b_nnz = b_indices->dim_size(0);
68+
const int64 sum_nnz = backprop_val_grad->NumElements();
69+
70+
const auto a_indices_mat = a_indices->matrix<int64>();
71+
const auto b_indices_mat = b_indices->matrix<int64>();
72+
const auto sum_indices_mat = sum_indices->matrix<int64>();
73+
74+
Tensor *a_val_grad, *b_val_grad;
75+
OP_REQUIRES_OK(ctx,
76+
ctx->allocate_output(0, TensorShape({a_nnz}), &a_val_grad));
77+
OP_REQUIRES_OK(ctx,
78+
ctx->allocate_output(1, TensorShape({b_nnz}), &b_val_grad));
79+
80+
T *a_val_grad_flat = a_val_grad->flat<T>().data();
81+
T *b_val_grad_flat = b_val_grad->flat<T>().data();
82+
const T *backprop_val_grad_flat = backprop_val_grad->flat<T>().data();
83+
memset(a_val_grad_flat, 0, sizeof(T) * a_nnz);
84+
memset(b_val_grad_flat, 0, sizeof(T) * b_nnz);
85+
86+
#define COMPARE(a_or_b, idx) \
87+
switch (sparse::DimComparator::cmp(a_or_b##_indices_mat, sum_indices_mat, \
88+
idx, k, num_dims)) { \
89+
case 0: \
90+
a_or_b##_val_grad_flat[idx] = backprop_val_grad_flat[k]; \
91+
++idx; \
92+
break; \
93+
case -1: \
94+
++idx; \
95+
a_or_b##_idx_geq = false; \
96+
break; \
97+
case 1: \
98+
break; \
99+
}
100+
101+
// Set-intersect the indices; fill in grads for positions in the
102+
// intersection.
103+
int64 i = 0, j = 0, k = 0;
104+
bool a_idx_geq, b_idx_geq;
105+
while (i < a_nnz && j < b_nnz && k < sum_nnz) {
106+
a_idx_geq = b_idx_geq = true;
107+
COMPARE(a, i);
108+
COMPARE(b, j);
109+
// increment pointer into sum_indices iff both the current A, B indices >=
110+
// the current sum index.
111+
if (a_idx_geq && b_idx_geq) ++k;
112+
}
113+
114+
// at most one loop below will run
115+
while (i < a_nnz && k < sum_nnz) {
116+
a_idx_geq = true;
117+
COMPARE(a, i);
118+
if (a_idx_geq) ++k;
119+
}
120+
while (j < b_nnz && k < sum_nnz) {
121+
b_idx_geq = true;
122+
COMPARE(b, j);
123+
if (b_idx_geq) ++k;
124+
}
125+
#undef COMPARE
126+
}
127+
};
128+
129+
#define REGISTER_KERNELS(type) \
130+
REGISTER_KERNEL_BUILDER( \
131+
Name("SparseAddGrad").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
132+
SparseAddGradOp<type>)
133+
134+
// This op should work for any T that SparseAdd is registered with.
135+
REGISTER_KERNELS(float);
136+
REGISTER_KERNELS(double);
137+
REGISTER_KERNELS(int64);
138+
REGISTER_KERNELS(int32);
139+
REGISTER_KERNELS(int16);
140+
REGISTER_KERNELS(int8);
141+
REGISTER_KERNELS(complex64);
142+
REGISTER_KERNELS(complex128);
143+
#undef REGISTER_KERNELS
144+
} // namespace tensorflow

‎tensorflow/core/kernels/sparse_add_op.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ class SparseAddOp : public OpKernel {
9595
std::vector<std::pair<bool, int64>> entries_to_copy; // from_a?, idx
9696
entries_to_copy.reserve(a_nnz + b_nnz);
9797
std::vector<T> out_values;
98-
const int64 num_dims = a_shape->dim_size(0);
98+
const int num_dims = a_shape->dim_size(0);
9999

100100
// The input and output sparse tensors are assumed to be ordered along
101101
// increasing dimension number.

‎tensorflow/core/ops/compat/ops_history.v0.pbtxt‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14681,6 +14681,55 @@ op {
1468114681
}
1468214682
}
1468314683
}
14684+
op {
14685+
name: "SparseAddGrad"
14686+
input_arg {
14687+
name: "backprop_val_grad"
14688+
type_attr: "T"
14689+
}
14690+
input_arg {
14691+
name: "a_indices"
14692+
type: DT_INT64
14693+
}
14694+
input_arg {
14695+
name: "b_indices"
14696+
type: DT_INT64
14697+
}
14698+
input_arg {
14699+
name: "sum_indices"
14700+
type: DT_INT64
14701+
}
14702+
output_arg {
14703+
name: "a_val_grad"
14704+
type_attr: "T"
14705+
}
14706+
output_arg {
14707+
name: "b_val_grad"
14708+
type_attr: "T"
14709+
}
14710+
attr {
14711+
name: "T"
14712+
type: "type"
14713+
allowed_values {
14714+
list {
14715+
type: DT_FLOAT
14716+
type: DT_DOUBLE
14717+
type: DT_INT64
14718+
type: DT_INT32
14719+
type: DT_UINT8
14720+
type: DT_UINT16
14721+
type: DT_INT16
14722+
type: DT_INT8
14723+
type: DT_COMPLEX64
14724+
type: DT_COMPLEX128
14725+
type: DT_QINT8
14726+
type: DT_QUINT8
14727+
type: DT_QINT32
14728+
type: DT_HALF
14729+
}
14730+
}
14731+
}
14732+
}
1468414733
op {
1468514734
name: "SparseApplyAdadelta"
1468614735
input_arg {

‎tensorflow/core/ops/ops.pbtxt‎

Lines changed: 61 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9446,7 +9446,7 @@ op {
94469446
name: "SparseAdd"
94479447
input_arg {
94489448
name: "a_indices"
9449-
description: "2-D. The `indices` of the first `SparseTensor`, size `[nnz, 2]` Matrix."
9449+
description: "2-D. The `indices` of the first `SparseTensor`, size `[nnz, ndims]` Matrix."
94509450
type: DT_INT64
94519451
}
94529452
input_arg {
@@ -9456,12 +9456,12 @@ op {
94569456
}
94579457
input_arg {
94589458
name: "a_shape"
9459-
description: "1-D. The `shape` of the first `SparseTensor`, size `[2]` Vector."
9459+
description: "1-D. The `shape` of the first `SparseTensor`, size `[ndims]` Vector."
94609460
type: DT_INT64
94619461
}
94629462
input_arg {
94639463
name: "b_indices"
9464-
description: "2-D. The `indices` of the second `SparseTensor`, size `[nnz, 2]` Matrix."
9464+
description: "2-D. The `indices` of the second `SparseTensor`, size `[nnz, ndims]` Matrix."
94659465
type: DT_INT64
94669466
}
94679467
input_arg {
@@ -9471,7 +9471,7 @@ op {
94719471
}
94729472
input_arg {
94739473
name: "b_shape"
9474-
description: "1-D. The `shape` of the second `SparseTensor`, size `[2]` Vector."
9474+
description: "1-D. The `shape` of the second `SparseTensor`, size `[ndims]` Vector."
94759475
type: DT_INT64
94769476
}
94779477
input_arg {
@@ -9533,6 +9533,63 @@ op {
95339533
summary: "Adds two `SparseTensor` objects to produce another `SparseTensor`."
95349534
description: "The input `SparseTensor` objects\' indices are assumed ordered in standard\nlexicographic order. If this is not the case, before this step run\n`SparseReorder` to restore index ordering.\n\nBy default, if two values sum to zero at some index, the output `SparseTensor`\nwould still include that particular location in its index, storing a zero in the\ncorresponding value slot. To override this, callers can specify `thresh`,\nindicating that if the sum has a magnitude strictly smaller than `thresh`, its\ncorresponding value and index would then not be included. In particular,\n`thresh == 0` (default) means everything is kept and actual thresholding happens\nonly for a positive value.\n\nIn the following shapes, `nnz` is the count after taking `thresh` into account."
95359535
}
9536+
op {
9537+
name: "SparseAddGrad"
9538+
input_arg {
9539+
name: "backprop_val_grad"
9540+
description: "1-D with shape `[nnz(sum)]`. The gradient with respect to\nthe non-empty values of the sum."
9541+
type_attr: "T"
9542+
}
9543+
input_arg {
9544+
name: "a_indices"
9545+
description: "2-D. The `indices` of the `SparseTensor` A, size `[nnz(A), ndims]`."
9546+
type: DT_INT64
9547+
}
9548+
input_arg {
9549+
name: "b_indices"
9550+
description: "2-D. The `indices` of the `SparseTensor` B, size `[nnz(B), ndims]`."
9551+
type: DT_INT64
9552+
}
9553+
input_arg {
9554+
name: "sum_indices"
9555+
description: "2-D. The `indices` of the sum `SparseTensor`, size\n`[nnz(sum), ndims]`."
9556+
type: DT_INT64
9557+
}
9558+
output_arg {
9559+
name: "a_val_grad"
9560+
description: "1-D with shape `[nnz(A)]`. The gradient with respect to the\nnon-empty values of A."
9561+
type_attr: "T"
9562+
}
9563+
output_arg {
9564+
name: "b_val_grad"
9565+
description: "1-D with shape `[nnz(B)]`. The gradient with respect to the\nnon-empty values of B."
9566+
type_attr: "T"
9567+
}
9568+
attr {
9569+
name: "T"
9570+
type: "type"
9571+
allowed_values {
9572+
list {
9573+
type: DT_FLOAT
9574+
type: DT_DOUBLE
9575+
type: DT_INT64
9576+
type: DT_INT32
9577+
type: DT_UINT8
9578+
type: DT_UINT16
9579+
type: DT_INT16
9580+
type: DT_INT8
9581+
type: DT_COMPLEX64
9582+
type: DT_COMPLEX128
9583+
type: DT_QINT8
9584+
type: DT_QUINT8
9585+
type: DT_QINT32
9586+
type: DT_HALF
9587+
}
9588+
}
9589+
}
9590+
summary: "The gradient operator for the SparseAdd op."
9591+
description: "The SparseAdd op calculates A + B, where A, B, and the sum are all represented\nas `SparseTensor` objects. This op takes in the upstream gradient w.r.t.\nnon-empty values of the sum, and outputs the gradients w.r.t. the non-empty\nvalues of A and B."
9592+
}
95369593
op {
95379594
name: "SparseApplyAdadelta"
95389595
input_arg {

‎tensorflow/core/ops/sparse_ops.cc‎

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,34 @@ limitations under the License.
1717

1818
namespace tensorflow {
1919

20+
REGISTER_OP("SparseAddGrad")
21+
.Input("backprop_val_grad: T")
22+
.Input("a_indices: int64")
23+
.Input("b_indices: int64")
24+
.Input("sum_indices: int64")
25+
.Output("a_val_grad: T")
26+
.Output("b_val_grad: T")
27+
.Attr("T: numbertype")
28+
.Doc(R"doc(
29+
The gradient operator for the SparseAdd op.
30+
31+
The SparseAdd op calculates A + B, where A, B, and the sum are all represented
32+
as `SparseTensor` objects. This op takes in the upstream gradient w.r.t.
33+
non-empty values of the sum, and outputs the gradients w.r.t. the non-empty
34+
values of A and B.
35+
36+
backprop_val_grad: 1-D with shape `[nnz(sum)]`. The gradient with respect to
37+
the non-empty values of the sum.
38+
a_indices: 2-D. The `indices` of the `SparseTensor` A, size `[nnz(A), ndims]`.
39+
b_indices: 2-D. The `indices` of the `SparseTensor` B, size `[nnz(B), ndims]`.
40+
sum_indices: 2-D. The `indices` of the sum `SparseTensor`, size
41+
`[nnz(sum), ndims]`.
42+
a_val_grad: 1-D with shape `[nnz(A)]`. The gradient with respect to the
43+
non-empty values of A.
44+
b_val_grad: 1-D with shape `[nnz(B)]`. The gradient with respect to the
45+
non-empty values of B.
46+
)doc");
47+
2048
REGISTER_OP("SparseAdd")
2149
.Input("a_indices: int64")
2250
.Input("a_values: T")
@@ -47,12 +75,12 @@ only for a positive value.
4775
4876
In the following shapes, `nnz` is the count after taking `thresh` into account.
4977
50-
a_indices: 2-D. The `indices` of the first `SparseTensor`, size `[nnz, 2]` Matrix.
78+
a_indices: 2-D. The `indices` of the first `SparseTensor`, size `[nnz, ndims]` Matrix.
5179
a_values: 1-D. The `values` of the first `SparseTensor`, size `[nnz]` Vector.
52-
a_shape: 1-D. The `shape` of the first `SparseTensor`, size `[2]` Vector.
53-
b_indices: 2-D. The `indices` of the second `SparseTensor`, size `[nnz, 2]` Matrix.
80+
a_shape: 1-D. The `shape` of the first `SparseTensor`, size `[ndims]` Vector.
81+
b_indices: 2-D. The `indices` of the second `SparseTensor`, size `[nnz, ndims]` Matrix.
5482
b_values: 1-D. The `values` of the second `SparseTensor`, size `[nnz]` Vector.
55-
b_shape: 1-D. The `shape` of the second `SparseTensor`, size `[2]` Vector.
83+
b_shape: 1-D. The `shape` of the second `SparseTensor`, size `[ndims]` Vector.
5684
thresh: 0-D. The magnitude threshold that determines if an output value/index
5785
pair takes space.
5886
)doc");

‎tensorflow/python/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -704,6 +704,7 @@ tf_gen_op_wrapper_py(
704704
"SerializeManySparse",
705705
"SerializeSparse",
706706
"SparseAdd",
707+
"SparseAddGrad",
707708
"SparseConcat",
708709
"SparseSplit",
709710
"SparseSelectLastK",

0 commit comments

Comments
 (0)
Back | FazBrowse Home | New Git URL