Flax is a neural network library for JAX that is designed for flexibility.
6.01k
stars
633
forks
source link
Replace `XLACompatibleSharding` with `jax.sharding.Sharding` since the former is deprecated and will be removed in the future. #3967
Closed
copybara-service[bot] closed 3 months ago
Replace
XLACompatibleShardingwithjax.sharding.Shardingsince the former is deprecated and will be removed in the future.