jax.lax.full_like

Contents

jax.lax.full_like#

jax.lax.full_like(x, fill_value, dtype=None, shape=None, sharding=None)[source]#

Create a full array like np.full based on the example array x.

Parameters:
  • x (ArrayLike | DuckTypedArray) – example array-like, used for shape and dtype information.

  • fill_value (ArrayLike) – a scalar value to fill the entries of the output array.

  • dtype (DTypeLike | None) – optional, a dtype parameter for the output ndarray.

  • shape (Shape | None) – optional, a shape parameter for the output ndarray.

  • sharding (Sharding | None) – an optional sharding specification for the resulting array. If not specified, the output will have the same sharding as the input, with a few exceptions/limitations in particular: 1. Sharding is not available during tracing, thus this will rely on jit. 2. If x is weakly typed or uncomitted, will use default sharding. 3. Shape is not None and is different from x.shape, default will be used.

Return type:

Array

Returns:

An ndarray with the same shape as x with its entries set equal to fill_value, similar to the output of np.full.