jax.numpy.astype

目录

jax.numpy.astype#

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

将数组转换为指定的数据类型。

JAX 实现 numpy.astype().

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

参数:

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

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

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

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

返回值:

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

返回类型:

Array

参见

示例

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