jax.numpy.arange#

jax.numpy.arange(start, stop=None, step=None, dtype=None, *, device=None, out_sharding=None)[source]#

创建一个等间隔数值数组。

numpy.arange() 在 JAX 中的实现，基于 jax.lax.iota()。

与 Python 的 range() 函数类似，它可以以几种不同的位置参数签名调用：

  • jnp.arange(stop)：生成从 0 到 stop（不含）的值，步长为 1。

  • jnp.arange(start, stop)：生成从 start 到 stop（不含）的值，步长为 1。

  • jnp.arange(start, stop, step)：生成从 start 到 stop（不含）的值，步长为 step。

与 Python 的 range() 函数一样，起始值是包含的，而终止值是不包含的。

参数:

  • start (ArrayLike | DimSize) – 区间的起始值，包含在内。

  • stop (ArrayLike | DimSize | None) – 可选的区间结束值，不包含在内。如果未指定，则 (start, stop) = (0, start)

  • step (ArrayLike | None) – 可选的区间步长。默认值为 1。

  • dtype (DTypeLike | None) – 返回数组的可选数据类型；如果未指定，将通过 start、stop 和 step 的类型提升来确定。

  • device (xc.Device | Sharding | None) – （可选）将创建的数组提交到的 Device 或 Sharding。

  • out_sharding (NamedSharding | P | None) – （可选）将创建的数组提交到的 NamedSharding 或 P。如果使用显式分片（https://jax.net.cn/en/latest/notebooks/explicit-sharding.html），请使用 out_sharding 参数。

返回:

从 start 到 stop（不含），以 step 为间隔的等距数值数组。

返回类型:

数组

注意

由于浮点数误差的累积，使用浮点数 step 参数调用 arange 可能会导致意想不到的结果，特别是在使用 float8_* 和 bfloat16 等低精度数据类型时。为避免精度错误，请考虑生成一个整数范围，然后将其缩放到所需的范围。例如，不要使用这种方式：

jnp.arange(-1, 1, 0.01, dtype='bfloat16')

生成整数序列并对其进行缩放会更精确

(jnp.arange(-100, 100) * 0.01).astype('bfloat16')

示例

单参数版本仅指定 stop 值

>>> jnp.arange(4)
Array([0, 1, 2, 3], dtype=int32)

传入浮点数 stop 值会得到浮点数结果

>>> jnp.arange(4.0)
Array([0., 1., 2., 3.], dtype=float32)

双参数版本指定 start 和 stop，且 step=1

>>> jnp.arange(1, 6)
Array([1, 2, 3, 4, 5], dtype=int32)

三参数版本指定 start、stop 和 step

>>> jnp.arange(0, 2, 0.5)
Array([0. , 0.5, 1. , 1.5], dtype=float32)

另请参阅