jax.numpy.full_like

目录

jax.numpy.full_like#

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

创建一个充满指定值的数组,其形状和数据类型与另一个数组相同。

JAX 实现的 numpy.full_like()

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

  • fill_value (ArrayLike) – 用于填充创建数组的标量或数组。

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

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

  • device (xc.Device | Sharding | None) – (可选) DeviceSharding ,创建的数组将被提交到该设备或分片。

返回:

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

返回类型:

Array

示例

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

fill_value 也可以是一个数组,该数组会被广播到指定的形状:

>>> x = jnp.arange(6).reshape(2, 3)
>>> jnp.full_like(x, fill_value=jnp.array([[1], [2]]))
Array([[1, 1, 1],
       [2, 2, 2]], dtype=int32)