jax-ml/jax · error · TypeError

reps length must be equal to the ndim of x, got {len(reps)=}

Error message

reps length must be equal to the ndim of x, got {len(reps)=} and {x.ndim=}.

What it means

lax.tile requires the reps sequence to have exactly one entry per dimension of x. A mismatched length makes the tiling specification ambiguous and is rejected.

Source

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

      x_aval.sharding.mesh.abstract_mesh,
      P(*tuple(s for d in x_aval.sharding.spec for s in [None, d])),
  )
  reshaped_aval = x_aval.update(shape=expand_shape, sharding=expand_sharding)
  reshaped = mlir.reshape(ctx, x, reshaped_aval)
  reshaped = mlir.lower_with_sharding_in_types(ctx, reshaped, reshaped_aval)
  broadcast_shape = tuple(k for pair in zip(reps, x_aval.shape) for k in pair)
  broadcasted_aval = x_aval.update(
      shape=broadcast_shape, sharding=expand_sharding)
  broadcasted = mlir.broadcast_in_dim(ctx, reshaped,
      broadcasted_aval, broadcast_dimensions=tuple(range(2 * x_aval.ndim)))
  broadcasted = mlir.lower_with_sharding_in_types(
      ctx, broadcasted, broadcasted_aval)
  out = mlir.reshape(ctx, broadcasted, aval_out)
  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]

def _tile_abstract_eval(x, reps):
  if x.ndim != len(reps):
    raise TypeError(
        f"reps length must be equal to the ndim of x, got {len(reps)=} "
        f"and {x.ndim=}.")
  for i, (r, sh) in enumerate(zip(reps, x.sharding.spec)):
    if r != 1 and sh is not None:
      raise core.ShardingTypeError(
          f'Operand cannot be sharded on dimension {i} when the tiling is'
          f' non-trivial. Got input type: {x} with reps: {reps}')
  return x.update(shape=tuple(np.multiply(x.shape, reps)))

def _tile_transpose_rule(ct, operand, *, reps):
  if type(ct) is ad_util.Zero:
    return [ad_util.Zero(operand.aval)]
  if not isinstance(operand, ad.UndefinedPrimal):
    return [None]  # transpose wrt literal
  out_spec = tuple(s for sp in operand.aval.sharding.spec for s in [None, sp])
  ct_reshaped = reshape(
      ct, tuple(k for pair in zip(reps, operand.aval.shape) for k in pair),
      out_sharding=operand.aval.sharding.update(spec=out_spec))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Supply one rep per dimension: for a 2-D array pass a length-2 sequence
2. Prefer jnp.tile (which mirrors numpy semantics) if you want numpy-style broadcasting of reps
3. Derive reps from x.ndim at runtime under jit with static shapes

Example fix

// before
x = jnp.zeros((4, 3))
y = lax.tile(x, [2])            # wrong: needs 2 entries
// after
y = lax.tile(x, [2, 1])         # tile dim0 twice, dim1 once

reps = tuple(reps)
if len(reps) < x.ndim: reps = (1,) * (x.ndim - len(reps)) + tuple(reps)
assert len(reps) == x.ndim

def tile_reps_valid(x, reps) -> bool:
    return len(tuple(reps)) == x.ndim

When it happens

Trigger: Calling jax.lax.tile(x, reps) with len(reps) != x.ndim, e.g. tile(x_2d, [3]) or tile(x_1d, [2, 2]).

Common situations: Coming from np.tile habits where numpy broadcasts shorter reps (np.tile(x, 3) is fine but lax.tile requires full length); dynamically-rank operands under jit where x.ndim changed.

Related errors

• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must
• conv_general_dilated window and window_strides must have the
• Wrong number of explicit pads for convolution: expected {},

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