Updated jnp.ceil/floor/trunc to preserve int dtypes

[vfdev-5]

Description:

• Updated jnp.ceil/floor/trunc to preserve int dtypes
• Updated tests
  • For integral dtypes but we can't yet today compare types vs numpy as numpy 2.0.0rc2 is not yet array api compliant in this case (related to ENH: Support integer dtype inputs in rounding functions numpy/numpy#26766):

>>> np.ceil(np.arange(5))
array([0., 1., 2., 3., 4.])
>>> np.arange(5)
array([0, 1, 2, 3, 4])
>>> np.__version__
'2.0.0rc2'

• Removed ceil, floor, trunc implemetnations from array_api and use jnp equivalent functions

Related to #21088

cc @Micky774

[vfdev-5]

@Micky774 can you review this PR please

[Micky774]

Overall looks good, it may be worth adding a quick changelog entry since this will result in different dtypes, what do you think @jakevdp

Either way, @vfdev-5 can you go ahead and push again to trigger CI?

Edit: Do we have test coverage for the integral dtype preservation of jnp.trunc? I'm not sure of the generic unary op check cares

[jakevdp]

Looks good, but I think we also need a changelog entry.

[vfdev-5]

I'll push an update now

[jakevdp]

Looks good!

[jakevdp]

Thanks!

[jakevdp]

It looks like this change is breaking tests of expected output dtypes when tested against older numpy versions. We'll have to account for that in the tests.

[jakevdp]

Looks great, thanks!

[jakevdp]

Actually, I have one more question here: what about booleans? The array API spec doesn't mention this explicitly, but it might be surprising if booleans are promoted to float when integers are not. I've asked the NumPy devs about this at numpy/numpy#26766 (comment).

[vfdev-5]

Actually, I have one more question here: what about booleans? The array API spec doesn't mention this explicitly, but it might be surprising if booleans are promoted to float when integers are not. I've asked the NumPy devs about this at numpy/numpy#26766 (comment).

I asked Array API folks about that and here is the answer on this question:

Those functions are only defined for numerical types. If a library chooses to support boolean dtypes in those functions, that is okay (e.g., for backward compat reasons), but that code is not guaranteed to be portable.

cc @kgryte

[jakevdp]

OK, in that case let's make the change for booleans as well. I find it strange that booleans are silently promoted to float, while integers keep their type. Alternatively, we could choose to promote booleans to int, by calling promote_args_numeric on the inputs.

[vfdev-5]

@jakevdp I included bools in the same route as for integers, let me know if we need to implement the alternative solution you suggested above.

[jakevdp]

Looks good, thanks!

[vfdev-5]

@jakevdp we should remove similarly ceil, floor, trunc from jax/experimental/array_api/_elementwise_functions.py and import them from jnp, right ?

[jakevdp]

@jakevdp we should remove similarly ceil, floor, trunc from jax/experimental/array_api/_elementwise_functions.py and import them from jnp, right ?

Yes, we should do that

[jakevdp]

Looks good!