jax
f210b0f9 - Add a flag `jax_require_devices_during_lowering` to control if physical devices are passed during lowering to stablehlo. This is temporary to unblock nvidia.

Commit

2 years ago

Add a flag `jax_require_devices_during_lowering` to control if physical devices are passed during lowering to stablehlo. This is temporary to unblock nvidia. PiperOrigin-RevId: 590318918

References

#18945 - Add a flag `jax_require_devices_during_lowering` to control if physical devices are passed during lowering to stablehlo. This is temporary to unblock nvidia.

Author

yashk2810

Committer

a-googler

Parents

1851447b

jax f210b0f9 - Add a flag `jax_require_devices_during_lowering` to control if physical devices are passed during lowering to stablehlo. This is temporary to unblock nvidia.

jax
f210b0f9 - Add a flag `jax_require_devices_during_lowering` to control if physical devices are passed during lowering to stablehlo. This is temporary to unblock nvidia.