pure_callback: 跨平台分片丢失了自己的错误消息(除了 IndexError 应该是 ValueError)
作者: rogeriojorge创建于 2026年9月16日更新于 2026年9月17日
当 `pure_callback` 为不在包含计算的设备分配中的设备提供 `SingleDeviceSharding` 时, `_callback_op_sharding` 会引发一个解释性的 `ValueError`。它不会: `tuple.index` 会引发 `ValueError`, 而不是 `IndexError`, 因此 `except` 子句永远不匹配,而 bare 信息逃逸。
内容来源: jax-ml/jax