mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 15:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
15,932 Commits 508 Branches 263 Tags
Commit Graph

12 Commits

Author SHA1 Message Date
jax authors
42ef649e65 Merge pull request #14475 from hawkinsp:openxla 
PiperOrigin-RevId: 516316330
 2023-03-13 14:04:41 -07:00
Peter Hawkins
172a831219 Switch JAX to use the OpenXLA repository. 2023-03-13 18:38:26 +00:00
Peter Hawkins
ba557d5e1b Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.". 
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.

PiperOrigin-RevId: 476167538
 2022-09-22 12:27:19 -07:00
Peter Hawkins
20475cc55c Add missing header to ducc build file. 2022-08-26 19:30:04 +00:00
jax authors
498fd2083e Merge pull request #12122 from hawkinsp:fft 
PiperOrigin-RevId: 470294824
 2022-08-26 11:32:07 -07:00
Peter Hawkins
b63801b4db Fixes for PocketFFT->ducc migration. 
* Rename modules from pocketfft to ducc.
* Fix up strides at their generation point rather than where they are
  consumed.
 2022-08-26 14:30:03 +00:00
Gordian Edenhofer
024ae47e79 Switch from pocketfft to ducc 
All credit goes to Martin Reinecke <martin@mpa-garching.mpg.de>.
 2022-08-26 13:36:25 +00:00
Peter Hawkins
caad556a03 Update PocketFFT archive to point to Github. 
Bazel has difficulty downloading from the gitlab instance, leading to a warning. Use the Github mirror instead.
 2021-05-19 09:24:17 -04:00
Peter Hawkins
1ed321cbcd Update PocketFFT version to fix crash due to undersized aligned allocations. 2021-03-12 08:46:01 -05:00
Skye Wanderman-Milne
ba4f7880a3 Replace use of TF's third_party_http_archive with http_archive in PocketFFT workspace. 
TF got rid of `third_party_http_archive`, and I'm not sure why we couldn't use Bazel's provided `http_archive` instead.
 2021-02-25 17:29:37 -08:00
Peter Hawkins
9b7fd6329b Add PocketFFT mirror to fix download problems for some Bazel versions. 2020-10-23 19:48:00 -04:00
Peter Hawkins
f58f1ee456 [JAX] Use PocketFFT for FFTs on CPU instead of Eigen. 
PocketFFT is the same FFT library used by NumPy (although we are using the C++ variant rather than the C variant.)

For the benchmark in #2952 on my workstation:

Before:
```
907.3490574884647
max:     4.362646594533903e-08
mean:    6.237288307614869e-09
min:     0.0
numpy fft execution time [ms]:   37.088446617126465
jax fft execution time [ms]:     74.93342399597168
```

After:
```
907.3490574884647
max:     1.9057386696477137e-12
mean:    3.9326737908882566e-13
min:     0.0
numpy fft execution time [ms]:   37.756404876708984
jax fft execution time [ms]:     28.128278255462646
```

Fixes https://github.com/google/jax/issues/2952

PiperOrigin-RevId: 338743753
 2020-10-23 14:20:32 -07:00