xla runtime error: Cannot remove instruction: %all-reduce

[jing-alice]

when i use xla client compile instructions, the following issues occurred:
E external/xla/xla/status_macros.cc:54] INTERNAL: RET_CHECK failure (external/xla/xla/hlo/ir/hlo_computation.cc:318) IsSafelyRemovable(instruction) Cannot remove instruction: %all-reduce = f32[32]{0} all-reduce(f32[32]{0} %reduce.3), channel_id=1, replica_groups={{0}}, to_apply=%region_0.44, metadata={op_name="parallelize(train_step_shard_parallel)/jit(main)/reduce_sum[axes=(0, 1, 2)];"}
What might be causing this problem?

jax==0.4.8
jaxlib==0.4.7+cuda12.cudnn88