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) – 可选地指定输出将提交到的设备。

返回:

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

返回类型:

数组

另请参阅

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