Skip to content

Add endpoint parameter to linspace - #4184

Draft
AKnassa wants to merge 1 commit into
ml-explore:mainfrom
AKnassa:add-linspace-endpoint
Draft

Add endpoint parameter to linspace#4184
AKnassa wants to merge 1 commit into
ml-explore:mainfrom
AKnassa:add-linspace-endpoint

Conversation

@AKnassa

@AKnassa AKnassa commented Aug 11, 2026

Copy link
Copy Markdown

Proposed changes

Fixes #4152.

mx.linspace always includes the stop value. numpy, jax and the array API all
let the caller drop it and space the samples over the half-open interval
[start, stop) instead. This adds the same option:

mx.linspace(0, 10, 5)                  # [0, 2.5, 5, 7.5, 10]
mx.linspace(0, 10, 5, endpoint=False)  # [0, 2, 4, 6, 8]   (matches numpy)

The implementation is the one-line change it looks like: divide by num instead
of num - 1 when the endpoint is excluded.

endpoint defaults to true, so existing behaviour is unchanged.

Placement of the argument

I put endpoint after dtype and before the stream, both in the C++ signature
and in the Python binding, so the stream stays last as it is everywhere else in
the API. That does shift the position of StreamOrDevice in the C++ signature,
but no in-tree caller passes the stream positionally to linspace, and the
Python stream argument is keyword-only in practice. Happy to move it if you
would rather it went elsewhere — the array API puts endpoint keyword-only
after dtype/device, while numpy has it fourth positional, so there is no
single obvious answer.

Testing

  • test_linspace_endpoint in python/tests/test_ops.py, checked against numpy
    for num in {0, 1, 2, 5, 50}, for negative ranges, and with an integer
    dtype. It fails before this change and passes after.
  • Extended the existing test linspace case in tests/ops_tests.cpp, since the
    C++ signature changed, including the num == 1 and num == 0 edges.
  • Full Python suite: 820 tests, no new failures.
  • Full C++ suite: 247 cases, all pass.
  • Docstring updated; the generated .pyi picks the new signature up on build.

Built CPU-only (MLX_BUILD_METAL=OFF); the GPU suite was not run on my machine.

Checklist

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

numpy, jax and the array API all let the caller drop the stop value
and space the samples over the half-open interval [start, stop)
instead. Add the same option, spacing by (stop - start) / num rather
than (stop - start) / (num - 1) when it is off.

endpoint defaults to true, so existing behaviour is unchanged.

Fixes ml-explore#4152
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Adding endpoint parameter in the mlx.core.linspace to make it more flexible and compliant with Array API

1 participant