Update activation sharding to shape mesh#63

wonjoo-wj wants to merge 1 commit intooptimize_spmd_shardingpytorch-tpu/llama:optimize_spmd_shardingfrom

wonjoo/optimize_spmd_shardingpytorch-tpu/llama:wonjoo/optimize_spmd_shardingCopy head branch name to clipboard