Commit Graph

6 Commits

Author SHA1 Message Date
Richard Wong 429e1742ab Feat: flax pjit example 2024-09-16 12:19:07 +09:00
Richard Wong ad5cf7735f Feat: fsdp demo
Refactor: pulling dataloader code into dataload.py
2024-09-15 22:41:00 +09:00
Richard Wong 005a1a5735 Feat: learn flax 2024-09-14 14:13:38 +09:00
Richard Wong d2dd72227f Feat: introduced efficient train data dtype, jit train step, bfloat16
mat mul
2024-09-14 02:02:45 +09:00
Richard Wong edd9c3551f Feat: implement working prediction 2024-09-12 22:57:19 +09:00
Richard Wong f523560141 Feat: jax implementation of t5 training and prediction 2024-09-11 08:17:02 +09:00