Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

segment_sum primitives / advanced indexing in jax.ops.index_add

See original GitHub issue

Would it be useful to others (aside from me) to have better support of sorted and unsorted segment_sums?

https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_sum https://www.tensorflow.org/api_docs/python/tf/math/segment_sum

In numpy one way to do unsorted segment sums is to sort then call np.add.reduceat, but this doesn’t seem to be in jax or autograd:

>>> import jax.numpy as np
np.add
>>> np.add
<function _one_to_one_binop.<locals>.<lambda> at 0x1430cdbf8>
>>> np.add.reduceat
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'function' object has no attribute 'reduceat'

Issue Analytics

State:
Created 4 years ago
Reactions:2
Comments:5 (5 by maintainers)

Top GitHub Comments

3reactions

hawkinspcommented, Apr 30, 2019

Yup, that’s the idea!

1reaction

hawkinspcommented, Apr 30, 2019

Ah, I guess it does not yet because it does not support advanced indexing. We should fix that.

Top Results From Across the Web

jax.ops.segment_sum - JAX documentation

GatherScatterMode value describing how out-of-bounds indices should be handled. By default, values outside of the range [0, num_segments) are dropped and do not ......