flax
a8bc3673 - Replaces pjit with jit in spmd.py

Commit
2 years ago
Replaces pjit with jit in spmd.py PiperOrigin-RevId: 574373680
Author
Committer
Parents
Loading