{"record":{"id":"fcbadb4d1e1e4a19","repo":"jax-ml/jax","slug":"grid-must-consist-of-python-integers-and-jax-array","errorCode":null,"errorMessage":"Grid must consist of Python integers and JAX Arrays: {grid_types}","messagePattern":"Grid must consist of Python integers and JAX Arrays: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":1755,"sourceCode":"    grid: a pallas grid definition.\n    in_specs: input pallas block specs\n    out_specs: output pallas block specs\n    tiling: optional tiling to assume for the refs.\n    dimension_semantics: optional tuple of GridDimensionSemantics (e.g. PARALLEL\n      or ARBITRARY).\n    trace_scopes: optional bool, indicates whether to annotate each region in\n      the pipeline using named_scope.\n    no_pipelining: If True, turns off pipelining and all copies will be made\n      synchronous. This is useful for debugging multiple-buffering related bugs.\n    num_cores: If set, the number of cores to partition the grid over.\n    core_id: If set, the core ID of the current core for partitioning the grid.\n    _explicit_indices: If True, the body will receive the iteration indices as\n      its first argument. This parameter is meant for internal use only.\n  \"\"\"\n\n  if any(not isinstance(d, (int, jax.Array)) for d in grid):\n    grid_types = tuple(type(d) for d in grid)\n    raise ValueError(\n        f\"Grid must consist of Python integers and JAX Arrays: {grid_types}\"\n    )\n  grid, grid_offsets = _partition_grid(grid, dimension_semantics,\n                                       num_cores, core_id)\n\n  num_steps = math.prod(grid)\n  in_specs = _normalize_specs(in_specs)\n  out_specs = _normalize_specs(out_specs)\n  get_buffer_count = lambda spec: (spec.pipeline_mode.buffer_count if\n    (spec is not None and spec.pipeline_mode is not None) else 2)\n  flattened_specs = jax.tree.leaves((in_specs, out_specs))\n  max_buffer_count = max((2, *map(get_buffer_count, flattened_specs)))\n\n  def pipeline(\n      *refs: Any,\n      scratches=None,\n      allocations=None,\n      body_prologue=None,","sourceCodeStart":1737,"sourceCodeEnd":1773,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L1737-L1773","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Coerce all grid entries: grid = tuple(int(g) for g in grid) where values are known","Use jnp.asarray(n) (a jax.Array) only for intentionally dynamic dimensions","Validate grid entries with isinstance(d, (int, jax.Array)) before the call"],"exampleFix":"# before\ngrid=(np.ceil(n / block), block)\n# after\ngrid=(int(np.ceil(n / block)), int(block))","handlingStrategy":"validation","validationCode":"import jax\nassert all(isinstance(d, int) or isinstance(d, jax.Array) for d in grid), \\\n    f'bad grid types: {tuple(type(d) for d in grid)}'","typeGuard":"def grid_valid(grid) -> bool:\n    import jax\n    return all(type(d) is int or isinstance(d, jax.Array) for d in grid)","tryCatchPattern":null,"preventionTips":["Normalize grid with tuple(int(g) for g in grid) when sizes are static","Beware float results from size math (ceil, division)"],"tags":["jax","pallas","grid","validation"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}