Add tf.SparseTensor and tf.IndexedSlices support to a number of T…
#18633
Conversation
Codecov ReportAttention:
Additional details and impacted files@@ Coverage Diff @@
## master #18633 +/- ##
==========================================
+ Coverage 78.48% 78.65% +0.16%
==========================================
Files 335 336 +1
Lines 33090 33369 +279
Branches 6486 6529 +43
==========================================
+ Hits 25972 26246 +274
- Misses 5546 5550 +4
- Partials 1572 1573 +1
Flags with carried forward coverage won't be shown. Click here to find out more.
☔ View full report in Codecov by Sentry. |
keras/backend/tensorflow/numpy.py
Outdated
| from keras.backend.tensorflow.core import convert_to_tensor | ||
|
|
||
|
|
||
| @sparse.element_wise_binary_on_union(tf.sparse.add) |
There was a problem hiding this comment.
Let's apply the following replacements:
element_wise -> elementwise (everywhere)
on_union -> just union (to be consistent with intersection)
There was a problem hiding this comment.
Question: is this new system really better than tf.sparse.add? Does it have performance implications? The previous "if + use tf.sparse" setup was very readable and easy to maintain. The new system is fairly high in abstraction and complexity. Are the benefits commensurate?
There was a problem hiding this comment.
Will do the replacements
is this new system really better than
tf.sparse.add?
So there are 2 times 3 cases: SparseTensor / IndexedSlices and union (add, maximum...), intersection (multiply), division (divide, true_divide, mod, floor_divide). tf.sparse happens to mostly support the "union" operators out of the box, so it was very little code. But it didn't support the other use cases and IndexedSlices is not supported for any operation. There is a lot of code for the 5 other use cases (for instance, this CL moved the long chunk of code to get multiply working for sparse from numpy.py to sparse.py).
One thing we could do differently is have the elementwise_binary_union decorator only add IndexedSlices support and not SparseTensors to have the tf.sparse.add code inline. I initially had it this way, but then I decided to make it look consistent.
Does it have performance implications?
For tf.sparse.add specifically, the only performance impact is the call to the decorator, which we need anyway for the IndexedSlices support. It does use tf.sparse.add directly, this is not re-implemented.
The new system is fairly high in abstraction and complexity.
I agree. After all was done to handle all the possible combinations, it was a lot more complex than initially expected.
Are the benefits commensurate?
Well the benefit is that with one decorator, you make an op SparseTensor and IndexedSlices compatible. And it's self documenting, you know which ops support SparseTensor and IndexedSlices (we could generate documentation based on the decorator). I could have decorated more ops actually (e.g. logical operations).
I want to look at JAX and see what approach makes sense for JAX. Maybe it will be similar, maybe it will make me rethink this whole thing.
There was a problem hiding this comment.
Thanks for the clarifications, this all sounds reasonable. Please let me know what you think about the JAX situation, and then we can finalize the design.
There was a problem hiding this comment.
Thanks!
One other thing is that implementing support in the ops directly adds complexity. For instance, in practice, this union / intersection business is never needed for IndexedSlices. In optimizers the only times IndexedSlices are combined, they always have the same indices. But in the ops, we can't just make this assumption.
|
To note, |
Interesting. Some of the |
hertschuh
force-pushed
the
indexed_slices_ops
branch
4 times, most recently
from
ab111e2 to
68c42aa
Compare
Looking into this a bit more, I found a way to check that the indices are the same at trace time. So, in the case of optimizers for instance, the |
…ensorFlow ops. The following element-wise unary ops now support `tf.SparseTensor` and `tf.IndexedSlices`. The output is of the same type as the input. - abs - absolute - arcsin - arcsinh - arctan - arctanh - ceil - conj - conjugate - copy - expm1 - floor - imag - log1p - negative - real - round - sign - sin - sinh - sqrt - square - tan - tanh The following element-wise unary ops now support `tf.SparseTensor` and `tf.IndexedSlices`. The output is dense. - arccos - arccosh - cos - cosh - exp - log - log10 - log2 - reciprocal The following element-wise binary ops now support `tf.SparseTensor` and `tf.IndexedSlices`. The output type depends on the two inputs and the op. - add - subtract - maximum - minimum - multiply - mod - divide - true_divide - floor_divide The following reduction op now supports `tf.IndexedSlices`. The output is an `tf.IndexedSlices` unless dimension 0 is reduced or the rank of the output is 1 or less. - mean This is in preparation for supporting sparse gradients in optimizers.
hertschuh
force-pushed
the
indexed_slices_ops
branch
from
68c42aa to
64ef001
Compare
google-ml-butler
bot
added
kokoro:force-run
ready to pull
google-ml-butler
bot
removed
ready to pull
…ensorFlow ops.
The following element-wise unary ops now support
tf.SparseTensorandtf.IndexedSlices. The output is of the same type as the input.The following element-wise unary ops now support
tf.SparseTensorandtf.IndexedSlices. The output is dense.The following element-wise binary ops now support
tf.SparseTensorandtf.IndexedSlices. The output type depends on the two inputs and the op.tf.SparseTensor)tf.SparseTensor)tf.SparseTensor)tf.SparseTensor)tf.SparseTensor)The following reduction op now supports
tf.IndexedSlices. The output is antf.IndexedSlicesunless dimension 0 is reduced or the rank of the output is 1 or less.This is in preparation for supporting sparse gradients in optimizers.