Complete bcomplex32 real/imag on NumPy 2.5+#383
Open
DeanTMaxim wants to merge 2 commits into
Open
Conversation
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.
This was referenced Jun 14, 2026
seberg
reviewed
Jun 14, 2026
| Returns: | ||
| The real part of the input array. | ||
| """ | ||
| _warn_old_numpy("real") |
Contributor
There was a problem hiding this comment.
Nice thought, but this functino is fine. We would basically have to use a module level __getattr__ and warn on first use of any of the complex type.
I am not sure that is actually worth it, though since it is also easy to miss or need to ignore and then does it globally.
(We could if others like it though, or just limit them to NumPy 2.5+ even if a bit of a shame.)
Author
There was a problem hiding this comment.
Dropped — relying on the docstring now (a module-level getattr would need the types lazily exposed to actually fire).
The warning fired inside ml_dtypes.real/imag, which are already correct on all NumPy versions, so it only reached users of the safe helpers -- not the arr.real/arr.imag users who hit the actual silent-wrong on <2.5. Drop it and rely on the docstring; a module-level __getattr__ guard would be the correct placement but needs the types exposed lazily to actually fire. Removes _warn_old_numpy, its calls, the warning test, the now-unneeded RuntimeWarning suppressions in the other tests, and the CHANGELOG entry.
DeanTMaxim
changed the title
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.
Takes over #375 per the discussion in #355.
What this changes
real/imagas ufunc loops onnp.real/np.imagforbcomplex32andcomplex32viaPyUFunc_AddLoopsFromSpecs(NumPy 2.4+ API), gated onPyArray_RUNTIME_VERSION >= 0x16so NumPy <2.5 is a no-op. On 2.5+,arr.real/arr.imagnow return correct results. The loop is a zero-copy view (NPY_NO_CASTING+view_offset0 / elsize) with a reinterpret-based strided fallback.bcomplex32.__doc__lists the other complex-aware builtins that don't recognize custom complex dtypes, with one-line workarounds.Differences from your draft (#375)
PyArrayMethod_Specuses field assignment instead of designated initializers, so there's no build-standard change.static_asserts in the strided loop guarding the[real, imag]no-padding layout the view/extraction relies on.<2.5helper warning — per review it only reached the already-correct helper, notarr.real.)Remaining limitations (documented, out of scope here)
np.iscomplex,np.vdot,np.linalg.norm,np.angle,np.linalg.det/invstill don't recognize these dtypes on any NumPy version. Root cause: NumPy keys "is complex" off the builtin type-number range (PyTypeNum_ISCOMPLEX), not the dtypekind; these are registered withkind='W'(using'c'would collidebcomplex32withcomplex32under same-kind+same-size equality — verified empirically that evenkind='c'leavesissubdtype(complexfloating)False). Documented inbcomplex32.__doc__; a permanent fix is NumPy-side (extending NEP 42 so user DTypes can declare complex-ness) and tracked separately.Status
< 0x16runtime gate:arr.realcorrect on 2.5.0rc1, silently wrong on 2.4.4.cc @seberg — happy to adjust. Intended for v0.5.5; unblocks jax-ml/jax#38392.
v0.5.5 release notes (draft)
arr.real/arr.imagnow work correctly forbcomplex32/complex32arrays on NumPy 2.5+.