mlx - implement segment_sum and segment_max

[lkarthee]

No description provided.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 0% with 22 lines in your changes are missing coverage. Please review.

❗ No coverage uploaded for pull request base (mlx@4c90dfb). Click here to learn what that means.

Files	Patch %	Lines
keras/src/backend/mlx/math.py	0.00%	22 Missing ⚠️

Additional details and impacted files

@@          Coverage Diff           @@
##             mlx   #19652   +/-   ##
======================================
  Coverage       ?   68.43%           
======================================
  Files          ?      506           
  Lines          ?    45959           
  Branches       ?     8496           
======================================
  Hits           ?    31451           
  Misses         ?    12857           
  Partials       ?     1651           

Flag	Coverage Δ
keras	68.35% <0.00%> (?)
keras-jax	58.82% <0.00%> (?)
keras-numpy	53.16% <0.00%> (?)
keras-tensorflow	59.97% <0.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[lkarthee]

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.

[lkarthee]

One way to solve this is using numpy.unique on segment_ids to determine num_segments. Let me know what to do.

[fchollet]

@angeloskath could you please advise here?

[Faisal-Alsrheed]

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.
#19619

I suggest creating a roadmap where we start with fundamental functions and then move up.

best

[lkarthee]

@Faisal-Alsrheed My thoughts regarding missing functions which can't be added due to design decisions/limitations of mlx:

• we have to check if there is a workaround.
• if there is no workaround, may be we have to fallback to using numpy or jax adding a TODO?

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.