jax.numpy.astype

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

将数组转换为指定的 dtype。

numpy.astype() 的 JAX 实现。

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

参数:

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

• dtype (DTypeLike | None) – 输出数据类型

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

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

返回:

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

返回类型:

数组

另请参阅

• jax.lax.convert_element_type(): 用于 XLA 风格数据类型转换的较低级别函数。

示例

>>> 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)