jit-ing a shard_map throws error - XlaRuntimeError: INTERNAL: RET_CHECK failure

[rupeshknn]

jit-ing a shard_map throws error:

XlaRuntimeError: INTERNAL: during context [hlo verifier]: RET_CHECK failure (external/xla/xla/service/hlo_verifier.cc:2852) check_inst->sharding() == common_sharding_inst->sharding() Inconsistent while sharding among instructions.

From the language in JEP/14273 it seemed like shard_map was jit composable. Is this not the case?

Full error

XlaRuntimeError: INTERNAL: during context [hlo verifier]: RET_CHECK failure (external/xla/xla/service/hlo_verifier.cc:2852) check_inst->sharding() == common_sharding_inst->sharding() Inconsistent while sharding among instructions: 
%while.2 = (s64[], f64[1]{0}, f64[], f64[1]{0}, f64[], /*index=5*/f64[27,27]{1,0}, f64[27,27]{1,0}, f64[27,27]{1,0}, f64[27,27]{1,0}, f64[], /*index=10*/f64[1]{0}, f64[]) while((s64[], f64[1]{0}, f64[], f64[1]{0}, f64[], /*index=5*/f64[27,27]{1,0}, f64[27,27]{1,0}, f64[27,27]{1,0}, f64[27,27]{1,0}, f64[], /*index=10*/f64[1]{0}, f64[]) %tuple.9), condition=%region_8.819, body=%region_7.761, sharding={{replicated}, {manual}, {manual}, {manual}, {manual}, /*index=5*/{manual}, {manual}, {manual}, {manual}, {manual}, /*index=10*/{manual}, {manual}}, metadata={op_name="jit(runthisfunc)/jit(main)/jit(shmap_body)/jit(multi_cap_func)/jit(letstry_alpha)/jit(update)/jit(<unnamed wrapped function>)/while[cond_nconsts=8 body_nconsts=3]" source_file="/home/rupeshknn/uw_research/twoq/tunable_coupler.py" source_line=335}
%arg_tuple.820 = (s64[], f64[1]{0}, f64[], f64[1]{0}, f64[], /*index=5*/f64[27,27]{1,0}, f64[27,27]{1,0}, f64[27,27]{1,0}, f64[27,27]{1,0}, f64[], /*index=10*/f64[1]{0}, f64[]) parameter(0), sharding={{manual}, {manual}, {manual}, {manual}, {manual}, /*index=5*/{manual}, {manual}, {manual}, {manual}, {manual}, /*index=10*/{manual}, {manual}}

Failed after sharding-propagation

source line 335 calls a jaxopt solver (gradient descent optimization)

running without jit-ing works fine. Any ideas guys?

[mattjj]

Thanks for the question!

From the language in JEP/14273 it seemed like shard_map was jit composable. Is this not the case?

We use shard_map inside jit all the time. But this looks like an XLA internal bug. Can you share a minimal repro?