jax.numpy.full_like

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

创建一个与给定数组具有相同形状和数据类型的数组，并用指定值填充。

JAX 实现的 numpy.full_like()。

参数：

• a (类似数组 | 鸭子类型数组) – 具有 shape 和 dtype 属性的类数组对象。

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

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

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

• device (xc.Device | Sharding | None | None) – (可选) Device 或 Sharding，用于指定所创建数组将被提交到的设备。

返回：

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

返回类型：

数组

另请参阅

• jax.numpy.full()

• jax.numpy.empty_like()

• jax.numpy.zeros_like()

• jax.numpy.ones_like()

示例

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