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

Asked By: Amey Agrawal

||

Source

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

Answered By: Lars Mennen

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.

Answered By: hpaulj

Answer 3

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

Answered By: Peter

Equivalent for np.add.at in tensorflow

Question:

Answers: