6 Commits

Author	SHA1	Message	Date
Angelos Katharopoulos	110d9b149d	Layer norm grad fix donation bug (#941) ... * add layer norm grad test * Fix donation bug in layernorm vjp --------- Co-authored-by: Awni Hannun <awni@apple.com>	2024-04-01 06:15:50 -07:00
Angelos Katharopoulos	29221fa238	Implement vjps for some primitives in the fast namespace (#883) ... * Implement rope vjp in terms of rope * RMSNormVJP primitive and kernel * Add LayerNormVJP primitive and kernel	2024-03-26 16:35:34 -07:00
Angelos Katharopoulos	fcda3a0e66	Increase test tolerance for fast.layer_norm (#880)	2024-03-22 12:10:27 -07:00
Angelos Katharopoulos	2225374060	Adds mx.fast.layer_norm (#870)	2024-03-21 13:55:51 -07:00
Awni Hannun	a54f06b16f	Fast RMS Norm (#862) ... * fast rmsnorm * no rms gpu * kernel * fix shared mem * looped rms and donation in softmax * Make the squaring in float32 to avoid underflow * Fix the default StreamOrDevice for rope and rms_norm in fast * nits --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2024-03-21 07:20:54 -07:00
Awni Hannun	ccf1645995	Custom primitive + RoPE fat op (#676) ... * extensions start * rope custom op * fix build * docs + rope benchmark * fix test * Add a Metal kernel for RoPE * Fix position of traditional * transform tests * Move rope computation to float and fix tests * Fix the test and a typo * change to fast * fix no metal build --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2024-02-14 14:04:25 -08:00