jax-ml/jax · error · ValueError

custom_call backend_config unexpected type: {backend_config}

Error message

custom_call backend_config unexpected type: {backend_config}

What it means

When emitting an MLIR stablehlo.custom_call, the backend_config parameter must be one of a few expected types (typically a string/bytes, None, or a DictAttr depending on api_version). This raise is the catch-all for any other Python type passed as backend_config.

Source

Thrown at jax/_src/interpreters/mlir.py:3404

  operands = list(operands)

  if backend_config is None:
    backend_config_attr = ir.StringAttr.get("")
  elif isinstance(backend_config, (str, bytes)):
    backend_config_attr = ir.StringAttr.get(backend_config)
  elif isinstance(backend_config, dict):
    # TODO(necula): it seems that the CustomCallOp constructor requires that
    # backend_config_attr be a string attribute, even though in some cases we
    # need it to be a DictAttr, e.g., for ApproxTopK on TPU.
    # "Verification failed: 'stablehlo.custom_call' op attribute 'backend_config' failed to satisfy constraint: string attribute"
    # To workaround this limitation we first set it to the empty string and we
    # use an unregistered attribute mhlo.backend_config to hold the DictAttr.
    # We must also use api_version=1 to ensure that mhlo.backend_config is
    # handled properly.
    backend_config_attr = ir.StringAttr.get("")
    api_version = 1
  else:
    raise ValueError("custom_call backend_config unexpected type: " + str(backend_config))
  attributes = dict(
      call_target_name=ir.StringAttr.get(call_target_name),
      has_side_effect=ir.BoolAttr.get(has_side_effect),
      backend_config=backend_config_attr,
      api_version=i32_attr(api_version),
      called_computations=ir.ArrayAttr.get(
          [ir.FlatSymbolRefAttr.get(name) for name in called_computations]
      ),
  )
  if operand_output_aliases is not None:
    attributes["output_operand_aliases"] = ir.ArrayAttr.get([
      hlo.OutputOperandAlias.get(
          # if len(result_types) == 1 then the aliasing refers implicitly to
          # the only output.
          output_tuple_indices=[output_idx] if len(result_types) > 1 else [],
          operand_index=input_idx,
          operand_tuple_indices=[],
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass backend_config as a str (serialized, e.g. JSON or protobuf text) or None
2. If passing dict-like config, serialize it to a string yourself rather than passing a Python dict
3. Check the api_version expectations: version 2 expects a string/bytes backend_config

Example fix

# before
backend_config = {'grid': (8, 8)}

# after
import json
backend_config = json.dumps({'grid': (8, 8)})

assert backend_config is None or isinstance(backend_config, (str, bytes)), \
    f'backend_config must be str/bytes/None, got {type(backend_config)}'

def is_valid_backend_config(bc) -> bool:
    return bc is None or isinstance(bc, (str, bytes))

When it happens

Trigger: Registering a custom-call target via ir_generation helpers where backend_config is passed as e.g. an int, list, or arbitrary object instead of a str/bytes/None/dict-attr. Common when hand-writing xla_client custom call rules or extending mlir.py's call_op builders.

Common situations: Third-party JAX extensions building custom calls with non-string backend configs; version upgrades that changed the accepted backend_config types (e.g. the migration to api_version=2 string configs).

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

• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li
• {prim_name} takes a scalar pred as argument, got {pred}
• `compute_on`'s compute_type argument must be a string.
• Cannot interpret value of type {typ} as an abstract array; i

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