jax-ml/jax · error · ValueError

vmem_limit_bytes must be an int: provided with a {type(vmem_

Error message

vmem_limit_bytes must be an int: provided with a {type(vmem_limit_bytes)}.

What it means

vmem_limit_bytes was passed with a non-int type (e.g. float, numpy scalar, string) to the TPU kernel lowering config.

Source

Thrown at jax/_src/tpu_custom_call.py:800

        ctx.module_context.pallas_collective_id_mapping.manual[key] = (
            collective_id
        )
        ctx.module_context.pallas_collective_id_mapping.all_ids.add(
            collective_id
        )

    if collective_id is None:
      raise ValueError(
          "collective_id has to be specified when using a custom barrier "
          "(cannot auto-allocate without lowering context)"
      )
  elif collective_id is not None and not allow_collective_id_without_custom_barrier:
    raise ValueError(
        "collective_id has to be unspecified or None when not using a custom"
        " barrier"
    )
  if vmem_limit_bytes is not None and not isinstance(vmem_limit_bytes, int):
    raise ValueError(
        "vmem_limit_bytes must be an int: provided with a"
        f" {type(vmem_limit_bytes)}."
    )
  if tiling is not None and  device_type != "sparsecore":
    raise ValueError(
        "explicit tiling is only supported for SparseCore kernels."
    )
  if opt_level is not None and device_type != "sparsecore":
    raise ValueError(
        "explicit opt_level is only supported for SparseCore kernels."
    )
  return CustomCallBackendConfig(
      lowered_module_asm,
      lowered_module_asm_version,
      has_communication,
      collective_id,
      device_type,
      cost_estimate,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Coerce to int: vmem_limit_bytes=int(value)
2. Validate config values before passing

Example fix

// before
cfg = dict(vmem_limit_bytes=1e6)
// after
cfg = dict(vmem_limit_bytes=int(1e6))

vmem_limit_bytes = int(vmem_limit_bytes) if vmem_limit_bytes is not None else None

def valid_vmem(v): return v is None or (isinstance(v, int) and not isinstance(v, bool))

When it happens

Trigger: Passing vmem_limit_bytes=1e6 or np.int64(...) or '1024' in compiler params.

Common situations: Reading config from YAML/JSON where numbers parse as float/str; numpy scalars from config math.

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

• {primitive_name}: Buffers with a memory space of HBM or ANY
• Vector clock size ({self.vector_clock_size}) must be greater
• Cannot store scalars to VMEM
• Cannot swap scalars to VMEM.
• load_staged_rhs must be an integer or None.

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