jax

5cc2dd8a - Add ROCm support to annotate_device_placement backward compatibility test

Add ROCm support to annotate_device_placement backward compatibility test This adds ROCm coverage for device placement annotation backward compatibility, ensuring the annotate_device_placement custom call remains stable across JAX versions.