repo tracking my jax experiments

Richard Wong 0762c02b31 Feat: implemented attention layer in equinox		2024-10-06 23:52:42 +09:00
equinox	Feat: implemented attention layer in equinox	2024-10-06 23:52:42 +09:00
learn_flax	Feat: learn flax	2024-09-14 14:13:38 +09:00
nnx	Feat: implemented attention layer in equinox	2024-10-06 23:52:42 +09:00
parallel	Feat: increased learning rate for effective large batch size learning	2024-09-22 22:28:41 +09:00
t5_model	Feat: increased learning rate for effective large batch size learning	2024-09-22 22:28:41 +09:00
.gitignore	Feat: increased learning rate for effective large batch size learning	2024-09-22 22:28:41 +09:00
check_time.py	Feat: implement working prediction	2024-09-12 22:57:19 +09:00
dataload.py	Feat: implemented attention layer in equinox	2024-10-06 23:52:42 +09:00
make_context_data.py	Feat: implemented attention layer in equinox	2024-10-06 23:52:42 +09:00
requirements.yaml	Feat: jax implementation of t5 training and prediction	2024-09-11 08:17:02 +09:00
t5_jax.py	Feat: increased learning rate for effective large batch size learning	2024-09-22 22:28:41 +09:00
t5_jax_parallel.py	Feat: t5_jax_simple_parallel implements a working example of fsdp	2024-09-20 23:42:51 +09:00
t5_jax_prediction.py	Feat: increased learning rate for effective large batch size learning	2024-09-22 22:28:41 +09:00
t5_jax_retrieval.py	Feat: jax implementation of t5 training and prediction	2024-09-11 08:17:02 +09:00
t5_jax_sfp_grad_accumulate.py	Feat: increased learning rate for effective large batch size learning	2024-09-22 22:28:41 +09:00
t5_jax_shmap.py	Feat: increased learning rate for effective large batch size learning	2024-09-22 22:28:41 +09:00
t5_jax_simple_parallel.py	Feat: increased learning rate for effective large batch size learning	2024-09-22 22:28:41 +09:00
t5_prediction_old.py	Feat: t5_jax_simple_parallel implements a working example of fsdp	2024-09-20 23:42:51 +09:00
t5_summarizer_flax.py	Feat: jax implementation of t5 training and prediction	2024-09-11 08:17:02 +09:00