How do I convert a np.add.at statement into tensorflow?
np.add.at(dW, self.x.ravel(), dout.reshape(-1, self.D))
Edit
self.dW.shape is (V, D), self.D.shape is (N, D) and self.x.size is N
Asked
Active
Viewed 995 times
4
Amey Agrawal
- 132
- 9
-
Expand on your question. Are you trying to understand what the `add.at` does? Or trying get something in `tensorflow` that does the same thing and at the same speed? `add.at` is use to speed up iterative problems where the standard buffer `add` produces the wrong result. – hpaulj Nov 02 '16 at 21:26
-
Search `SO` for '[numpy] add.at' to see how `add.at` has been used solve various `numpy` problems. – hpaulj Nov 02 '16 at 21:35
-
Do you know if there are duplicate indices in `self.x`? That's when `at` makes a difference. – hpaulj Nov 02 '16 at 22:25
-
I do understand what the code does. I am looking for an tensorflow alternative. – Amey Agrawal Nov 03 '16 at 00:54
-
If anybody, who is open to do it in pytorch, comes here looking for a solution, check this out: https://stackoverflow.com/a/65584479/3337089 – Nagabhushan S N Jan 05 '21 at 18:36
4 Answers
2
For np.add.at, you probably want to look at tf.SparseTensor, which represents a tensor by a list of values and a list of indices (which is more suitable for sparse data, hence the name).
So for your example:
np.add.at(dW, self.x.ravel(), dout.reshape(-1, self.D))
that would be (assuming dW, x and dout are tensors):
tf.sparse_add(dW, tf.SparseTensor(x, tf.reshape(dout, [-1])))
This is assuming x is of shape [n, nDims] (i.e. x is a 'list' of n indices, each of dimension nDims), and dout has shape [n].
Lars Mennen
- 544
- 4
- 4
-
Actually in the original numpy code, dout is of shape (N, D). self.x.reshape(-1) gives array of shape (N, ). While dW has shape (V, D). – Amey Agrawal Nov 03 '16 at 00:52
-
What is the shape of self.x? (before reshape) Your edit mentions it is of size N, but in case it is of shape N, why do you do .ravel()? – Lars Mennen Nov 03 '16 at 08:25
-
-
-
If this sparse tensor is modeled on the scipy csr sparse matrix, duplicate indices are summed. That's the kind of behavior that the unbuffered `add.at` is supposed to produce. – hpaulj Nov 04 '16 at 08:15
-
-
Tensorflow now has `tf.tensor_scatter_nd_add` - I recommend using this instead of `sparse_add`. Anecdotally, it seems to be 10x faster. – Peter Aug 16 '22 at 18:17
0
Here's an example of what np.add.at does:
In [324]: a=np.ones((10,))
In [325]: x=np.array([1,2,3,1,4,5])
In [326]: b=np.array([1,1,1,1,1,1])
In [327]: np.add.at(a,x,b)
In [328]: a
Out[328]: array([ 1., 3., 2., 2., 2., 2., 1., 1., 1., 1.])
If instead I use +=
In [331]: a1=np.ones((10,))
In [332]: a1[x]+=b
In [333]: a1
Out[333]: array([ 1., 2., 2., 2., 2., 2., 1., 1., 1., 1.])
note that a1[1] is 2, not 3.
If instead I use an iterative solution
In [334]: a2=np.ones((10,))
In [335]: for i,j in zip(x,b):
...: a2[i]+=j
...:
In [336]: a2
Out[336]: array([ 1., 3., 2., 2., 2., 2., 1., 1., 1., 1.])
it matches.
If x does not have duplicates then += works just fine. But with the duplicates, the add.at is required to match the iterative solution.
hpaulj
- 221,503
- 14
- 230
- 353
0
Edit Do not use this one, use tf.tensor_scatter_nd_add instead, as shown in this answer. I'll leave this here for posterity.
Here's Lars's answer, updated to match current TensorFlow api (2.8.1), made to handle multi-dimensional case, and turned into a complete function:
def tf_add_at(
values: tf.Tensor, # A length-N vector of values
indices: tf.Tensor, # A shape (N, ) or (N, D) vector of indices (where D is dimension of redsult)
result_shape: Optional[Tuple[int, ...]] = None # Shape of result array (e.g. (D, ), or (D1, D2)). If null, take smallest that fits ixs.
) -> tf.Tensor: # Resulting tensor, which will have result_shape, contain data in vals.
""" Add the values, grouping by indices, into a tensor with the given shape """
indices = tf.cast(tf.reshape(indices, (indices.shape[0], -1)), dtype=tf.int64) # Make sure ixs hav shape (n_indices, n_dims) - this line handles the 1d case
if result_shape is None:
result_shape = tf.reduce_max(indices, axis=0) + 1
return tf.sparse.reduce_sum(SparseTensor(
indices=tf.concat([indices, tf.range(len(indices), dtype=tf.int64)[:, None]], axis=1),
values=values,
dense_shape=tf.concat([result_shape, [len(indices)]], axis=0)
), axis=-1)
Which passes test:
def test_tf_add_at():
# 1d case
vals = tf.constant([2, 5, 7, 2, 0, 8, 3, 5])
ixs = tf.constant([0, 2, 2, 2, 0, 4, 4, 3])
desired = tf.constant([2+0, 0, 5+7+2, 5, 8+3])
result = tf_add_at(vals=vals, ixs=ixs)
assert np.array_equal(result.numpy(), desired.numpy())
# 2d case
vals = tf.constant([2, 5, 7, 2, 0, 8, 3, 5])
ixs = tf.constant([(0, 0), (0, 2), (0, 2), (0, 2), (0, 0), (1, 1), (1, 1), (1, 0)])
result = tf_add_at(vals=vals, ixs=ixs)
desired = tf.constant([[2+0, 0, 5+7+2], [5, 8+3, 0]])
assert np.array_equal(result.numpy(), desired.numpy())
Peter
- 12,274
- 9
- 71
- 86
-
Actually I just discovered `tf.math.segment_sum` https://www.tensorflow.org/api_docs/python/tf/math/segment_sum - which does this already, but only for the 1d case, and it expects indices to be sorted. – Peter Aug 04 '22 at 17:33
0
Use tf.tensor_scatter_nd_add.
Example use in test:
def test_tf_scatter_add():
# 1d case
vals = tf.constant([2, 5, 7, 2, 0, 8, 3, 5])
ixs = tf.constant([0, 2, 2, 2, 0, 4, 4, 3])
desired = tf.constant([2 + 0, 0, 5 + 7 + 2, 5, 8 + 3])
result = tf.tensor_scatter_nd_add(tensor=tf.zeros(5, dtype=vals.dtype), indices=tf.reshape(ixs, (-1, 1)), updates=vals)
assert np.array_equal(result.numpy(), desired.numpy())
# 2d case
vals = tf.constant([2, 5, 7, 2, 0, 8, 3, 5])
ixs = tf.constant([(0, 0), (0, 2), (0, 2), (0, 2), (0, 0), (1, 1), (1, 1), (1, 0)])
result = tf.tensor_scatter_nd_add(tensor=tf.zeros((2, 3), dtype=vals.dtype), indices=ixs, updates=vals)
desired = tf.constant([[2 + 0, 0, 5 + 7 + 2], [5, 8 + 3, 0]])
assert np.array_equal(result.numpy(), desired.numpy())
Peter
- 12,274
- 9
- 71
- 86