Hist_index or scatter and reduce, but not histc or scatter_reduce

1

I’m looking to express the following op in a nice manner:

# compute index[i, ...] for all data[i, ...]
hist_value[i, index[i, ...]] += data[i, ...]

This is almost like the scatter_reduce_ with add operation, but that would require:

  • additional memory allocation for hist_value AND index to be of the same shape as data
  • additional reduction (sum) over the dimensions we had to add

It’s almost like histc, but not quite, since

  • indices are known
  • Can’t use the default algorithm for bucket calculation since the bucket index computation is slightly tricky

I know the min/max/number of buckets so an API similar to histc might be the best fit here (but I’m not sure)

I don’t know if this is doable via Einstein notation. What do you people think?