heat
heat copied to clipboard
ht.average() should allow weighted average across multiple axes
Feature functionality
ht.average(x, weights=None, axis=None, returned=False)
(PR #352) has the same constraints as numpy.average(), among them the following:
- if axis is not None, the weights tensor must be 1D and of the same size as x.gshape[axis].
This makes sense as long as axis is not a tuple. numpy.average() supports tuple axis, but as far as I understand, in this case np.average() only works if weights.shape = x.shape (this is not clear from the numpy documentation, it's what I gather from testing).
I see two possible solutions:
- enforce weights.shape = x.shape if axis is tuple (and make it clear in the docs)
- allow weights = tuple of ht.tensors of the appropriate size and in the appropriate order for the respective axis.
At the moment, ht.average() raises a NotImplementedError if axis is a tuple.
I would like to work on this.
@ClaudiaComito I have coded up the second solution and also made the small change required for solution 1. I have opened PRs for both. Let me know which one would be the preferred solution. 1 --> #953 2 --> #952