Skip to content

Commit

Permalink
Add support for copy kwarg in astype to match Array API
Browse files Browse the repository at this point in the history
  • Loading branch information
Micky774 committed Apr 11, 2024
1 parent 36bedee commit 2215021
Show file tree
Hide file tree
Showing 6 changed files with 81 additions and 9 deletions.
6 changes: 6 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,12 @@ Remember to align the itemized text with the first line of an item within a list
* The `device()` method of JAX arrays has been removed, after being deprecated
since JAX v0.4.21. Use `arr.devices()` instead.

* Bug fixes
* {func}`jax.numpy.astype` will now always return a copy when `copy=True`.
Previously, no copy would be made when the output array would have the same
dtype as the input array. This may result in some increased memory usage.
To prevent copying when possible, set `copy=None`. To error when a copy is
required, set `copy=False`.

## jaxlib 0.4.27

Expand Down
6 changes: 4 additions & 2 deletions jax/_src/numpy/array_methods.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,11 +31,13 @@
import numpy as np
import jax
from jax import lax
from jax.sharding import Sharding
from jax._src import core
from jax._src import dtypes
from jax._src.api_util import _ensure_index_tuple
from jax._src.array import ArrayImpl
from jax._src.lax import lax as lax_internal
from jax._src.lib import xla_client as xc
from jax._src.numpy import lax_numpy
from jax._src.numpy import reductions
from jax._src.numpy import ufuncs
Expand All @@ -55,15 +57,15 @@
# functions, which can themselves handle instances from any of these classes.


def _astype(arr: ArrayLike, dtype: DTypeLike) -> Array:
def _astype(arr: ArrayLike, dtype: DTypeLike, copy: bool = True, device: xc.Device | Sharding | None = None) -> Array:
"""Copy the array and cast to a specified dtype.
This is implemented via :func:`jax.lax.convert_element_type`, which may
have slightly different behavior than :meth:`numpy.ndarray.astype` in
some cases. In particular, the details of float-to-int and int-to-float
casts are implementation dependent.
"""
return lax_numpy.astype(arr, dtype)
return lax_numpy.astype(arr, dtype, copy=copy, device=device)


