Adds jax.lax.scaled_dot for scaled dot products.

This change introduces a new `jax.lax.scaled_dot` function, which computes a dot product where the inputs can be float8 types. It produces the Composite op that could be lowered to the triton, cuBLAS, cuDNN, or rewritten as the regular dot.

The fallback: If the scaled-dot is not enabled the composite call gets inlined as a sequence of the ops: the float8 inputs and scales are converted to bfloat16, the scales broadcasted and then multiplied with the corresponding operands elementwise, then passed to `jax.lax.dot_general`.

The function includes input validation for shapes and dtypes, and supports both 2D and 3D (batched) dot products. Tests are added to cover various scenarios, including error conditions and jit compilation.

PiperOrigin-RevId: 821573195

Author: loislo

jax/_src/lax/__init__.py

@@ -314,3 +314,4 @@
 from jax._src.lax.ann import (
   approx_top_k_p as approx_top_k_p
 )
+from jax._src.lax.scaled_dot import scaled_dot as scaled_dot