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)];"}

[huhuiqi7]

I am testing the auto sharding function of the CNN network and encountered an error. I am using Python to call the xla_extension library compiled in XLA.

error log:
(MeshHostWorker pid=15170) 2024-07-15 11:02:36.208517: 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)];"}

python code:
compiled = backend.compile(xla_computation_to_mlir_text(hlo.get_computation()), compile_options)

build command:
HERMETIC_PYTHON_VERSION=3.9 HTTP_PROXY=http://10.88.206.35:31166 HTTPS_PROXY=http://10.88.206.35:31166 bazel build --repo_env=HERMETIC_PYTHON_VERSION --test_output=all --spawn_strategy=sandboxed //xla/python:xla_extension --python_path=/root/miniconda3/envs/py39/bin/python --sandbox_debug

environment:
cuda=12.2, nccl=2.19.3-1+cuda12.2, clang=17.0.2, python=3.9.19, Ubuntu=20.04

[huhuiqi7]

Fixed by updating jax version from 0.4.8 to 0.4.10