Fix complex arr.real/arr.imag on NumPy 2.5+#375
Closed
seberg wants to merge 1 commit into
Closed
Conversation
This adds the new mechanism for defining .imag/.real. Admittedly, we don't actually need the full loop (and if we do more new DType things, the looping part should be a helper construct). I.e. `.imag` and `.real` itself will only ever use the resolution function here and return a view. This is a draft as it can only be merged after jax-mlgh-357 (and requires doc fixup).
DeanTMaxim
added a commit
to DeanTMaxim/ml_dtypes
that referenced
this pull request
Jun 14, 2026
Takes over jax-ml#375 per issue jax-ml#355. Registers real/imag as ufunc loops on the np.real/np.imag ufuncs via PyUFunc_AddLoopsFromSpecs (the NumPy 2.4+ convenience API), gated on PyArray_RUNTIME_VERSION >= 0x16 (NPY_2_5_API_VERSION) so NumPy <2.5 skips registration. On NumPy 2.5+ this makes arr.real / arr.imag return correct results for bcomplex32 and complex32, which NumPy otherwise cannot recognize as complex (their dtype kind is 'W', not 'c'). The registered loop returns NPY_NO_CASTING with a view_offset (0 for the real half, elsize for the imaginary half), so the result is a zero-copy view into the input buffer, with a reinterpret-based strided loop as the fallback. static_asserts guard the [real, imag] no-padding layout that the view/extraction relies on. PyUFunc_AddLoopsFromSpecs is backported in numpy.h (slot 47) so a single wheel runs across NumPy versions. On NumPy <2.5 the native arr.real/arr.imag are silently wrong, so the ml_dtypes.real/imag Python helpers (correct on all versions) emit a RuntimeWarning steering users toward the helpers or an upgrade. Builds as C++17: the PyArrayMethod_Spec is constructed by field assignment rather than designated initializers, so no build-standard change is needed. Adds a NumPy 2.5 pre-release job to the CI matrix. Several other complex-aware builtins (np.iscomplex, np.vdot, np.linalg.norm, np.angle, np.linalg.det/inv) do not recognize these custom complex dtypes on any NumPy version, because NumPy keys complex-ness off the builtin type-number range rather than the dtype kind; bcomplex32.__doc__ documents these with one-line workarounds. Tested locally against NumPy 2.5.0rc1.
DeanTMaxim
added a commit
to DeanTMaxim/ml_dtypes
that referenced
this pull request
Jun 14, 2026
Takes over jax-ml#375 per issue jax-ml#355. Registers real/imag as ufunc loops on the np.real/np.imag ufuncs via PyUFunc_AddLoopsFromSpecs (the NumPy 2.4+ convenience API), gated on PyArray_RUNTIME_VERSION >= 0x16 (NPY_2_5_API_VERSION) so NumPy <2.5 skips registration. On NumPy 2.5+ this makes arr.real / arr.imag return correct results for bcomplex32 and complex32, which NumPy otherwise cannot recognize as complex (their dtype kind is 'W', not 'c'). The registered loop returns NPY_NO_CASTING with a view_offset (0 for the real half, elsize for the imaginary half), so the result is a zero-copy view into the input buffer, with a reinterpret-based strided loop as the fallback. static_asserts guard the [real, imag] no-padding layout that the view/extraction relies on. PyUFunc_AddLoopsFromSpecs is backported in numpy.h (slot 47) so a single wheel runs across NumPy versions. On NumPy <2.5 the native arr.real/arr.imag are silently wrong, so the ml_dtypes.real/imag Python helpers (correct on all versions) emit a RuntimeWarning steering users toward the helpers or an upgrade. Builds as C++17: the PyArrayMethod_Spec is constructed by field assignment rather than designated initializers, so no build-standard change is needed. Adds a NumPy 2.5 pre-release job to the CI matrix. Several other complex-aware builtins (np.iscomplex, np.vdot, np.linalg.norm, np.angle, np.linalg.det/inv) do not recognize these custom complex dtypes on any NumPy version, because NumPy keys complex-ness off the builtin type-number range rather than the dtype kind; bcomplex32.__doc__ documents these with one-line workarounds. Tested locally against NumPy 2.5.0rc1.
|
Taking this over in #383 (builds as C++17, adds the <2.5 warning). |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This adds the new mechanism for defining .imag/.real. Admittedly, we don't actually need the full loop (and if we do more new DType things, the looping part should be a helper construct). I.e.
.imagand.realitself will only ever use the resolution function here and return a view.This is a draft as it can only be merged after gh-357 (and requires doc fixup). It also requires an update after NumPy did the necessary C-versioning bump on main (missing until now).