def _nbytes(arr: ArrayLike) -> int:
Expand Down
34 changes: 30 additions & 4 deletions jax/_src/numpy/lax_numpy.py
Original file line number Diff line number Diff line change
Expand Up @@ -2262,17 +2262,43 @@ def _convert_to_array_if_dtype_fails(x: ArrayLike) -> ArrayLike:
In particular, the details of float-to-int and int-to-float casts are
implementation dependent.
""")
def astype(x: ArrayLike, dtype: DTypeLike | None, /, *, copy: bool = True) -> Array:
def astype(x: ArrayLike, dtype: DTypeLike | None,
/, *, copy: bool = True,
device: xc.Device | Sharding | None = None) -> Array:
util.check_arraylike("astype", x)
x_arr = asarray(x)
del copy # unused in JAX
if dtype is None:
dtype = dtypes.canonicalize_dtype(float_)
dtypes.check_user_dtype_supported(dtype, "astype")
if (
issubdtype(x_arr.dtype, complexfloating)
and dtypes.isdtype(dtype, ("integral", "real floating"))
):
warnings.warn(
"Casting from complex to real dtypes will soon raise a ValueError. "
"Please first use jnp.real or jnp.imag to take the real/imaginary "
"component of your input.",
DeprecationWarning, stacklevel=2
)

out = x_arr
# convert_element_type(complex, bool) has the wrong semantics.
if np.dtype(dtype) == bool and issubdtype(x_arr.dtype, complexfloating):
return (x_arr != _lax_const(x_arr, 0))
return lax.convert_element_type(x_arr, dtype)
out = (x_arr != _lax_const(x_arr, 0))
out = _place_array(out, device=device, copy=copy)

# We offer a more specific warning than the usual ComplexWarning so we prefer
# to issue our warning.
with warnings.catch_warnings():
warnings.simplefilter("ignore", ComplexWarning)
return lax.convert_element_type(out, dtype)

def _place_array(x, device=None, copy=None):
# TODO(micky774): Implement in future PRs as we formalize device placement
# semantics
if copy:
return _array_copy(x)
return x


@util.implements(np.asarray, lax_description=_ARRAY_DOC)
Expand Down
20 changes: 18 additions & 2 deletions jax/experimental/array_api/_data_type_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,18 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import builtins
import functools
from typing import NamedTuple
import jax
import jax.numpy as jnp


from jax._src.lib import xla_client as xc
from jax._src.sharding import Sharding
from jax._src import dtypes as _dtypes
from jax.experimental.array_api._dtypes import (
bool, int8, int16, int32, int64, uint8, uint16, uint32, uint64,
float32, float64, complex64, complex128
Expand Down Expand Up @@ -124,8 +129,19 @@ def _promote_types(t1, t2):
raise ValueError("No promotion path for {t1} & {t2}")


def astype(x, dtype, /, *, copy=True):
return jnp.array(x, dtype=dtype, copy=copy)
def astype(x, dtype, /, *, copy: builtins.bool = True, device: xc.Device | Sharding | None = None):
src_dtype = x.dtype if hasattr(x, "dtype") else _dtypes.dtype(x)
if (
src_dtype is not None
and _dtypes.isdtype(src_dtype, "complex floating")
and _dtypes.isdtype(dtype, ("integral", "real floating"))
):
raise ValueError(
"Casting from complex to non-complex dtypes is not permitted. Please "
"first use jnp.real or jnp.imag to take the real/imaginary component of "
"your input."
)
return jnp.astype(x, dtype, copy=copy, device=device)


def can_cast(from_, to, /):
Expand Down
4 changes: 3 additions & 1 deletion jax/numpy/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ from jax._src.typing import (
Array, ArrayLike, DType, DTypeLike,
DimSize, DuckTypedArray, Shape, DeprecatedArg
)
from jax._src.sharding import Sharding
from jax._src.lib import xla_client as xc
from jax.numpy import fft as fft, linalg as linalg
from jax.sharding import Sharding as _Sharding
import numpy as _np
Expand Down Expand Up @@ -115,7 +117,7 @@ def asarray(
) -> Array: ...
def asin(x: ArrayLike, /) -> Array: ...
def asinh(x: ArrayLike, /) -> Array: ...
def astype(a: ArrayLike, dtype: Optional[DTypeLike], /, *, copy: builtins.bool = ...) -> Array: ...
def astype(a: ArrayLike, dtype: Optional[DTypeLike], /, *, copy: builtins.bool = ..., device: xc.Device | Sharding | None = ...) -> Array: ...
def atan(x: ArrayLike, /) -> Array: ...
def atan2(x: ArrayLike, y: ArrayLike, /) -> Array: ...
def atanh(x: ArrayLike, /) -> Array: ...
Expand Down
20 changes: 20 additions & 0 deletions tests/lax_numpy_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -3840,6 +3840,26 @@ def testAstypeBool(self, from_dtype, use_method, to_dtype='bool'):
self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
self._CompileAndCheck(jnp_op, args_maker)

@jtu.sample_product(
change_dtype=[True, False],
copy=[True, False],
)
def testAstypeCopy(self, change_dtype, copy):
dtype = 'float32' if change_dtype else 'int32'
expect_copy = change_dtype or copy
x = jnp.arange(5, dtype='int32')
y = x.astype(dtype, copy=copy)

assert y.dtype == dtype
y.delete()
assert x.is_deleted() != expect_copy

def testAstypeComplexDowncast(self):
x = jnp.array(2.0+1.5j, dtype='complex64')
msg = "Casting from complex to non-complex dtypes will soon raise "
with self.assertWarns(DeprecationWarning, msg=msg):
x.astype('float32')

def testAstypeInt4(self):
# Test converting from int4 to int8
x = np.array([1, -2, -3, 4, -8, 7], dtype=jnp.int4)
Expand Down

0 comments on commit 2215021

Please sign in to comment.