jax-ml/jax · error · TypeError

full must be called with scalar fill_value, got fill_value.s

Error message

full must be called with scalar fill_value, got fill_value.shape {}.

What it means

jax.lax.full requires fill_value to be a scalar (a 0-d value). If np.shape(fill_value) is non-empty, JAX raises this TypeError because full broadcasts a single scalar into a shape, unlike numpy.full which accepts array fill values.

Source

Thrown at jax/_src/lax/lax.py:3622


def full(shape: Shape, fill_value: ArrayLike, dtype: DTypeLike | None = None, *,
         sharding: Sharding | None = None) -> Array:
  """Returns an array of `shape` filled with `fill_value`.

  Args:
    shape: sequence of integers, describing the shape of the output array.
    fill_value: the value to fill the new array with.
    dtype: the type of the output array, or `None`. If not `None`, `fill_value`
      will be cast to `dtype`.
    sharding: an optional sharding specification for the resulting array,
      note, sharding will currently be ignored in jitted mode, this might change
      in the future.
  """
  shape = canonicalize_shape(shape)
  if np.shape(fill_value):
    msg = "full must be called with scalar fill_value, got fill_value.shape {}."
    raise TypeError(msg.format(np.shape(fill_value)))
  if dtype is None:
    weak_type = dtypes.is_weakly_typed(fill_value)
    fill_dtype = _dtype(fill_value)
  else:
    if isinstance(dtype, dtypes.ExtendedDType):
      return dtype._rules.full(shape, fill_value, dtype)
    weak_type = False
    fill_dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "full")
  fill_value = _convert_element_type(fill_value, fill_dtype, weak_type)
  if (sharding is not None and
      isinstance(fill_value, array.ArrayImpl) and sharding._is_concrete):
    broadcast_shape = sharding.shard_shape(shape)
    shard = broadcast(fill_value, broadcast_shape)
    shard = shard.addressable_data(0)
    return array.make_array_from_callback(
        shape, sharding, lambda _: shard, dtype=fill_dtype)

  if sharding is not None and not sharding._is_concrete:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pass a scalar fill_value: extract it with .item() or float(...) if it is a 1-element array
  2. If you need an array broadcast into a shape, use jnp.broadcast_to(arr, shape) or jnp.where with a mask instead of lax.full
  3. Wrap incoming values with jnp.asarray(fill_value).reshape(()) when they are guaranteed single-element

Example fix

// before
julia = lax.full((3, 3), jnp.array([7.0]))
// after
julia = lax.full((3, 3), 7.0)
# or broadcast an array:
# arr = jnp.broadcast_to(jnp.array([7.0]), (3, 3))
Defensive patterns

Strategy: validation

Validate before calling

import numpy as np
if np.ndim(fill_value) != 0:
    fill_value = np.asarray(fill_value).reshape(-1)[0]  # or raise
out = lax.full(shape, fill_value)

Type guard

def is_scalar(v) -> bool:
    return np.ndim(v) == 0

Prevention

When it happens

Trigger: Calling lax.full(shape, fill_value) (or lax.zeros/ones helpers built on it) with an array/tensor fill_value, e.g. lax.full((3,3), jnp.array([1,2])); also iota/eye-style helpers that route through full.

Common situations: Porting numpy code np.full(shape, some_array) to JAX; passing a computed vector (e.g. a per-channel constant) where a scalar was assumed; dynamically typed fill_value from user config.

Related errors


AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/0f745ed526c48884. Report an issue: GitHub.