mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 20:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
23,091 Commits 504 Branches 263 Tags
Commit Graph

29 Commits

Author SHA1 Message Date
Michael Hudgins
d4d1518c3d Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax 
PiperOrigin-RevId: 676843138
 2024-09-20 07:52:33 -07:00
Jake VanderPlas
568987af23 Finalize deprecation of batched keys to PRNG functions 
PiperOrigin-RevId: 636196573
 2024-05-22 09:40:32 -07:00
Jake VanderPlas
6f7ebff585 random_lax_test: fix kstest for newer NumPy 2024-05-06 09:20:07 -07:00
Jake VanderPlas
c7b7b01e63 remove test of deprecated jax.random.shuffle API 
PiperOrigin-RevId: 623499655
 2024-04-10 08:19:12 -07:00
Jake VanderPlas
31e2358887 test: work around issue with kstest in scipy>1.12 2024-04-03 11:17:56 -07:00
Peter Hawkins
b7401872d5 Disable a random distribution test that appears to fail under scipy 1.13.0rc1. 
PiperOrigin-RevId: 621352047
 2024-04-02 18:11:30 -07:00
Jake VanderPlas
8949a63ce1 [key reuse] rename flag to jax_debug_key_reuse 2024-03-22 05:37:30 -07:00
Roy Frostig
477a5aa148 set a lower pval for vmap-of-rbg based uniform statistical test 
PiperOrigin-RevId: 614434781
 2024-03-10 10:15:06 -07:00
Roy Frostig
29edfd8925 define a loop-free untrue batching rule for rng_bit_generator 2024-03-08 13:13:03 -08:00
Jake VanderPlas
d08e9a03d8 [key reuse] add eager checks 2024-02-29 15:30:19 -08:00
Jake VanderPlas
03ce8ca0ca jax.random: deprecate passing of batched keys to APIs 2024-01-17 12:53:24 -08:00
David Majnemer
e089c9b84f Internal only changes. 
Reverts 4347950d9d018a254fd00bded54ae79df2e71556

PiperOrigin-RevId: 592679160
 2023-12-20 15:55:53 -08:00
Roy Frostig
66e86a36ee decrease sample size of dirichlet test 
Tweak the tolerance threshold instead.

PiperOrigin-RevId: 592051942
 2023-12-18 17:14:09 -08:00
Roy Frostig
3380b9feee split the random generalized normal test and skip its K-S half 
It is key-sensitive and sometimes slow.

PiperOrigin-RevId: 590756597
 2023-12-13 17:01:19 -08:00
Roy Frostig
8dcc079dd8 increase sample sizes on noisy random tests 
... revealed by the partitionable Threefry upgrade; see
https://github.com/google/jax/discussions/18480

PiperOrigin-RevId: 590327421
 2023-12-12 14:00:03 -08:00
David Majnemer
4347950d9d Internal only changes. 
PiperOrigin-RevId: 590071197
 2023-12-11 22:46:15 -08:00
Jake VanderPlas
a52d18781e Add experimental static key reuse checking 2023-12-11 12:03:48 -08:00
David Majnemer
ad14478dd9 Further enhance the seed 
This choice passes more tests for more targets.

PiperOrigin-RevId: 588395200
 2023-12-06 05:47:20 -08:00
David Majnemer
f29ec4ed45 Change seed for a test 
We got unlucky and hit a seed which happens to fail the KS test.

PiperOrigin-RevId: 587885112
 2023-12-04 17:06:17 -08:00
jiayaobo
ae2387dc27 add random.binomial 
update

update

modify
 2023-11-19 14:51:10 +08:00
Jake VanderPlas
5f7335fb55 Deprecate jax.random.shuffle 
This has been long deprecated, but this PR uses the standard deprecation
framework to make it easier to finalize.
 2023-11-06 12:21:56 -08:00
Jake VanderPlas
96d9f89415 [random] better errors for unsupported operations on prng keys 2023-11-03 19:23:18 -07:00
Jake VanderPlas
6da4750c3b [random] remove internal uses of deprecated prng.seed_with_impl() 2023-10-17 13:18:08 -07:00
Sergei Lebedev
cbcaac2756 MAINT Migrate remaining internal/test modules to use state objects 
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.
 2023-10-12 17:32:15 +01:00
Sergei Lebedev
2f70ae700a Migrate another subset of internal modules to use state objects 
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.

PiperOrigin-RevId: 572587137
 2023-10-11 08:46:06 -07:00
Blake Hechtman
a10854786b [JAX] randint goes from [min, max) so [min - N, max + M] actually maps to [min, 
max]

PiperOrigin-RevId: 572034040
 2023-10-09 14:05:39 -07:00
Jake VanderPlas
b407620d1d random_test: fix deprecation warnings for key tests 
Some versions of numpy on some platforms raise warnings when custom PRNG keys
are passed to np.assert_array_equal. Address this by creating a specific function
for comparing key values.
 2023-10-06 13:21:23 -07:00
Jake VanderPlas
3d503e01dc random_test: remove unnecessary test utilities 2023-10-05 15:33:14 -07:00
jax authors
8f911e1512 random_test: Split into two so that each target is small enough to fit within a medium timeout. 
PiperOrigin-RevId: 571146766
 2023-10-05 15:28:51 -07:00