Equivalent for np.add.at in tensorflow

Question

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

Boyhood

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. – Alie

Search SO for '[numpy] add.at' to see how add.at has been used solve various numpy problems. – Alie

Do you know if there are duplicate indices in self.x? That's when at makes a difference. – Alie

I do understand what the code does. I am looking for an tensorflow alternative. – Boyhood

If anybody, who is open to do it in pytorch, comes here looking for a solution, check this out: https://mcmap.net/q/1472906/-add-a-index-selected-tensor-to-another-tensor-with-overlapping-indices-in-pytorch – Bazooka

Answer 1

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].

Pickel

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). – Boyhood

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()? – Pickel

Actually the shape of x is different but it's size is N. – Boyhood

You could do tf.reshape(x, [-1]) to get a shape [N] tensor then. – Pickel

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. – Alie

Thank you! Sorry for the late response. – Boyhood

Tensorflow now has tf.tensor_scatter_nd_add - I recommend using this instead of sparse_add. Anecdotally, it seems to be 10x faster. – Sarcomatosis

Answer 2

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.

Alie

Answer 3

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())

Sarcomatosis

Actually I just discovered tf.math.segment_sum 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. – Sarcomatosis

Answer 4

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())

Sarcomatosis