pyro-ppl / numpyro / 18989813418

push

github

chore: update `numpyro.util.set_platform` for `JAX_PLATFORM_NAME` and `JAX_PLATFORMS` (#2091)

* chore: update `numpyro.util.set_platform` for `JAX_PLATFORM_NAME` and `JAX_PLATFORMS`

`JAX_PLATFORM_NAME` has been deprecated in favour of `JAX_PLATFORMS` (ref: https://github.com/jax-ml/jax/issues/32980#issuecomment-3463782922), therefore updating the `numpyro.util.set_platform` with deprecation warning and error handling.

* fix: update default values for `flag_name` and `valid_platforms` in `numpyro.util.set_platform` function

* chore: update type hints for `flag_name` and `valid_platforms` in `numpyro.util.set_platform` function

* chore: update comment for valid platforms in `numpyro.util.set_platform` function

1 of 11 new or added lines in 1 file covered. (9.09%)

13704 of 14830 relevant lines covered (92.41%)

1.69 hits per line