jax-ml/jax · error · ValueError

batch dimension sizes must match; got {sorted_arr_aval.shape

Error message

batch dimension sizes must match; got {sorted_arr_aval.shape[:batch_dims]} != {query_aval.shape[:batch_dims]}

What it means

Raised by the SearchSorted HiJAX primitive when the leading batch dimensions of sorted_arr and query differ. When batch_dims > 0, both arrays must carry identical leading shape[0:batch_dims] so the primitive can vmap over them consistently.

Source

Thrown at jax/_src/numpy/hijax.py:92

      raise ValueError(
          f"dimension={dimension} must be in range [{batch_dims},"
          f" {sorted_arr_aval.ndim})"
      )
    if sorted_arr_aval.dtype != query_aval.dtype:
      raise ValueError(
          "dtypes of sorted_arr and query must match; got "
          f"{sorted_arr_aval.dtype} and {query_aval.dtype}"
      )
    if side not in ["left", "right"]:
      raise ValueError(
          f"invalid argument side={side!r}, expected 'left' or 'right'"
      )
    if method not in self.valid_methods:
      raise ValueError(
          f"invalid argument {method=}, expected one of {list(self.valid_methods)}"
      )
    if sorted_arr_aval.shape[:batch_dims] != query_aval.shape[:batch_dims]:
      raise ValueError(
          "batch dimension sizes must match; got"
          f" {sorted_arr_aval.shape[:batch_dims]} != {query_aval.shape[:batch_dims]}"
      )
    if not dtypes.issubdtype(out_dtype, np.integer):
      raise ValueError(f"out_dtype should be an integer type; got {out_dtype}")
    # Attempt this here to catch overflow errors early.
    out_dtype.type(sorted_arr_aval.shape[dimension])
    self.in_avals = (sorted_arr_aval, query_aval)
    self.out_aval = core.typeof(api.eval_shape(
      functools.partial(_searchsorted_impl,
        dimension=dimension, batch_dims=batch_dims, side=side,
        dtype=out_dtype, method=method),
        sorted_arr_aval, query_aval))
    self.params = dict(
      side=side,
      dimension=dimension,
      batch_dims=batch_dims,
      method=method,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Align batch shapes: query = jnp.broadcast_to(query, sorted_arr.shape[:batch_dims] + query.shape[batch_dims:]) or reshape the sorted array
Check which operand gained/lost a batch dim under vmap and use in_axes appropriately
Print sorted_arr.shape[:batch_dims] and query.shape[:batch_dims] right before the call

Example fix

# before
out = searchsorted_batched(sorted_arr, query)  # (8,100) vs (16,)
# after
query = jnp.broadcast_to(query[None], (8, query.shape[0]))
out = searchsorted_batched(sorted_arr, query)

Defensive patterns

Strategy: validation

Validate before calling

assert sorted_arr.shape[:batch_dims] == query.shape[:batch_dims], (sorted_arr.shape, query.shape)

Type guard

def batch_shapes_match(a, q, k: int) -> bool:
    return a.shape[:k] == q.shape[:k]

Prevention

Print both shapes before batched search calls
Use jnp.broadcast_to on the smaller operand when batch dims legitimately differ
Be explicit with vmap in_axes when only one operand is batched

When it happens

Trigger: Calling batched searchsorted (batch_dims=k) where, e.g., sorted_arr.shape=(B1, N, M) and query.shape=(B2, M) with B1 != B2; often caused by mismatched leading axes after reshaping or vmap.

Common situations: Using vmap/pmap where one operand got an extra or different batch axis; a reshape that merged or split the batch dim in only one operand.

Related errors

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