flax

Fix shard_map documentation link in compilation.py

#5038

Merged