jax.numpy.astype

jax.numpy.astype(x, dtype, /, *, copy=False, device=None)

将数组转换为指定的 dtype。

JAX 对 numpy.astype() 的实现。

这是通过 jax.lax.convert_element_type() 实现的，在某些情况下，其行为可能与 numpy.astype() 略有不同。特别是，浮点数到整数和整数到浮点数的转换细节取决于具体的实现。

参数:

• x (ArrayLike) – 要转换的输入数组

• dtype (DTypeLike | None) – 输出 dtype

• copy (bool) – 如果为 True，则始终返回副本。如果为 False（默认值），则仅在必要时返回副本。

• device (xc.Device | Sharding | None | None) – 可选地指定输出将提交到的设备。

返回:

一个与 x 形状相同的数组，包含指定 dtype 的值。

返回类型:

Array

另请参阅

• jax.lax.convert_element_type()：用于 XLA 样式 dtype 转换的较低级别函数。

示例

>>> x = jnp.array([0, 1, 2, 3])
>>> x
Array([0, 1, 2, 3], dtype=int32)
>>> x.astype('float32')
Array([0.0, 1.0, 2.0, 3.0], dtype=float32)

>>> y = jnp.array([0.0, 0.5, 1.0])
>>> y.astype(int)  # truncates fractional values
Array([0, 0, 1], dtype=int32)