Use jax.lax.scaled_dot for MXFP dot_general in Qwix instead of jax.nn.scaled_matmul that binds to cuDNN - #369

Use jax.lax.scaled_dot for MXFP dot_general in Qwix instead of jax.nn.scaled_matmul that binds to cuDNN#369

copybara-service[bot] wants to merge 1 commit into

maingoogle/qwix:mainfrom

test_981031906google/qwix:test_981031906Copy head branch name to clipboard