feat: Implement BatchWeightedSum and BatchWeightedMean#14
Merged
Conversation
… and merge support This commit introduces `BatchWeightedSum` and `BatchWeightedMean` classes to compute weighted statistics over batches of data. This final version incorporates several rounds of feedback to improve robustness and design: - **Core Implementation**: Added `BatchWeightedSum` and `BatchWeightedMean` classes. - **Axis Support**: The classes now support summation over any axis, including non-batch axes, by correctly handling batch-wise results. - **Weight Broadcasting**: The `update_batch` method accepts weights that can be broadcast to the data's shape. - **Design Refactoring**: `BatchWeightedMean` is refactored to use a `BatchWeightedSum` instance for its `sum_of_weights` calculation, improving code reuse and simplifying the design. - **Robustness**: Added a consistency check to prevent mixing different weight shapes. Division-by-zero in `BatchWeightedMean` is handled gracefully. - **Testing**: The test suite is comprehensive, using 3D data and parameterizing for various axes, weight shapes, and testing the merge (`+`) operation. - **Documentation**: All new classes are included in the documentation.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This change adds
BatchWeightedSumandBatchWeightedMeanclasses to thebatchstatslibrary, allowing for batch-wise calculation of weighted sum and mean. Theupdate_batchmethods are updated to accept weights. The implementation includes tests and documentation.