Closed phu0ngng closed 3 weeks ago
Is this a duplicate with #883?
Hi @denera, This is not a duplicate of #883. I apologize for the very confusing names. I updated the name of the old PR to make it more clear. I also rebased so that the changes in the old PR are not included in the PR. Happy review ^^!
/te-ci jax
/te-ci jax
Description
jax/cpp_extensions.py
into multiple files by category. New files are now in thejax/cpp_extensions/
dir.tex
module.mlp.py
andfused_attn.py
tolayernorm_mlp.py
andattention.py
. (I may move these module files into a newmodules
folder to be consistent with TE/PyTorch).There is no change in the implementation.
Type of change
Checklist: