jax-ml/jax · error · ValueError

Grid must consist of Python integers and JAX Arrays: {grid_t

Error message

Grid must consist of Python integers and JAX Arrays: {grid_types}

What it means

The pipeline entry point validates that every grid element is a Python int or a jax.Array (dynamic dimension). Floats, numpy scalars, None, or other objects cannot describe a grid, and the error reports the offending types tuple.

Source

Thrown at jax/_src/pallas/mosaic/pipeline.py:1755

    grid: a pallas grid definition.
    in_specs: input pallas block specs
    out_specs: output pallas block specs
    tiling: optional tiling to assume for the refs.
    dimension_semantics: optional tuple of GridDimensionSemantics (e.g. PARALLEL
      or ARBITRARY).
    trace_scopes: optional bool, indicates whether to annotate each region in
      the pipeline using named_scope.
    no_pipelining: If True, turns off pipelining and all copies will be made
      synchronous. This is useful for debugging multiple-buffering related bugs.
    num_cores: If set, the number of cores to partition the grid over.
    core_id: If set, the core ID of the current core for partitioning the grid.
    _explicit_indices: If True, the body will receive the iteration indices as
      its first argument. This parameter is meant for internal use only.
  """

  if any(not isinstance(d, (int, jax.Array)) for d in grid):
    grid_types = tuple(type(d) for d in grid)
    raise ValueError(
        f"Grid must consist of Python integers and JAX Arrays: {grid_types}"
    )
  grid, grid_offsets = _partition_grid(grid, dimension_semantics,
                                       num_cores, core_id)

  num_steps = math.prod(grid)
  in_specs = _normalize_specs(in_specs)
  out_specs = _normalize_specs(out_specs)
  get_buffer_count = lambda spec: (spec.pipeline_mode.buffer_count if
    (spec is not None and spec.pipeline_mode is not None) else 2)
  flattened_specs = jax.tree.leaves((in_specs, out_specs))
  max_buffer_count = max((2, *map(get_buffer_count, flattened_specs)))

  def pipeline(
      *refs: Any,
      scratches=None,
      allocations=None,
      body_prologue=None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Coerce all grid entries: grid = tuple(int(g) for g in grid) where values are known
2. Use jnp.asarray(n) (a jax.Array) only for intentionally dynamic dimensions
3. Validate grid entries with isinstance(d, (int, jax.Array)) before the call

Example fix

# before
grid=(np.ceil(n / block), block)
# after
grid=(int(np.ceil(n / block)), int(block))

import jax
assert all(isinstance(d, int) or isinstance(d, jax.Array) for d in grid), \
    f'bad grid types: {tuple(type(d) for d in grid)}'

def grid_valid(grid) -> bool:
    import jax
    return all(type(d) is int or isinstance(d, jax.Array) for d in grid)

When it happens

Trigger: Passing grid elements like 1.0, np.int64(n), a shape from an external framework, or a string, e.g. grid=(batch_size_float, 128).

Common situations: Sizes computed as floats (e.g. from ceil division or config parsing), numpy integers from tensor shapes, or None from an optional dimension leaking into grid.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• All elements in the grid must be strictly positive, but got
• Axis {axis} is out of bounds for grid {self.grid}
• Sum of sizes {n} must be equal to dimension {axis} of the op
• Cannot pull iota along dimension {dimension} with None block
• Invalid memory space: {memory_space!r}

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