mlx - implement segment_sum and segment_max #19652
Conversation
lkarthee
changed the title
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## mlx #19652 +/- ##
======================================
Coverage ? 68.43%
======================================
Files ? 506
Lines ? 45959
Branches ? 8496
======================================
Hits ? 31451
Misses ? 12857
Partials ? 1651
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. |
| @@ -6,11 +6,37 @@ | |||
|
|
|||
|
|
|||
| def segment_sum(data, segment_ids, num_segments=None, sorted=False): | |||
| raise NotImplementedError("segment_sum is not implemented for mlx") | |||
| data = convert_to_tensor(data) | |||
There was a problem hiding this comment.
mlx does not have unique and boolean indices ml-explore/mlx#246 and ml-explore/mlx#377 . So the implementation is buggy with varying values of num_segments and negative segment_ids.
There was a problem hiding this comment.
One way to solve this is using numpy.unique on segment_ids to determine num_segments. Let me know what to do.
There was a problem hiding this comment.
@angeloskath could you please advise here?
There was a problem hiding this comment.
Hi @angeloskath Any update on above comment from @fchollet ? Please. Thank you!
|
Hi @lkarthee, I have noticed the same issue: there are missing and needed functions in MLX. What is the best strategy here? Should we use NumPy temporarily and add a TODO comment? I have decided to start with core files and functions. I suggest creating a roadmap where we start with fundamental functions and then move up. best |
|
@Faisal-Alsrheed My thoughts regarding missing functions which can't be added due to design decisions/limitations of mlx:
Regarding fixing core first, I have been doing that in my PRs. I tried to implement Pooling and CNN related funcs, realised many tests are failing and started fixing test cases. @fchollet any thoughts on what to use as fallback - numpy or jax. |
|
Hi @fchollet Any update on this PR? Please. Thank you! |
Sorry, just checking this now. I would say that JAX is preferable as fallback since numpy would not be sufficiently performant. |
|
Hi @lkarthee Any update on this PR? Please. Thank you! |
|
This PR is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you. |
|
This PR was closed because it has been inactive for 28 days. Please reopen if you'd like to work on this further. |
No description provided.