[ronaldmannak]

This PR implements #134 except for the null-stream changes that I believe need to be generated.

Added

• mlx_distributed_clear_backends() binds clear_backends().
• mlx_distributed_init_with_all_gather(...) binds init(strict, bk, factory)
  for JACCL, with the factory flattened into one callback plus payload and
  dtor, as in mlx_closure_new_func_payload. A non-zero return from the
  callback fails the init. dtor runs exactly once, including on failure.
• mlx_distributed_group_get_rank() and mlx_distributed_group_get_size()
  return a status and write the value to res.

Note that ml-explore/mlx#4558 changes one case where mlx_distributed_clear_backends() is needed. On 0.32.3, a non-strict init also caches its size-one fallback, so a later init returns that fallback until the backends are cleared. From the next MLX release the fallback is no longer cached. Clearing is still needed to drop a backend that did initialize.

Not included: null stream in the mlx_distributed_* ops

The C++ backend automatically chooses the communication stream (CPU for Ring, MPI, and JACCL, and GPU for NCCL) when the stream parameter is unspecified. Currently, mlx-c requires an explicit stream. We kind of lessened the impact by making .cpu the default stream in mlx-swift but it would be helpful to update the eight distributed functions to convert an empty mlx_stream handle (ctx == NULL) into StreamOrDevice{}, while passing explicit streams through unchanged. I can then update mlx-swift to default to nil instead of .cpu. Since this is generated code, I assume this must be done by the maintainers.

Testing

Compiled against MLX v0.32.2 and checked from a C99 caller with
-pedantic -Wall -Wextra; both files pass clang-format. The functions have
not been tested at runtime.

[andresy]

@ronaldmannak would you mind allowing edits from maintainers? There are few changes I would like to make before merging

[ronaldmannak]

@andresy Yes, absolutely. No problem. I'm a bit lost what should be generated and what not, so go ahead. Thanks!

[andresy]

Hmm somehow I cannot push changes there - maybe because it is an organization-own fork?

[ronaldmannak]

Weird. I don't know how GitHub handles orgs. But I see you fixed it by creating a new PR. Perfect