* Add AdamW jit kernel

* Second implementation

* Add missing header

* Correct number of jit kernels in the test

* Added adam kernel

* CI rerun