mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 20:06:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax
History
Sharad Vikram 143ed40a78 Add collect_profile script
2022-06-03 17:56:17 -07:00
..
_src
Merge pull request #10937 from carlosgmartin:random_gnormal_ball
2022-06-03 16:44:03 -07:00
example_libraries
feat: refactor code using pyupgrade
2022-05-17 22:14:05 +01:00
experimental
[sparse] Track unique_indices in sparse transform.
2022-06-03 14:51:26 -07:00
image
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
interpreters
[ROCm] Enable XLA python callback on ROCm.
2022-06-03 17:52:34 +00:00
lax
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
lib
Reexport jaxlib.__version as jax.lib.__version__.
2022-04-28 10:25:06 -07:00
nn
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
numpy
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
ops
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
scipy
Change implementation of jax.scipy.linalg.polar() and jax._src.scipy.eigh to use the QDWH decomposition from jax._src.lax.qdwh.
2022-05-12 07:20:52 -07:00
tools
Drop Bazel visibility of //jaxlib.
2022-04-28 13:16:54 -07:00
__init__.py
Merge pull request #9118 from skye:device_context_manager
2022-06-02 10:33:53 -07:00
abstract_arrays.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
ad_checkpoint.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
api_util.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
cloud_tpu_init.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
collect_profile.py
Add collect_profile script
2022-06-03 17:56:17 -07:00
config.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
core.py
Initialize the thread-local compilation context when undefined in new threads.
2022-05-31 12:57:48 -07:00
custom_batching.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
custom_derivatives.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
custom_transpose.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
distributed.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
dlpack.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
dtypes.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
errors.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
flatten_util.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
jaxpr_util.py
Fix or ignore some pytype errors.
2022-05-31 21:25:55 -07:00
linear_util.py
Merge pull request #9118 from skye:device_context_manager
2022-06-02 10:33:53 -07:00
prng.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
profiler.py
Store profiler server as a global variable and add a stop_server function
2022-06-01 17:50:06 -07:00
py.typed
Improve type annotations
2021-01-13 10:26:35 -05:00
random.py
Added random.generalized_normal and random.ball.
2022-06-03 15:11:29 -04:00
stages.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
test_util.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
tree_util.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
util.py
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
version.py
Update version + CHANGELOG for jax 0.3.13 release
2022-05-16 12:17:07 -07:00