jax.numpy.linspace

jax.numpy.linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0, *, device=None)

在区间内返回均匀间隔的数字。

JAX 对 numpy.linspace() 的实现。

参数:

• start (Array | ndarray | bool | number | bool | int | float | complex | TypedNdArray) – 起始值的标量或数组。

• stop (Array | ndarray | bool | number | bool | int | float | complex | TypedNdArray) – 结束值的标量或数组。

• num (int) – 要生成的值得数量。默认为 50。

• endpoint (bool) – 如果为 True（默认），则结果中包含 stop 值。如果为 False，则不包含 stop 值。

• retstep (bool) – 如果为 True，则返回一个 (result, step) 元组，其中 step 是 result 中相邻值之间的间隔。

• axis (int) – 要沿其生成 linspace 的整数轴。默认为零。

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

• dtype (str | type[Any] | dtype | SupportsDType | None)

返回:

• values 是从 start 到 stop 的均匀间隔值的数组。

• step 是相邻值之间的间隔。

返回类型:

一个数组 values，或者如果 retstep 为 True，则是一个元组 (values, step)，其中：

另请参阅

• jax.numpy.arange(): 给定起始点和步长，生成 N 个均匀间隔的值。

• jax.numpy.logspace(): 生成对数间隔的值。

• jax.numpy.geomspace(): 生成几何间隔的值。

示例

0 到 10 之间的 5 个值的列表

>>> jnp.linspace(0, 10, 5)
Array([ 0. ,  2.5,  5. ,  7.5, 10. ], dtype=float32)

0 到 10 之间的 8 个值的列表，不包括结束点

>>> jnp.linspace(0, 10, 8, endpoint=False)
Array([0.  , 1.25, 2.5 , 3.75, 5.  , 6.25, 7.5 , 8.75], dtype=float32)

值列表及其之间的步长

>>> vals, step = jnp.linspace(0, 10, 9, retstep=True)
>>> vals
Array([ 0.  ,  1.25,  2.5 ,  3.75,  5.  ,  6.25,  7.5 ,  8.75, 10.  ],      dtype=float32)
>>> step
Array(1.25, dtype=float32)

多维 linspace

>>> start = jnp.array([0, 5])
>>> stop = jnp.array([5, 10])
>>> jnp.linspace(start, stop, 5)
Array([[ 0.  ,  5.  ],
       [ 1.25,  6.25],
       [ 2.5 ,  7.5 ],
       [ 3.75,  8.75],
       [ 5.  , 10.  ]], dtype=float32)