jax.numpy.ones_like

jax.numpy.ones_like(a, dtype=None, shape=None, *, device=None)

创建一个与给定数组具有相同形状和数据类型的值为 1 的数组。

JAX 对 numpy.ones_like() 的实现。

参数:

• a (Array | ndarray | bool | number | bool | int | float | complex | DuckTypedArray) – 具有 shape 和 dtype 属性的类数组对象。

• shape (Any | None) – 可选地覆盖创建的数组的形状。

• dtype (str | type[Any] | dtype | SupportsDType | None) – 可选地覆盖创建的数组的数据类型。

• device (Device | Sharding | None) – （可选）Device 或 Sharding，创建的数组将提交到该设备或分片。

返回:

具有指定形状和数据类型的数组，如果指定了设备，则位于指定设备上。

返回类型:

Array

参见

• jax.numpy.empty()

• jax.numpy.zeros_like()

• jax.numpy.ones_like()

• jax.numpy.full_like()

示例

>>> x = jnp.arange(4)
>>> jnp.ones_like(x)
Array([1, 1, 1, 1], dtype=int32)
>>> jnp.ones_like(x, dtype=bool)
Array([ True,  True,  True,  True], dtype=bool)
>>> jnp.ones_like(x, shape=(2, 3))
Array([[1, 1, 1],
       [1, 1, 1]], dtype=int32)