Skip to content

Use new DType API with backcompat path and fix result_type() - #374

Open
seberg wants to merge 10 commits into
jax-ml:mainfrom
seberg:try-new-dtype-hack
Open

Use new DType API with backcompat path and fix result_type()#374
seberg wants to merge 10 commits into
jax-ml:mainfrom
seberg:try-new-dtype-hack

Conversation

@seberg

@seberg seberg commented Apr 20, 2026

Copy link
Copy Markdown
Contributor

gh-360 is a little bit hard to actually pull of. In NumPy 2.5+ (backport vendorered to compile on older NumPy versions), it is now allowed to create a "new style DType" that still is a legacy dtype.
That means we inherit a few existing quirks, but it mostly means that ml_dtypes can start using new API conveniently without much concern for backwards compatibility issues.

The current state is:

  • result_type() now does a better job/the right thing (yes this adds a bunch of code).
  • There is a small annoyance that arrays (at least on <2.5) print with array(..., dtype=dtype('bfloat16')) rather than just bfloat16. (I am restoring that in NumPy for 2.5+, although we need a better solution).
  • NumPy wants the copy cast defined, so I did that. Probably speeds up strided copies, but it is more of a requirement.

Next steps, unlocked things:

  • We can implement np.finfo() (given one backport PR from me in NumPy) for NumPy 2.5+.
  • I think it would be nice to move all casts over but it won't do much besides speeding up strided casts quite a lot.
  • There may be some other new API now accessible, although I can't think of anything super interesting right now. But of course if we add things in NumPy it'll now be easier to use it quickly (similar to the finfo).

I agree that the long list of dtypes is a bit annoying for CommonDType. After NumPy has a better DType hierarchy, we could allow creating your own baseclass there easily which would simplify this (although maybe not speed up in practice).

