jax.numpy.ones_like#

jax.numpy.ones_like(a, dtype=None, shape=None, *, device=None)[源代码]#

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

numpy.ones_like() 的 JAX 实现。

参数:
  • a (类数组 | 鸭子类型数组) – 具有 shapedtype 属性的类数组对象。

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

  • dtype (DTypeLike | None | None) – 可选地覆盖所创建数组的数据类型。

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

返回:

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

返回类型:

数组

示例

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