jax.numpy.full_like

内容

jax.numpy.full_like#

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

创建一个与数组具有相同形状和数据类型的数组,该数组全部填充指定的值。

JAX 实现的 numpy.full_like().

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

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

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

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

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

返回值:

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

返回类型:

数组

示例

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

填充值 也可以是广播到指定形状的数组

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