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. Tests are added to cover various scenarios, including error conditions and jit compilation.
PiperOrigin-RevId: 863699668
