Commit 8f2773b

authored and

committed

Migrate Pipeline to use NNX module.

1 parent 6ce880e commit 8f2773bCopy full SHA for 8f2773b

6 files changed

+2818

-590

lines changed

src/MaxText/layers
tests
- pipeline_parallelism_test.py
- pipeline_parallelism_test_linen.py

6 files changed

+2818

-590

lines changed

Comments

(0)