Scaled dot fixes: 4D tensor handling

[pemeliya]

Cherry-picked from: #940

Also, fixing compilation under rocm 7.13:

external/com_google_absl/absl/base/casts.h:183:23: note: while substituting prior template arguments into non-type template parameter [with Dest = _Float16, Source = uint16_t]
  183 | inline constexpr Dest bit_cast(const Source& source) {
      |                       ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  184 |   return __builtin_bit_cast(Dest, source);
      |   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  185 | }
      | ~
xla/stream_executor/rocm/buffer_debug_float_check_kernel_rocm.cu.cc:45:5: note: while substituting deduced template arguments into function template 'bit_cast' [with Dest = _Float16, Source = uint16_t, $2 = (no value)]
   45 |     absl::bit_cast<_Float16>(uint16_t{0x7C00});
      |     ^

[i-chaochen]

We also have

0.10.0: https://github.com/ROCm/xla/tree/rocm-jaxlib-v0.10.0
0.10.1: https://github.com/ROCm/xla/tree/rocm-jaxlib-v0.10.1

as well.

So it's best to upstream it first and let JAX team to decide which release branches they will pick up