mlx - implement segment_sum and segment_max

[codecov-commenter]

Codecov Report

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

:exclamation: 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 :warning:

Additional details and impacted files

```diff @@ Coverage Diff @@ ## mlx #19652 +/- ## ====================================== Coverage ? 68.43% ====================================== Files ? 506 Lines ? 45959 Branches ? 8496 ====================================== Hits ? 31451 Misses ? 12857 Partials ? 1651 ``` | [Flag](https://app.codecov.io/gh/keras-team/keras/pull/19652/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | Coverage Δ | | |---|---|---| | [keras](https://app.codecov.io/gh/keras-team/keras/pull/19652/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | `68.35% <0.00%> (?)` | | | [keras-jax](https://app.codecov.io/gh/keras-team/keras/pull/19652/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | `58.82% <0.00%> (?)` | | | [keras-numpy](https://app.codecov.io/gh/keras-team/keras/pull/19652/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | `53.16% <0.00%> (?)` | | | [keras-tensorflow](https://app.codecov.io/gh/keras-team/keras/pull/19652/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | `59.97% <0.00%> (?)` | | Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team#carryforward-flags-in-the-pull-request-comment) to find out more.

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it 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. https://github.com/keras-team/keras/pull/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.