(heavy use of claude to spit out code, but of course with absolute design micro-managing in many relevant parts -- but I'll need to go through once more myself).

@seberg
seberg force-pushed the try-new-dtype-hack branch 2 times, most recently from de3a5a8 to eff09d7 Compare April 20, 2026 12:58
@seberg

seberg commented Apr 21, 2026

Copy link
Copy Markdown
Contributor Author

@hawkinsp just in case you have a quick thought here. I tried to rewrite things to just use the new API but keep it a "legacy" dtype to some degree, so that there should be no real regressions but at the same time it works fine all the way back to NumPy 2 (the only regression I noticed it that arrays print a bit less nice).

However, this uses PyType_Ready. NumPy predated PyType_FromMetaclass and it still implements e.g. tp_new making it's use incompatible.
So that is a bigger downside with this approach: adopting the Python stable API may be hard or indefinitely deferred, because one would need to hack around this for old NumPy versions.
(I am pretty sure it is possible to hack around it, NumPy effectively does it and I think so does pybind11 probably. But it may be pretty ugly...)

The alternative I can currently think off is to just allow PyArrayInitDTypeMeta_FromSpec to amend the current legacy dtype.
Less forward looking and also I liked how this backported while the amending pattern is nice for new NumPy versions but backports worse, I expect (mainly because of cast definition patterns).

@hawkinsp

Copy link
Copy Markdown
Collaborator

I haven't looked yet, but

However, this uses PyType_Ready. NumPy predated PyType_FromMetaclass and it still implements e.g. tp_new making it's use incompatible.
So that is a bigger downside with this approach: adopting the Python stable API may be hard or indefinitely deferred, because one would need to hack around this for old NumPy versions.

It's not the end of the world to have to build ml_dtypes per Python version: it's a small enough package. I had previously abandoned trying to use the limited dtype API for similar reasons (#195). And eventually when the oldest supported NumPy ages off our support matrix, we can switch.

@seberg

seberg commented Apr 22, 2026

Copy link
Copy Markdown
Contributor Author

OK, cool, then I think I'll pursue this, we need a better way to transition a package like ml_dtypes and I think this is viable.

Long term for the stable API: I suspect the right thing will be to have a new DType creation function, that creates the full heap-type for you based on the spec.
(That way, even if we need a bit crazy things, that can live in NumPy. I.e. PyArrayDTypeMeta_FromSpecs(module, type_slots, dtype_slots), but that'll be a NumPy 2.6 thing at best -- I am also very curious about the Python stable API developments around this. Back in the day, I stole their ideas, but they seemed to have improved on them quite a lot!)

@seberg seberg changed the title EXPLORATORY: Add a path that allows using new DType API Use new-style DType API with backcompat path and fix result_type() Jun 5, 2026
@seberg seberg changed the title Use new-style DType API with backcompat path and fix result_type() Use new DType API with backcompat path and fix result_type() Jun 5, 2026
@seberg
seberg force-pushed the try-new-dtype-hack branch from c67581e to 800f815 Compare June 5, 2026 10:48
@seberg

seberg commented Jun 5, 2026

Copy link
Copy Markdown
Contributor Author

I made a big pass on this cleaning things up and hopefully fixing some issues (byte-swapping and an incorrect result_type() branching).

This should now actually be OK, although it does require NumPy 2.0+.

With this being merged in NumPy and working out fine here, this is now actually ready. I'll note again the one little annoyance I have found for now (hopefully the only one), and that is that when printing arrays we now print dtype=dtype('bfloat16') in long-form.
For NumPy 2.5, I'll try to just preserve the current behavior (even if not perfect for byte-swapping), then we can see.
(In theory we could monkey-patch NumPy if we wanted to.)

Unlike gh-360 this doesn't make it obvious to just not set a character for example but it does ensure that old-style code paths are taken so that actual regressions are unlikely (while it seems it'll be a long whack-a-mole with gh-360).

@seberg
seberg marked this pull request as ready for review June 5, 2026 10:52
@seberg

seberg commented Jun 5, 2026

Copy link
Copy Markdown
Contributor Author

This also now fixes gh-301 which I suspect is the magic thing that might make my CuPy "test everything" attempt feasible (to the point I am considering if we should just hack that flag if this PR isn't so easy).

(Nevermin, older NumPy will still not support it nicely of course. If we want to improve this without the flag, I think the solution might be to check if isfinite is defined for the dtype or not.)

@hawkinsp

Copy link
Copy Markdown
Collaborator

Sorry this had dropped off my radar of things to look at.

I asked Gemini to review this and it pointed out some issues that I think may be real at least in part. I asked it to make test cases.

# Copyright 2026 The ml_dtypes Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for np.result_type() across ml_dtypes custom DTypes."""

import itertools

import ml_dtypes
import numpy as np
import pytest

# Short aliases for readability in parametrize lists
bf16 = ml_dtypes.bfloat16
f4   = ml_dtypes.float4_e2m1fn
f6_e2m3 = ml_dtypes.float6_e2m3fn
f6_e3m2 = ml_dtypes.float6_e3m2fn
f8_e3m4 = ml_dtypes.float8_e3m4
f8_e4m3 = ml_dtypes.float8_e4m3
f8_e4m3fn   = ml_dtypes.float8_e4m3fn
f8_e4m3fnuz = ml_dtypes.float8_e4m3fnuz
f8_e4m3b11  = ml_dtypes.float8_e4m3b11fnuz
f8_e5m2     = ml_dtypes.float8_e5m2
f8_e5m2fnuz = ml_dtypes.float8_e5m2fnuz
f8_e8m0     = ml_dtypes.float8_e8m0fnu
bc32 = ml_dtypes.bcomplex32
c32  = ml_dtypes.complex32
i1, i2, i4 = ml_dtypes.int1,  ml_dtypes.int2,  ml_dtypes.int4
u1, u2, u4 = ml_dtypes.uint1, ml_dtypes.uint2, ml_dtypes.uint4

ALL_CUSTOM_FLOATS = [bf16, f4, f6_e2m3, f6_e3m2,
                     f8_e3m4, f8_e4m3, f8_e4m3fn, f8_e4m3fnuz,
                     f8_e4m3b11, f8_e5m2, f8_e5m2fnuz, f8_e8m0]
ALL_INTN = [i1, i2, i4, u1, u2, u4]
ALL_CUSTOM_COMPLEX = [bc32, c32]


def rt(a, b):
  return np.result_type(a, b)


# ---------------------------------------------------------------------------
# Custom float vs NumPy built-in types
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("a, b, expected", [
    # ---- bool: custom float always wins ----
    (bf16,      np.bool_, bf16),
    (f8_e4m3fn, np.bool_, f8_e4m3fn),
    (f4,        np.bool_, f4),
    # ---- floats: pick the wider ----
    (f4,        np.float16, np.float16),      # f4 fits in float16
    (f8_e4m3fn, np.float16, np.float16),      # float8 fits in float16
    (f8_e5m2,   np.float16, np.float16),      # float8 fits in float16
    (bf16,      np.float16, np.float32),      # incomparable → float32
    (f8_e4m3fn, np.float32, np.float32),      # all custom floats fit in float32
    (bf16,      np.float32, np.float32),
    (bf16,      np.float64, np.float64),
    (f8_e4m3fn, np.float64, np.float64),
    # ---- integers: PyArray_CommonDType decides ----
    (bf16,      np.int8,  bf16),              # bfloat16 has 8 sig bits, int8 needs 7 → bf16 wins
    (bf16,      np.int16, np.float64),        # bfloat16 has 8 sig bits, int16 needs 15 → float64
    (f8_e4m3fn, np.int8,  np.float64),        # float8 can't represent all int8 values
    (f8_e4m3fn, np.int32, np.float64),        # float8 can't represent all int32 values
    # ---- complex: other always wins ----
    (bf16,      np.complex64,  np.complex64),
    (f8_e4m3fn, np.complex64,  np.complex64),
    (bf16,      np.complex128, np.complex128),
    (f8_e4m3fn, np.complex128, np.complex128),
])
def test_custom_float_vs_numpy(a, b, expected):
  assert rt(a, b) == np.dtype(expected)
  assert rt(b, a) == np.dtype(expected)  # must be symmetric


# ---------------------------------------------------------------------------
# Custom float vs custom float
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("a, b, expected", [
    # ---- same type ----
    (bf16,      bf16,      bf16),
    (f8_e4m3fn, f8_e4m3fn, f8_e4m3fn),
    (f4,        f4,        f4),
    # ---- narrower fits safely into wider ----
    (f4,        f6_e2m3,   f6_e2m3),          # f4 ⊂ f6_e2m3 (more exp + mantissa)
    (f4,        f8_e4m3fn, f8_e4m3fn),        # f4 fits in every float8+
    (f4,        bf16,      bf16),             # f4 fits in bfloat16
    (f8_e4m3fn, bf16,      bf16),             # float8 fits in bfloat16
    (f8_e5m2,   bf16,      bf16),             # float8 fits in bfloat16
    (f8_e3m4,   bf16,      bf16),             # float8 fits in bfloat16
    # ---- incomparable: one has more exp, other more mantissa → float32 ----
    (bf16,      f8_e5m2,   bf16),             # f8_e5m2 fits in bf16 (bf16 > in all dims)
    (f8_e4m3fn, f8_e5m2,   np.float32),       # e4m3 has more mantissa, e5m2 has more exp
    (f8_e4m3fn, f8_e4m3fnuz, f8_e4m3fn),      # same digits/max_exp → numeric_limits match; fn wins
    (f6_e2m3,   f6_e3m2,   np.float32),       # one has more mantissa, other more exp
])
def test_custom_float_vs_custom_float(a, b, expected):
  assert rt(a, b) == np.dtype(expected)
  assert rt(b, a) == np.dtype(expected)  # must be symmetric


# ---------------------------------------------------------------------------
# Custom float vs custom int  (float always dominates)
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("float_t, int_t", [
    (bf16,      i4),
    (bf16,      u4),
    (bf16,      i1),
    (f8_e4m3fn, i4),
    (f8_e4m3fn, u4),
    (f8_e5m2,   i2),
    (f4,        i1),
])
def test_custom_float_beats_custom_int(float_t, int_t):
  assert rt(float_t, int_t) == np.dtype(float_t)
  assert rt(int_t, float_t) == np.dtype(float_t)  # symmetric


# ---------------------------------------------------------------------------
# Custom int vs NumPy built-in types
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("a, b, expected", [
    # ---- bool: custom int always wins ----
    (i4,  np.bool_, i4),
    (u4,  np.bool_, u4),
    (i1,  np.bool_, i1),
    # ---- all other NumPy types: return other (intN is always smaller) ----
    (i4,  np.int8,    np.int8),
    (i4,  np.int16,   np.int16),
    (i4,  np.int32,   np.int32),
    (i4,  np.uint8,   np.uint8),
    (u4,  np.int8,    np.int8),
    (i2,  np.int8,    np.int8),
    (i4,  np.float16, np.float16),
    (i4,  np.float32, np.float32),
    (i4,  np.float64, np.float64),
    (i4,  np.complex64,  np.complex64),
    (i4,  np.complex128, np.complex128),
    (u4,  np.float32,    np.float32),
])
def test_custom_int_vs_numpy(a, b, expected):
  assert rt(a, b) == np.dtype(expected)
  assert rt(b, a) == np.dtype(expected)  # must be symmetric


# ---------------------------------------------------------------------------
# Custom int vs custom int
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("a, b, expected", [
    # ---- same type ----
    (i4, i4, i4),
    (u4, u4, u4),
    # ---- mixed sign: neither fits the other → int16 ----
    (i4, u4, np.int16),
    (i2, u2, np.int16),
    (i1, u1, np.int16),
    # ---- same sign, different width → int16 fallback ----
    (i2, i4, np.int16),
    (u2, u4, np.int16),
    (i1, i4, np.int16),
])
def test_custom_int_vs_custom_int(a, b, expected):
  assert rt(a, b) == np.dtype(expected)
  assert rt(b, a) == np.dtype(expected)  # must be symmetric


# ---------------------------------------------------------------------------
# Custom complex vs NumPy built-in types
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("a, b, expected", [
    # ---- bool + integers: wrap in cfloat ----
    (bc32, np.bool_,   np.complex64),
    (bc32, np.int8,    np.complex64),
    (bc32, np.int32,   np.complex64),
    (c32,  np.bool_,   np.complex64),
    (c32,  np.int8,    np.complex64),
    # ---- floats ≤ float32: wrap in cfloat ----
    (bc32, np.float16, np.complex64),
    (bc32, np.float32, np.complex64),
    (c32,  np.float16, np.complex64),
    (c32,  np.float32, np.complex64),
    # ---- float64+: need cdouble ----
    (bc32, np.float64,    np.complex128),
    (bc32, np.longdouble, np.clongdouble),
    (c32,  np.float64,    np.complex128),
    # ---- built-in complex: other always wins ----
    (bc32, np.complex64,  np.complex64),
    (bc32, np.complex128, np.complex128),
    (c32,  np.complex64,  np.complex64),
    (c32,  np.complex128, np.complex128),
])
def test_custom_complex_vs_numpy(a, b, expected):
  assert rt(a, b) == np.dtype(expected)
  assert rt(b, a) == np.dtype(expected)  # must be symmetric


# ---------------------------------------------------------------------------
# Custom complex vs custom float / custom int
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("a, b, expected", [
    # ---- custom floats: all fit in cfloat alongside our complex ----
    (bc32, bf16,      np.complex64),
    (bc32, f8_e4m3fn, np.complex64),
    (bc32, f8_e5m2,   np.complex64),
    (bc32, f4,        np.complex64),
    (c32,  bf16,      np.complex64),
    (c32,  f8_e4m3fn, np.complex64),
    # ---- custom ints: all tiny, fit in cfloat ----
    (bc32, i4,  np.complex64),
    (bc32, u4,  np.complex64),
    (bc32, i1,  np.complex64),
    (c32,  i4,  np.complex64),
    # ---- two custom complex types ----
    (bc32, c32,  np.complex64),
    (bc32, bc32, bc32),
    (c32,  c32,  c32),
])
def test_custom_complex_vs_custom(a, b, expected):
  assert rt(a, b) == np.dtype(expected)
  assert rt(b, a) == np.dtype(expected)  # must be symmetric


# ---------------------------------------------------------------------------
# Python scalars: 0, 0.0, 0.0j  (abstract types)
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("dtype, scalar, expected", [
    # ---- custom floats dominate Python int and Python float ----
    (bf16,      0,    bf16),
    (bf16,      0.0,  bf16),
    (f8_e4m3fn, 0,    f8_e4m3fn),
    (f8_e4m3fn, 0.0,  f8_e4m3fn),
    (f4,        0,    f4),
    (f4,        0.0,  f4),
    # ---- custom float + Python complex → cfloat ----
    (bf16,      0.0j, np.complex64),
    (f8_e4m3fn, 0.0j, np.complex64),
    (f4,        0.0j, np.complex64),
    # ---- custom ints: a Python int defers to the int dtype, but a Python
    #      float/complex crosses the integer kind and promotes to the default
    #      float64 / complex128 (matching NumPy's built-in integers) ----
    (i4,  0,    i4),
    (i4,  0.0,  np.float64),
    (i4,  0.0j, np.complex128),
    (u4,  0,    u4),
    (u4,  0.0,  np.float64),
    (u4,  0.0j, np.complex128),
    (i1,  0,    i1),
    (i1,  0.0,  np.float64),
    (i2,  0.0j, np.complex128),
    # ---- custom complex dominates all Python scalars ----
    (bc32, 0,    bc32),
    (bc32, 0.0,  bc32),
    (bc32, 0.0j, bc32),
    (c32,  0,    c32),
    (c32,  0.0,  c32),
    (c32,  0.0j, c32),
])
def test_python_scalars(dtype, scalar, expected):
  assert rt(dtype, scalar) == np.dtype(expected)


@pytest.mark.parametrize("int_t", ALL_INTN)
@pytest.mark.parametrize("scalar, concrete, expected", [
    (1.0,  np.float16,   np.float16),
    (1.0,  np.float32,   np.float32),
    (1.0,  np.float64,   np.float64),
    (1.0,  np.complex64, np.complex64),
    (1.0j, np.float16,   np.complex64),
    (1.0j, np.float32,   np.complex64),
    (1.0j, np.float64,   np.complex128),
    (1.0j, np.complex64, np.complex64),
])
def test_weak_scalar_stays_weak(int_t, scalar, concrete, expected):
  # Sanity check that no-matter the order the concrete precision wins
  # (i.e. promoting int + pyfloat -> pyfloat).
  for args in itertools.permutations([int_t, scalar, concrete]):
    assert np.result_type(*args) == np.dtype(expected), args


# ---------------------------------------------------------------------------
# Problem 1: Signed custom int vs unsigned NumPy built-ins
#
# NPyIntN_CommonDType returns `other` for all non-userdef NumPy types
# (!PyTypeNum_ISUSERDEF(other->type_num)). When `cls` is a signed type
# (int1, int2, int4) and `other` is an unsigned integer (uint8, uint16, etc.),
# returning `other` produces an unsigned result. Unsigned types cannot
# represent negative values of signed intN (e.g. result_type(int4, uint8) ->
# uint8 wraps -1 to 255).
#
# Expected: For signed intN with unsigned built-in types, delegate to
# PyArray_CommonDType(&PyArray_ByteDType, other) to follow standard NumPy
# signed+unsigned promotion (int4 + uint8 -> int16, int4 + uint16 -> int32, etc.).
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("signed_int", [i1, i2, i4])
@pytest.mark.parametrize("unsigned_builtin, expected", [
    (np.uint8,  np.int16),
    (np.uint16, np.int32),
    (np.uint32, np.int64),
    (np.uint64, np.float64),
])
def test_signed_intn_vs_unsigned_builtins(signed_int, unsigned_builtin, expected):
  assert rt(signed_int, unsigned_builtin) == np.dtype(expected)
  assert rt(unsigned_builtin, signed_int) == np.dtype(expected)


# ---------------------------------------------------------------------------
# Problem 2: Custom float commutativity & asymmetry in TRY_CUSTOM_FLOAT
#
# In TRY_CUSTOM_FLOAT (custom_float.h), if two distinct custom float types have
# identical precision and exponent range (e.g. float8_e5m2 and float8_e5m2fnuz,
# both with digits=3 and max_exponent=16), CustomFloatSafeTo<T, OtherT>()
# evaluates to true in both directions. Because the macro returns `other` on
# the first match, CommonDType(A, B) returns B while CommonDType(B, A) returns A,
# breaking commutativity (result_type(a, b) != result_type(b, a)).
#
# Expected: All pairs of custom float types must be commutative.
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("a, b", itertools.combinations(ALL_CUSTOM_FLOATS, 2))
def test_custom_float_commutativity(a, b):
  assert rt(a, b) == rt(b, a)


# ---------------------------------------------------------------------------
# Problem 2 (cont.) & Problem 3: Incomparable floats & float8_e8m0fnu vs float16
#
# Problem 2 (cont.): Distinct float formats with equal precision/range without
# strict containment (e.g. float8_e5m2 vs float8_e5m2fnuz) should promote to
# a common supertype (float32).
#
# Problem 3: In NPyCustomFloat_CommonDType, case NPY_HALF only checks
# `if (is_bfloat16)` before falling through to return `other` (float16).
# However, float8_e8m0fnu also has max_exponent=128 (dynamic range up to 2^127),
# which exceeds float16's max value (~65504). Falling through to float16 causes
# float8_e8m0 values to overflow.
#
# Expected: float8_e8m0fnu vs float16 should promote to float32 (matching
# bfloat16 vs float16).
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("a, b, expected", [
    (f8_e5m2, f8_e5m2fnuz, np.float32),
    (f8_e8m0, np.float16,   np.float32),
])
def test_incomparable_custom_float_promotions(a, b, expected):
  assert rt(a, b) == np.dtype(expected)
  assert rt(b, a) == np.dtype(expected)


# ---------------------------------------------------------------------------
# Problem 4a: Custom complex vs bool over-promotion
#
# In NPyCustomComplex_CommonDType, NPY_BOOL is grouped with NPY_BYTE/NPY_INT/
# NPY_FLOAT and returns cfloat (complex64). Since booleans (0 and 1) fit
# exactly in 16-bit complex scalars, returning complex64 unnecessarily widens
# the dtype (unlike custom float and custom int types which return cls).
#
# Expected: bool + bcomplex32 -> bcomplex32, bool + complex32 -> complex32.
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("complex_t, bool_expected", [
    (bc32, bc32),
    (c32,  c32),
])
def test_custom_complex_vs_bool(complex_t, bool_expected):
  assert rt(complex_t, np.bool_) == np.dtype(bool_expected)
  assert rt(np.bool_, complex_t) == np.dtype(bool_expected)


# ---------------------------------------------------------------------------
# Problem 4b: Component float + custom complex over-promotion
#
# In NPyCustomComplex_CommonDType, combining bfloat16 with bcomplex32 or
# float16 with complex32 widens to complex64. Because bfloat16 is the exact
# real component of bcomplex32 (and float16 is the real component of complex32),
# they should stay in the matching 32-bit complex type without widening to
# 64-bit (analogous to float32 + complex64 -> complex64).
# Likewise, narrower floats (float4, float8) fit within complex32.
#
# Expected: bfloat16 + bcomplex32 -> bcomplex32, float16/float8 + complex32 -> complex32.
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("complex_t, float_t, expected", [
    (bc32, bf16,        bc32),
    (c32,  np.float16,  c32),
    (c32,  f4,          c32),
    (c32,  f8_e4m3fn,   c32),
    (c32,  f8_e5m2,     c32),
])
def test_custom_complex_vs_matching_float(complex_t, float_t, expected):
  assert rt(complex_t, float_t) == np.dtype(expected)
  assert rt(float_t, complex_t) == np.dtype(expected)


# ---------------------------------------------------------------------------
# Problem 4b (cont.): Custom int vs custom complex over-promotion
#
# Custom intN types (int1..4, uint1..4) all have values in [-8, 15], which fit
# comfortably within bfloat16 (8 bits precision) and float16 (11 bits precision).
# Combining custom intN with bcomplex32/complex32 should not widen to complex64.
#
# Expected: intN + bcomplex32 -> bcomplex32, intN + complex32 -> complex32.
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("complex_t", ALL_CUSTOM_COMPLEX)
@pytest.mark.parametrize("int_t", [i1, i2, i4, u1, u2, u4])
def test_custom_complex_vs_custom_int(complex_t, int_t):
  assert rt(complex_t, int_t) == np.dtype(complex_t)
  assert rt(int_t, complex_t) == np.dtype(complex_t)


# ---------------------------------------------------------------------------
# Problem 4c: 64-bit integer vs custom complex precision loss
#
# In NPyCustomComplex_CommonDType, case NPY_LONGLONG / NPY_ULONGLONG returns
# PyArray_CFloatDType (complex64). However, float32 has only 24 bits of mantissa
# precision and cannot represent all 64-bit integer values without loss.
# Standard NumPy promotion rules promote 64-bit integers with complex to
# complex128 (cdouble).
#
# Expected: int64/uint64 + bcomplex32/complex32 -> complex128.
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("complex_t", ALL_CUSTOM_COMPLEX)
@pytest.mark.parametrize("int64_t", [np.int64, np.uint64])
def test_custom_complex_vs_64bit_int(complex_t, int64_t):
  assert rt(complex_t, int64_t) == np.dtype(np.complex128)
  assert rt(int64_t, complex_t) == np.dtype(np.complex128)


WDYT?

@hawkinsp

Copy link
Copy Markdown
Collaborator

#360 is looking more plausible to merge, but I'm wondering if it makes sense to take some of the backward compat ideas from this PR.

@seberg

seberg commented Aug 19, 2026

Copy link
Copy Markdown
Contributor Author

I'll have a look very soon. I think this is the necessary approach, but it is very possible that #360 is getting those promotion cases more correct. I think I basically left some existing over-promotions untouched here (but don't exactly recall).

This uses the new-style API for NumPy 2.0+.  NumPy 2.5 ships some
backport compatibility hacks that we vendor here to compile on
older NumPy versions as well.

This requires some churn, the biggest one being settuping a within
DType casting implementation.

However, it allows using new API optionally the biggest thing being
that `result_type()` can now do the right thing.

The one downside is that the `dtype=` is printing not so nice
for NumPy <2.5.
@seberg
seberg force-pushed the try-new-dtype-hack branch from 418cf95 to 3cbcc72 Compare September 3, 2026 15:38
@seberg

seberg commented Sep 3, 2026

Copy link
Copy Markdown
Contributor Author

OK, asked my agent to adopt those and fixed with a short review (I should have another look through for sanity, although the tests look good).
Overall, the points looked right, most were just not covered and incorrect.

Added a few more smaller fixes, this probably needs an iteration, but it may work now.

(Again about the other PR, it does more I am sure, but I think backwards compat mean we have to go this route, then basically continue the work there to use the new API as much as possible. That makes sense, but I wanted to not do everything in one go and the result-type is the most interesting fix, I think.)

(I'll work on the failures, there are some more follow-ups here probably.)

@seberg
seberg force-pushed the try-new-dtype-hack branch from 8ddc579 to 623c66c Compare September 4, 2026 09:23
@seberg
seberg force-pushed the try-new-dtype-hack branch from 54c74b4 to 2127317 Compare September 4, 2026 12:38
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.

2 participants