jax-ml/jax · error · TypeError

rule must be a str, but got {type(rule)}

Error message

rule must be a str, but got {type(rule)}

What it means

str_to_sdy_sharding_rule (and therefore def_partition) requires the rule to be a Python str describing operands->results. Any other type (bytes, list, SdyShardingRule object) raises TypeError immediately.

Source

Thrown at jax/_src/custom_partitioning_sharding_rule.py:336

                             **factor_sizes: int) -> SdyShardingRule:
  """Constructs a SdyShardingRule object from the Einsum notation like string.

  This is done by verifying that the input Einsum notation like string and
  with optional special factors and factor sizes represents a valid sharding
  rule and converting it to an internal representation.

  Args:
    rule: The Einsum notation like string for an operation.
    reduction_factors: A tuple of factors that are reduction factors.
    need_replication_factors: A tuple of factors that are need_replication factors.
    permutation_factors: A tuple of factors that are permutation factors.
    **factor_sizes: The optional factor sizes.

  Raises:
    ValueError: If there is any problem with the rule or factor_sizes.
  """
  if not isinstance(rule, str):
    raise TypeError(f"rule must be a str, but got {type(rule)}")
  if not all(isinstance(size, int) for size in factor_sizes.values()):
    raise TypeError(
        f"factor_sizes must be a dict of str to int, but got {factor_sizes}")

  # Replace ... with a single char to simplify parsing.
  if BATCHING in rule:
    raise ValueError(f"Unknown character '{BATCHING}'")
  if "." in rule:
    rule = rule.replace("...", BATCHING)
    if "." in rule:
      raise ValueError("Character '.' must be used inside ellipsis '...'")

  try:
    operands, results = rule.split("->")
  except ValueError as e:
    raise ValueError(f"There is no -> in rule: '{rule}'") from e

  operand_mappings = _parse_values(operands)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert the value to str before calling, e.g. rule.decode() for bytes or '->'.join(parts) for structured forms
2. If you already hold an SdyShardingRule, pass it where the rule object is accepted rather than the str API

Example fix

# before
rule = b'i,j->j'
# after
rule = 'i,j->j'

assert isinstance(rule, str), f'rule must be str, got {type(rule)}'

def is_rule_str(rule) -> bool:
    return isinstance(rule, str)

When it happens

Trigger: Passing rule=b'i,j->j', rule=['i','j'], or an already-parsed SdyShardingRule to def_partition.

Common situations: Loading rules from config/protobuf and forgetting to decode; refactoring from parsed objects back to strings.

Related errors

• Expected unreduced_kind to be of type `jax.sharding.Unreduce
• {name} argument of ManualAxisType should of type `frozenset`
• {kind} must be a tuple of factors
• Compound factors should be one level, nested brackets are no
• Brackets are not balanced

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