jax

5a1619fb - Canonicalize ShapeDtypeStruct dtypes at JAX API boundaries using check_and_canonicalize_user_dtype.

Canonicalize ShapeDtypeStruct dtypes at JAX API boundaries using check_and_canonicalize_user_dtype. Previously, `_sds_aval_mapping` called `canonicalize_dtype`, which silently downcast 64-bit dtypes without warning when `enable_x64=False` and ignored `explicit_x64_dtypes="allow"`. * Updates `_sds_aval_mapping` and registers a `canonicalize_value` handler for `ShapeDtypeStruct` using `check_and_canonicalize_user_dtype(..., allow_non_jax_dtypes=True)`. * Updates `ffi.py` so `ffi_call` output dtypes are abstractified under `explicit_x64_dtypes("allow")`. * Fixes unit tests for 32-bit warning compliance. PiperOrigin-RevId: 955934253