jax.Array.astype

abstract Array.astype(dtype, copy=False, device=None)

复制数组并转换为指定的数据类型。

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

参数:

• self (Array)

• dtype (DTypeLike | None)

• copy (bool)

• device (xc.Device | Sharding | None)

返回类型:

Array