pure_callback: cross-platform sharding loses its own error message (except IndexError should be ValueError)

When pure_callback is given a SingleDeviceSharding for a device that is not in the enclosing computation's device assignment, _callback_op_sharding is meant to raise an explanatory ValueError. It does not: tuple.index raises ValueError, not IndexError, so the except clause never matches and the bare message escapes.

jax/_src/callback.py (0.11.1):

try:
    device_index = device_assignment.index(device)
except IndexError as e:
    raise ValueError(
        "Sharding provided to pure_callback specifies a device"
        f" {device} that is not in the device assignment"
        f" ({device_assignment})") from e

Reproducer

On any machine with a GPU:

import jax
import jax.numpy as jnp

cpu = jax.devices("cpu")[0]

def f(x):
    return jax.pure_callback(
        lambda v: v * 2.0,
        jax.ShapeDtypeStruct(x.shape, x.dtype),
        x,
        sharding=jax.sharding.SingleDeviceSharding(cpu),
    )

jax.jit(f)(jnp.ones(3))

jax 0.11.1, 2x RTX A4000, CUDA driver 580.173.02:

ValueError: tuple.index(x): x not in tuple

Expected: the message the code above writes, naming the device and the assignment.

Why it matters

The traceback contains no frame from the calling library, and the message names neither pure_callback nor a device, so the cause is not recoverable from the error. A library that deliberately runs part of a computation on the host CPU while the caller jits on an accelerator hits this on every GPU machine and on no CPU machine.

except ValueError (or if device not in device_assignment:) restores the intended message.