Skip to content

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

Merged
copybara-service[bot] merged 1 commit into
jax-ml:mainfrom
seberg:try-new-dtype-hack
Oct 22, 2026
Merged

copybara-service[bot] merged 1 commit 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).

@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 2 times, most recently from 54c74b4 to 2127317 Compare September 4, 2026 12:38
@seberg

seberg commented Sep 29, 2026

Copy link
Copy Markdown
Contributor Author

Remembered, this so just a gentle ping (actually I remembered more than this, that I need the character change in CuPy, because without that promotion isn't just weird but very broken because can_cast() guesses incorrect).

@hawkinsp

Copy link
Copy Markdown
Collaborator

Thanks for the ping, I'll take a look.

@hawkinsp

hawkinsp commented Oct 1, 2026

Copy link
Copy Markdown
Collaborator

I haven't had time to stare deeply yet myself, but AI says this:


Problem 1: 3+ Argument np.result_type Raises DTypePromotionError When Mixing Custom Floats with Python complex (1j)

Root Cause

In NPyCustomFloat_CommonDType, Python int and float abstract scalars are handled, but Python complex (PyArray_PyComplexDType) was left as a comment:

// Python abstract scalars defer to the concrete type. (should add complex
// here)
if (other == &PyArray_PyLongDType || other == &PyArray_PyFloatDType) {
  Py_INCREF(cls);
  return cls;
}

When other == &PyArray_PyComplexDType, NPyCustomFloat_CommonDType falls through to return Py_NotImplemented.

  • 2 arguments: np.result_type(bfloat16, 1j) succeeds only because PyArray_CommonDType retries in the reverse order (PyComplexDType.__common_dtype__(bfloat16)), which probes bfloat16.__common_dtype__(CFloatDType) and returns complex64.
  • 3+ arguments: PyArray_PromoteDTypeSequence runs reduce_dtypes_to_most_knowledgeable to elect a main_dtype at dtypes[0], and then requires NPY_DT_CALL_common_dtype(main_dtype, dtypes[i]) to succeed in one direction (without reverse fallback) for all remaining elements.
    • Because CustomFloat returns Py_NotImplemented for PyComplexDType, PyComplexDType returns Py_NotImplemented for built-in float/complex DTypes (Float16DType, Float32DType, Float64DType, CFloatDType, CDoubleDType), and built-in DTypes return Py_NotImplemented for CustomFloat, they form a deferral cycle:
      $$\text{CustomFloat} \prec \text{PyComplexDType} \prec \text{BuiltinFloat} \prec \text{CustomFloat}$$
    • Every permutation of (custom_float, builtin_float_or_complex, 1j) elects either BuiltinFloat, CustomFloat, or PyComplexDType as main_dtype, which then fails on the third element in the cycle and raises np.exceptions.DTypePromotionError.
    • Similarly, np.result_type(int4, bfloat16, 1j) elects bfloat16 as main_dtype (since int4 handles PyComplexDType directly and defers to bfloat16), and bfloat16.__common_dtype__(PyComplexDType) returns Py_NotImplemented, raising DTypePromotionError.

Suggested Fix

In NPyCustomFloat_CommonDType, handle PyArray_PyComplexDType:

if (other == &PyArray_PyComplexDType) {
  Py_INCREF(&PyArray_CFloatDType);
  return &PyArray_CFloatDType;
}

Reproducing Test Case

@pytest.mark.parametrize("float_t", ALL_CUSTOM_FLOATS)
@pytest.mark.parametrize(
    "concrete, expected",
    [
        (np.float16, np.complex64),
        (np.float32, np.complex64),
        (np.float64, np.complex128),
        (np.complex64, np.complex64),
        (np.complex128, np.complex128),
    ],
)
def test_custom_float_with_python_complex_and_concrete(
    float_t, concrete, expected
):
  for args in itertools.permutations([float_t, 1.0j, concrete]):
    assert np.result_type(*args) == np.dtype(expected), args


@pytest.mark.parametrize("float_t", ALL_CUSTOM_FLOATS)
@pytest.mark.parametrize("int_t", ALL_INTN)
def test_custom_float_with_custom_int_and_python_complex(float_t, int_t):
  for args in itertools.permutations([float_t, int_t, 1.0j]):
    assert np.result_type(*args) == np.dtype(np.complex64), args

Problem 2: Non-Commutative Promotion When a Python Scalar (0, 0.0, 0.0j) Precedes a Custom DType Array or Scalar Instance

Root Cause

All ml_dtypes DTypes are initialized using NPY_DT_legacy_descriptor_proto, which sets DType->flags |= NPY_DT_LEGACY (1) in dtypemeta.c (and npy_2_compat_new_dtypes.h). Consequently, NPY_DT_is_legacy(other) is true and other->type_num >= 256.

In NumPy's abstractdtypes.c, int_common_dtype (PyLongDType), float_common_dtype (PyFloatDType), and complex_common_dtype (PyComplexDType) have a legacy fallback else if (NPY_DT_is_legacy(other)) that does not return Py_NotImplemented:

  • PyLongDType.__common_dtype__(other) calls other.__common_dtype__(UInt8DType) (lines 176–188).
  • PyFloatDType.__common_dtype__(other) calls other.__common_dtype__(HalfDType) (lines 227–239).
  • PyComplexDType.__common_dtype__(other) calls other.__common_dtype__(CFloatDType) (lines 260–272).

Why existing tests in result_type_test.py missed this:

  1. test_python_scalars only tested rt(dtype, scalar) and never rt(scalar, ...) with array or scalar instances.
  2. In array_result_type and PyArray_ResultType, NumPy partitions arguments into descrs[] (DType objects) and arrs[] (arrays, NumPy scalars, and Python scalars) and always places descrs[] before arrs[] in all_DTypes.
  3. Whenever the custom DType operand is an array (np.zeros(1, dtype=dt)) or scalar instance (dt(1)) and the Python scalar is passed first (e.g. np.result_type(0.0, np.zeros(1, dtype=dt)) or in ufuncs like 0.0 + np.zeros(1, dtype=dt)), all_DTypes[0] is the Python abstract DType and all_DTypes[1] is dt. NumPy's abstractdtypes.c legacy probe then runs instead of dt's CommonDType, producing completely different results depending on argument order:
Expression (Custom Array First) Result Expression (Python Scalar First) Result (Bug)
np.result_type(np.zeros(1, f8_e4m3fn), 0) float8_e4m3fn np.result_type(0, np.zeros(1, f8_e4m3fn)) float64 (probed with uint8)
np.result_type(np.zeros(1, bf16), 0.0) bfloat16 np.result_type(0.0, np.zeros(1, bf16)) float32 (probed with float16)
np.result_type(np.zeros(1, f8_e4m3fn), 0.0) float8_e4m3fn np.result_type(0.0, np.zeros(1, f8_e4m3fn)) float16 (probed with float16)
np.result_type(np.zeros(1, i4), 0) int4 np.result_type(0, np.zeros(1, i4)) int16 (probed with uint8)
np.result_type(np.zeros(1, u4), 0) uint4 np.result_type(0, np.zeros(1, u4)) uint8 (probed with uint8)
np.result_type(np.zeros(1, i4), 0.0) float64 np.result_type(0.0, np.zeros(1, i4)) float16 (probed with float16)
np.result_type(np.zeros(1, i4), 0.0j) complex128 np.result_type(0.0j, np.zeros(1, i4)) complex64 (probed with complex64)
np.result_type(np.zeros(1, bc32), 0) bcomplex32 np.result_type(0, np.zeros(1, bc32)) complex64 (probed with uint8)
np.result_type(np.zeros(1, bc32), 0.0) bcomplex32 np.result_type(0.0, np.zeros(1, bc32)) complex64 (probed with float16)
np.result_type(np.zeros(1, bc32), 0.0j) bcomplex32 np.result_type(0.0j, np.zeros(1, bc32)) complex64 (probed with complex64)

Furthermore, for 3 arguments with arrays (e.g., np.result_type(1.0, np.zeros(1, dtype=np.float64), np.zeros(1, dtype=i4))), reduce_dtypes_to_most_knowledgeable compares low=0 (PyFloatDType) with high=2 (i4) via float_common_dtype(PyFloatDType, i4) -> float16 (so PyFloatDType is not swapped with i4), then swaps PyFloatDType with Float64DType, electing Float64DType as main_dtype, which then fails to promote i4 and raises DTypePromotionError.

Reproducing Test Case

@pytest.mark.parametrize(
    "dtype, scalar, expected",
    [
        (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),
        (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),
        (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_with_arrays_commutative(dtype, scalar, expected):
  arr = np.zeros(1, dtype=dtype)
  assert np.result_type(arr, scalar) == np.dtype(expected)
  assert np.result_type(scalar, arr) == 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.0j, np.float32, np.complex64),
        (1.0j, np.float64, np.complex128),
    ],
)
def test_weak_scalar_stays_weak_with_arrays(int_t, scalar, concrete, expected):
  int_arr = np.zeros(1, dtype=int_t)
  concrete_arr = np.zeros(1, dtype=concrete)
  for args in itertools.permutations([int_arr, scalar, concrete_arr]):
    assert np.result_type(*args) == np.dtype(expected), args

Problem 3: NPyIntN_CommonDType Promotes Custom Ints to Non-Numeric Built-In DTypes and New-Style DTypes (StringDType)

Root Cause

In ints.cc, NPyIntN_CommonDType has an unguarded check:

if (!PyTypeNum_ISUSERDEF(other->type_num)) {
  Py_INCREF(other);
  return other;
}

Because PyTypeNum_ISUSERDEF(type) is ((type) >= 256) && ((type) < 512) (ndarraytypes.h), !PyTypeNum_ISUSERDEF(other->type_num) is true for:

  1. np.dtypes.StringDType() (type_num == NPY_VSTRING == 2056) and any new-style PyArrayDTypeMeta_Spec DType (type_num == -1):
    • Because StringDType is non-parametric (NPY_DT_is_parametric is false), np.result_type(int4, np.dtypes.StringDType()) succeeds and returns StringDType() instead of raising DTypePromotionError!
  2. Built-in non-numeric legacy DTypes (NPY_STRING, NPY_UNICODE, NPY_VOID, NPY_DATETIME, NPY_TIMEDELTA):
    • NPyIntN_CommonDType returns StrDType / BytesDType / VoidDType / DateTime64DType / TimeDelta64DType instead of Py_NotImplemented.
    • PyArray_ResultType then attempts PyArray_CastDescrToDType(int4_descr, common_dtype), raising a generic TypeError: cannot cast dtype int4 to <class 'numpy.dtypes.StrDType'> instead of np.exceptions.DTypePromotionError.

Suggested Fix

Restrict the built-in branch in NPyIntN_CommonDType to legacy numeric DTypes:

if (NPY_DT_is_legacy(other) && other->type_num < NPY_NTYPES_LEGACY) {
  if (other->type_num == NPY_BOOL) {
    Py_INCREF(cls);
    return cls;
  }
  if constexpr (std::numeric_limits<T>::is_signed) {
    if (PyTypeNum_ISUNSIGNED(other->type_num)) {
      return PyArray_CommonDType(&PyArray_ByteDType, other);
    }
  }
  if (PyTypeNum_ISNUMBER(other->type_num)) {
    Py_INCREF(other);
    return other;
  }
}

Reproducing Test Case

@pytest.mark.parametrize(
    "custom_t", ALL_CUSTOM_FLOATS + ALL_INTN + ALL_CUSTOM_COMPLEX
)
@pytest.mark.parametrize(
    "incompat_t",
    [
        np.dtype("U8"),
        np.dtype("S8"),
        np.dtype("V8"),
        np.dtype("M8[D]"),
        np.dtype("m8[s]"),
        np.dtypes.StringDType(),
    ],
)
def test_incompatible_dtypes_raise_promotion_error(custom_t, incompat_t):
  with pytest.raises(np.exceptions.DTypePromotionError):
    np.result_type(custom_t, incompat_t)
  with pytest.raises(np.exceptions.DTypePromotionError):
    np.result_type(incompat_t, custom_t)

Problem 4: np.float16 (NPY_HALF) Non-Associativity and Missing PyArray_RegisterCanCast to NPY_HALF

Root Cause

  1. Non-associativity between custom floats that fit in float16:
    • In floats.cc, NPyCustomFloat_CommonDType promotes any custom float T where CustomFloatSafeTo<T, half>() is true (f4, f6_e2m3, f6_e3m2, f8_e3m4, f8_e4m3, f8_e4m3fn, f8_e5m2) with np.float16 to np.float16.
    • However, in floats.h, CommonCustomFloatDTypeImpl promotes two incomparable custom floats (such as (f8_e4m3fn, f8_e5m2) or (f6_e2m3, f6_e3m2)) directly to np.float32, even when both fit in np.float16 (and even though bfloat16 contains both f8_e4m3fn and f8_e5m2!).
    • As a result:
      • np.result_type(np.result_type(f8_e4m3fn, np.float16), f8_e5m2) == np.result_type(np.float16, f8_e5m2) == dtype('float16'), whereas
      • np.result_type(np.result_type(f8_e4m3fn, f8_e5m2), np.float16) == np.result_type(np.float32, np.float16) == dtype('float32').
  2. Missing PyArray_RegisterCanCast to NPY_HALF:
    • In RegisterCustomFloatCasts, PyArray_RegisterCanCast is called for NPY_FLOAT, NPY_DOUBLE, NPY_LONGDOUBLE, NPY_CFLOAT, NPY_CDOUBLE, and NPY_CLONGDOUBLE, but never for NPY_HALF, even when CustomFloatSafeTo<T, half>() is true.
    • Because custom floats have kind == 'V', NumPy marks the T -> float16 cast as NPY_UNSAFE_CASTING.
    • Consequently:
      • np.result_type(f8_e4m3fn, np.float16) is dtype('float16'), yet np.can_cast(f8_e4m3fn, np.float16, casting='safe') and np.can_cast(f8_e4m3fn, np.float16, casting='same_kind') are False.
      • Calling any ufunc or np.concatenate on (f8_e4m3fn, np.float16) promotes the inputs to float16 via default_ufunc_promoter and then fails casting validation (casting='same_kind') with UFuncTypeError: Cannot cast ufunc 'add' input 0 from dtype(float8_e4m3fn) to dtype('float16') with casting rule 'same_kind'.

Suggested Fix

  1. In RegisterCustomFloatCasts, register NPY_HALF when CustomFloatSafeTo<T, half>() holds:
    if constexpr (CustomFloatSafeTo<T, half>()) {
      if (PyArray_RegisterCanCast(CustomFloatType<T>::npy_descr, NPY_HALF,
                                  NPY_NOSCALAR) < 0) {
        return false;
      }
    }
  2. Make CommonCustomFloatDTypeImpl and NPyCustomFloat_CommonDType consistent on whether custom floats that fit in float16 (or bfloat16) promote to float16 or float32.

Reproducing Test Case

@pytest.mark.parametrize(
    "float_t",
    [f4, f6_e2m3, f6_e3m2, f8_e3m4, f8_e4m3, f8_e4m3fn, f8_e5m2],
)
def test_custom_float_safe_to_float16_can_cast_and_ufunc(float_t):
  assert np.result_type(float_t, np.float16) == np.dtype(np.float16)
  assert np.can_cast(float_t, np.float16, casting="safe")
  a = np.ones(2, dtype=float_t)
  b = np.ones(2, dtype=np.float16)
  assert np.add(a, b).dtype == np.dtype(np.float16)
  assert np.concatenate([a, b]).dtype == np.dtype(np.float16)


def test_float8_with_float16_associativity():
  ab_c = np.result_type(np.result_type(f8_e4m3fn, f8_e5m2), np.float16)
  a_bc = np.result_type(f8_e4m3fn, np.result_type(f8_e5m2, np.float16))
  assert ab_c == a_bc

Problem 5: CommonDType Promotes Type Pairs That Have No Registered Cast (or Only unsafe Casts) in casts.cc

Root Cause

  1. (custom_int, float8_e8m0fnu) has no cast registered at all (and float8_e8m0fnu cannot represent 0 or negative numbers):
    • NPyCustomFloat_CommonDType unconditionally returns cls for any IsCustomIntDType(other).
    • However, float8_e8m0fnu is an unsigned exponent-only format ($2^E$, no zero, no negative numbers), and casts.cc registers no cast between any custom int (i1, i2, i4, u1, u2, u4) and float8_e8m0fnu.
    • np.result_type(i4, f8_e8m0) returns dtype(float8_e8m0fnu), but np.can_cast(i4, f8_e8m0, casting='unsafe') is False and np.zeros(1, dtype=i4).astype(f8_e8m0) raises TypeError.
  2. (int4 | uint4, float6_e2m3fn) has no cast registered at all (and other narrow floats cannot represent all int4/uint4 values):
    • casts.cc explicitly omits int4 <-> float6_e2m3fn and uint4 <-> float6_e2m3fn because float6_e2m3fn has max value 7.5 (// int4 -> float6_e2m3fn is not safe and we only register safe casts).
    • Yet np.result_type(i4, f6_e2m3) and np.result_type(u4, f6_e2m3) return dtype(float6_e2m3fn), which cannot even be cast to with casting='unsafe'.
    • Likewise, float4_e2m1fn (max 6.0, digits = 2) cannot represent int4 values -8, -7, 7 or uint4 values 5, 7, 9..15, and float6_e3m2fn, float8_e5m2, float8_e5m2fnuz, float8_e3m4 (digits = 3) cannot represent odd uint4 values 9, 11, 13, 15.
  3. (float8_e8m0fnu, bcomplex32) has no cast registered at all:
    • In complex.cc, NPyCustomComplex_CommonDType returns bcomplex32 when other is float8_e8m0fnu (since CustomFloatSafeTo<float8_e8m0fnu, bfloat16>() is true).
    • However, casts.cc only registers casts between float8_e8m0fnu and bfloat16 (and registers that same cast twice on lines 109 and 110), never registering a cast between float8_e8m0fnu and bcomplex32 (or complex64 / float64).
    • Thus np.result_type(f8_e8m0, bc32) returns dtype(bcomplex32), while np.can_cast(f8_e8m0, bc32, casting='unsafe') is False and np.ones(1, dtype=f8_e8m0).astype(bc32) raises TypeError.
  4. Missing PyArray_RegisterCanCast for cross-custom-type casts breaks ufuncs and np.concatenate under default casting='same_kind':
    • In casts.cc, RegisterTwoWayCustomCast and RegisterOneWayCustomCast only call PyArray_RegisterCastFunc and never PyArray_RegisterCanCast.
    • Because all custom DTypes have kind == 'V', PyArray_AddLegacyWrapping_CastingImpl assigns NPY_UNSAFE_CASTING to every cross-itemsize cast between custom types (e.g. f8_e4m3fn -> bf16, f4 -> bf16, i4 -> bf16, bf16 -> bc32).
    • Whenever np.result_type(a, b) promotes two distinct custom DTypes to a wider custom DType, np.add(arr_a, arr_b) and np.concatenate([arr_a, arr_b]) fail under default casting='same_kind' with TypeError: Cannot cast ... with casting rule 'same_kind'.

Reproducing Test Case

ALL_CUSTOM_DTYPES = ALL_CUSTOM_FLOATS + ALL_INTN + ALL_CUSTOM_COMPLEX


@pytest.mark.parametrize(
    "a, b", list(itertools.combinations(ALL_CUSTOM_DTYPES, 2))
)
def test_all_promoted_custom_pairs_can_cast_and_concatenate(a, b):
  common = np.result_type(a, b)
  assert np.can_cast(a, common, casting="safe"), (a, b, common)
  assert np.can_cast(b, common, casting="safe"), (a, b, common)
  arr_a = np.ones(1, dtype=a)
  arr_b = np.ones(1, dtype=b)
  assert np.concatenate([arr_a, arr_b]).dtype == common
  assert np.add(arr_a, arr_b).dtype == common

Problem 6: NPyIntN_CommonDType Promotes Distinct Custom Ints to int16 Instead of int8, Breaking Associativity

Root Cause

In ints.cc, promoting any two distinct custom ints (int1, int2, int4, uint1, uint2, uint4) returns &PyArray_Int16DType (np.int16), even though:

  • Every custom int has values in $[-8, 15]$, which fits in np.int8 ($[-128, 127]$).
  • ints.cc registers NPY_INT8 as a safe cast target for all six custom int types.
  • np.result_type(int_t, np.int8) returns np.int8 for every custom int int_t (including uint4).

Promoting (i2, i4) or (i4, u4) to np.int16 instead of np.int8 (or i4 for (i2, i4)) breaks associativity with np.int8 and bfloat16:

  • With np.int8:
    • np.result_type(np.result_type(i2, np.int8), i4) == np.result_type(np.int8, i4) == dtype('int8')
    • np.result_type(np.result_type(i2, i4), np.int8) == np.result_type(np.int16, np.int8) == dtype('int16')
  • With bfloat16:
    • np.result_type(i2, np.result_type(i4, bf16)) == np.result_type(i2, bf16) == dtype(bfloat16)
    • np.result_type(np.result_type(i2, i4), bf16) == np.result_type(np.int16, bf16) == dtype('float64')

Reproducing Test Case

@pytest.mark.parametrize("a, b", list(itertools.combinations(ALL_INTN, 2)))
@pytest.mark.parametrize("c", [np.int8, bf16])
def test_custom_int_promotion_associativity(a, b, c):
  ab_c = np.result_type(np.result_type(a, b), c)
  a_bc = np.result_type(a, np.result_type(b, c))
  abc = np.result_type(a, b, c)
  assert ab_c == a_bc == abc, (a, b, c, ab_c, a_bc, abc)

Problem 7: Leftover Safe-Cast Registration from int8/uint8 to Narrow Custom Floats, Duplicate Cast Registration, and Unchecked spec->slots

Root Cause

  1. Bogus PyArray_RegisterCanCast from int8/uint8 to all custom floats in floats.cc:
    • RegisterCustomFloatCasts unconditionally registers NPY_UBYTE (uint8) and NPY_BYTE (int8) as safely castable to every CustomFloatType<T>.
    • Only bfloat16 (8 significand bits, signed, wide exponent range) can represent all int8 and uint8 values. Narrower floats (float4_e2m1fn, float6_*, float8_*) cannot represent all int8/uint8 values (and float8_e8m0fnu cannot represent 0 or negative numbers).
    • As a result, np.can_cast(np.int8, f4, casting='safe') and np.can_cast(np.uint8, f8_e8m0, casting='safe') return True, contradicting np.result_type(np.int8, f4) == np.dtype(np.float64).
  2. Duplicate cast registration in casts.cc:
    • RegisterTwoWayCustomCast<float8_e8m0fnu, bfloat16, float>() already registers both directions (float8_e8m0fnu -> bfloat16 and bfloat16 -> float8_e8m0fnu), so line 110 (RegisterTwoWayCustomCast<bfloat16, float8_e8m0fnu, float>()) is a duplicate registration.
  3. Unchecked spec->slots pointer dereference in npy_2_compat_new_dtypes.h:
    • Line 50 checks if (spec->slots != NULL ...), but lines 61–66 unconditionally dereference slot = spec->slots in while (slot->slot != 0 || slot->pfunc != NULL).

Reproducing Test Case

@pytest.mark.parametrize(
    "float_t", [t for t in ALL_CUSTOM_FLOATS if t is not bf16]
)
@pytest.mark.parametrize("int_t", [np.int8, np.uint8])
def test_int8_uint8_not_safe_to_narrow_custom_floats(float_t, int_t):
  assert not np.can_cast(int_t, float_t, casting="safe")

@hawkinsp

hawkinsp commented Oct 4, 2026 •

Copy link
Copy Markdown
Collaborator

I ran our internal tests against this PR and this behavior change causes at least one test failure:

Before: np.ones((1024, 64), dtype=ml_dtypes.float8_e4m3fn) + 1 preserved dtype float8_e4m3fn
After: np.ones((1024, 64), dtype=ml_dtypes.float8_e4m3fn) + 1 returns float32

This one certainly seems like one that should be fixed: floats should win against a Python int scalar.

AI suggests:

Why it breaks: When evaluating arr_float8 + 1 (or arr_bfloat16 + 1.0), NumPy calls PyArray_CommonDType(PyLongDType, float8_e4m3fn) with the Python scalar's abstract DType first.

Because ml_dtypes uses NPY_DT_legacy_descriptor_proto, NPY_DT_is_legacy(other) is true. In abstractdtypes.c, PyLongDType, PyFloatDType, and PyComplexDType assume any NPY_DT_is_legacy user dtype has no custom common_dtype for abstract scalars, so instead of returning Py_NotImplemented (which would delegate to ml_dtypes's common_dtype), they probe other->common_dtype(other, uint8 / float16 / complex64)—which promotes float8 + uint8 to float64/float16 and bfloat16 + float16 to float32.

The fix:

In NumPy (abstractdtypes.c): In int_common_dtype, float_common_dtype, and complex_common_dtype, inside else if (NPY_DT_is_legacy(other)), call NPY_DT_CALL_common_dtype(other, cls) first and return its result unless it is Py_NotImplemented, only falling back to the uint8/float16/complex64 probe when it returns Py_NotImplemented.

In ml_dtypes (to work on already-released NumPy 2.0–2.4 without waiting for the NumPy fix): At module init in dtypes.cc, wrap the common_dtype slot (((PyArrayDTypeMeta_CommonDType**)dt.dt_slots)[NPY_DT_common_dtype - 1]) on PyArray_PyLongDType, PyArray_PyFloatDType, and PyArray_PyComplexDType so that when other is an ml_dtypes DType, it returns Py_NotImplemented. That forces PyArray_CommonDType to call other->common_dtype(other, cls) (your NPyCustomFloat_CommonDType / NPyCustomInt_CommonDType / NPyCustomComplex_CommonDType).

One small bug in ints.cc: In NPyCustomInt_CommonDType, other == &PyArray_PyComplexDType currently returns NPY_DOUBLE instead of NPY_CDOUBLE.

@seberg

seberg commented Oct 7, 2026

Copy link
Copy Markdown
Contributor Author

Grrr, yeah. Some of those earlier reports seemed a bit like the agent wanted to solve everything (rather than just improve things considerably), but not all and not the last thing.

Yes, I had not adjusted all the safe-casting logic, I had considered it a follow-up, but yeah, maybe this is large enough to just pile on a bit. That said, due to the effect below, I don't think we should change it yet.

Disregarding the AI suggestions (yes, NumPy should not use it's "legacy hack" on the abstract DTypes when downstream explicitly implement CommonDType; but hell will have to freeze over for me to consider monkey-patching)...

I designed this NumPy path explicitly to allow improving things while not breaking existing behavior, because fixing everything at once isn't viable.
So, unfortunately, here we have to cut our losses and stay with bad existing behavior:
ml_dtypes.float8_* promotes with uint8 to give float8_ right now. That is wrong but makes result_type(ml_dtypes.float8_*, 1) do the right thing.
IMO the solution is to just keep it wrong until NumPy 2.6 (probably) returns NotImplemented in its PyInt promoter (because it sees a custom CommonDType on our "legacy" dtype).

I have to check the notes on non-commutative promotion. Many of those seem like embarrassing misses in testing mostly; but there are some limitations on that when it comes to int/uint mixing e.g.. I have logic in NumPy that should make it work within some limitations, but let's see.

@seberg

seberg commented Oct 20, 2026

Copy link
Copy Markdown
Contributor Author

OK... this grew again, there is a lot of agent work there, but I did manually polish the actual common-dtype operations quite a lot as well at some point along the way.

The current state is that:

  • Everything should be correct on NumPy 2.6 (there are two PRs open that will do this). But there are also the in4+uint+float16/bfloat16 tests that are simply xfailed until such a time.
  • Byte-swapping is now implemented, I don't know, it just happened along the way, but is pretty small.
  • I added the fill slow because my agent complained that complex linspace works less well (but it was always broken since it promoted).

The safe casts are fixed up, but I did not attempt to fix the fact that float8 and float4 are marked as safe because they are both kind='f' with an itemsize of 1. Custom cast registrations with ArrayMethods can fix those, but I don't want to do that here.

Another thing I ignored is that type(np.dtype("bfloat16"))() fails on old NumPy versions. That is code that nobody has in practice. We could add a custom __new__ to allow that, but I dunno if it's worthwhile.

@hawkinsp

Copy link
Copy Markdown
Collaborator

Thanks! I'll pull this in and run our tests as well.

@seberg

seberg commented Oct 20, 2026

Copy link
Copy Markdown
Contributor Author

Hmmm, the nightly is failing because I tested with the NumPy fixes applied and the xfail is wrong on dev without those. If we want to push things, could make that either wait for the final release or add a probe.

@hawkinsp

Copy link
Copy Markdown
Collaborator

Happy either way. I know you probably want to get this in, up to you.

@hawkinsp

Copy link
Copy Markdown
Collaborator

From the point of view of our tests this seems ok. I'm tempted to land and iterate at this point.

(There's a merge conflict but I have resolved it internally and can do it in the merge.)

Is this good to go from your perspective?

@seberg

seberg commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor Author

Thanks, I pushed a fix for the nightly failure (one of the PRs is merged already, which made it leaner as well).
I also pushed the merge, which added a few tests for the new dtype (agent, but it seemed to have done a good job).

Is this good to go from your perspective?

Yeah, I think so. As I mentioned there are follow-ups that are good and I am sure we could squabble about details, but I think the direction is sound and the backwards compatibility changes not too bad here. (Of course the promotions are changed, so not a bug fix.)

@hawkinsp

Copy link
Copy Markdown
Collaborator

Would you mind squashing your commits? Thanks!

@seberg
seberg force-pushed the try-new-dtype-hack branch from 8730b84 to 44565ce Compare October 22, 2026 13:36
Register dtypes through NumPy's new-style DType API while retaining
legacy dtype compatibility. To do this the NumPy 2.5 compatibility
headers are vendored (to not require NumPy 2.5+ for compilation).

The major new feature is to use the common-dtype new-style method
to fix many `np.result_type()` results, especially when mixing with
Python scalars. (Some of these fixes need NumPy 2.6+ for full consistency)
Related to this some safe-casting rules are fixed (but not all, this
will require using new-style casts throughout).
Additinoall, byte-swapping and complex-to-complex conversions are now
fixed.
@seberg
seberg force-pushed the try-new-dtype-hack branch from 44565ce to 4603e1d Compare October 22, 2026 13:42
@seberg

seberg commented Oct 22, 2026

Copy link
Copy Markdown
Contributor Author

Yup, done! Got too used to squash merging (or not caring) ;).

@copybara-service
copybara-service Bot merged commit 6b27d48 into jax-ml:main Oct 22, 2026
7 of 8 checks passed
@seberg
seberg deleted the try-new-dtype-hack branch October 23, 2026 07:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants