`@parameterized.parameters` interacts very badly with `JaxTestCase.rng`:
the RNG seed is derived from the test method and the test method name
can change if additional test cases are inserted before it. This can cause
CI failures in functions that are completely unrelated to the change that
introduces the breakage.
We should seriously reconsider this strategy. Either all instances of
`parameters` + `self.rng` should be removed or we should find an alternative
strategy for seeding.
PiperOrigin-RevId: 589798050
The motivation here is to gradually replace all dynamic lookups on `jax.config`
with statically-typed state objects, which are more type checker/IDE friendly.
This is a follow up to #18008.
This change prepares for allowing more flexible tag matching. For example, we may want to write "gpu" in a test and have it match both "cuda" and "rocm" devices, which we cannot do under the current API but can easily do under this design.
Replace uses of device_under_test() in a context that performs an equality test with a call to test_device_matches().
Replace uses of if_device_under_test() with test_device_matches() and delete if_device_under_test().
PiperOrigin-RevId: 568923117
Their behavior is the same as functions in scipy.special. The only small
difference is in rel_entr function, which unlike scipy.special does not
take the optional parameter 'out'.
Resolves#16630
- Add gamma fcn api in scipy.special
- Add tests for this purpose
- Add function to the docs
Currently, there is no implementation of the gamma function in jax
but there is one in scipy.special. This breaks some higher level
jit-compilation like in the blackjax backend for pymc. This commit
adds the missing gamma function.
Resolves: #15409
The goal is to ensure that all shards fit into a medium timeout in sanitizer
configurations.
Running 256 entry vectors in spectral_dac is too slow, so let's replace that
with a smaller vector that isn't a power of 2. Avoiding a power of 2 requires
us to widen the tolerance a bit due to vectorization changes.
While here, specify deps a little more precisely as well.
PiperOrigin-RevId: 514440062
The spectral_dac tests are also shrunk because running the full suite on 256-entry vectors is too slow.
This allows them to run in ASAN in more situations.
While here, specify deps a little more precisely as well.
PiperOrigin-RevId: 511829646
