mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
23,091 Commits 500 Branches 263 Tags
Commit Graph

51 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
09fd345de9 pre-commit: update hooks & pin using hashes 2024-08-27 15:23:13 -07:00
Peter Hawkins
7f4ef63cd8 Run pyupgrade --py310-plus. 
Also apply manual fixes to import sorting and unused imports.
 2024-06-26 16:10:18 -04:00
jax authors
fc1e1d4a65 Add freshness metablock to JAX OSS docs. 
PiperOrigin-RevId: 645508135
 2024-06-21 14:50:49 -07:00
Sai-Suraj-27
29def4eefa Updated all the pre-commit hooks versions. 2024-04-08 00:59:02 +05:30
Jake VanderPlas
10eae3f93a CI: update jupytext version 2024-01-09 14:34:21 -08:00
Jan Hrček
4da56dcdd7 Fix duplicate word occurrences 2023-12-19 06:15:30 +01:00
Jake VanderPlas
389eb97a7c CI: update pre-commit hooks to latest version 2023-10-30 09:12:24 -07:00
Peter Hawkins
2c32660a8f Replace references to DeviceArray with Array. 
A number of stale references are lurking in our documentation.
 2023-08-18 17:46:00 -04:00
Jake VanderPlas
7bb8312f82 CI: update jupytext to v0.14.7 2023-07-24 11:51:45 -07:00
Peter Hawkins
319ab98980 Apply pyupgrade --py39-plus. 
Notable changes:
* use PEP 585 type names
* use PEP 604 type union syntax where `from __future__ import annotations` is present.
* use f-strings in more places.
* remove redundant arguments to open().
 2023-07-21 14:49:44 -04:00
Kevin Millikin
921fd222bf Refer to the original map/zip classes via builtins 
Referring to them as simply `map` or `zip` will create recursive
reimplementations (with no base case!) if the cell is reevaluated in
the same runtime.
 2023-05-24 07:47:50 +01:00
jax authors
af4d4943a7 Merge pull request #8633 from shawwn:2021-11-19/autodidax-fix-jaxpr-subcomp-return-type 
PiperOrigin-RevId: 519745476
 2023-03-27 09:52:20 -07:00
Peter Hawkins
f66f6ec98a [JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge. 
Limit jax._src.lib to shims around jaxlib and nothing else.

The goal of this change is to avoid a dependency cycle between the rest of jax and jax._src.lib in a Bazel build. This allows the types for jax._src.lib to be inferred by pytype in isolation without referring to the rest of JAX.

PiperOrigin-RevId: 512922397
 2023-02-28 07:01:57 -08:00
Jake VanderPlas
87051ba875 DOC: fix autodidax 2023-01-26 17:23:07 -08:00
Jake VanderPlas
20b55a119e CI: update jupytext version 2023-01-23 14:42:03 -08: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
5527966b27 [JAX] Deprecate .to_py() property on arrays. Implement __array__ instead. 
.to_py() was something of an accidental export from the JAX array classes. There are other mechanisms to turn a JAX array into a NumPy array, including `np.asarray(x)` and `jax.device_get(x)`. Deprecate this mechanism because it is redundant.

PiperOrigin-RevId: 469984029
 2022-08-25 07:28:27 -07:00
jax authors
2eb5a6e65b Merge pull request #11927 from NeilGirdhar:fix_typo 
PiperOrigin-RevId: 467949795
 2022-08-16 09:49:59 -07:00
Neil Girdhar
ad38a6bb28 Fix common typo: Tuple[X] -> Tuple[X, ...] 2022-08-16 11:47:22 -04:00
Jake VanderPlas
eeb9b5f1f6 pre-commit hook: update flake8, mypy, & jupytext 2022-08-15 15:32:45 -07:00
Jeppe Klitgaard
17de89b16a feat: refactor code using pyupgrade 
This PR upgrades legacy Python code to 3.7+ code using pyupgrade:
```sh
pyupgrade --py37-plus --keep-runtime-typing **.py
```

a
 2022-05-17 22:14:05 +01:00
Jake VanderPlas
8a3bfe0300 DOC: add references for haskell-style signatures 2022-05-13 12:35:28 -07:00
Jake VanderPlas
1134ff5871 CI: update jupytext to 0.13.8 2022-05-02 11:22:31 -07:00
Matthew Johnson
43036e1316 improvements to autodidax reduce_sum 
* generalize reduce_sum to handle multiple axes
* add reduce_sum transpose rule

also fix bug in AD jaxpr formation related to deduplicating consts
 2022-03-15 20:22:55 -07:00
Jake VanderPlas
b9b79bab31 maint: update pre-commit package versions & fix new mypy errors 2022-01-31 13:39:11 -08:00
Peter Hawkins
68e9e1c26d Consolidate more XLA-lowering logic between jit, pmap, and xmap. 
Move remaining functions relating to building XLA HLO IR out of xla_bridge.py and into jax.interpreters.xla.

PiperOrigin-RevId: 413244450
 2021-11-30 14:24:33 -08:00
Shawn Presser
2d61a5fedb autodidax: fix jaxpr_subcomp return type annotation 2021-11-19 22:19:55 -06:00
Peter Hawkins
1a73743610 Move xla_bridge.constant to jax.interpreter.xla.pyval_to_ir_constant. 
This is a more descriptive name and a better location (next to other facilities for building XLA IR).

Quite a few users of the former xla_bridge.constant() didn't need anything other than uncanonicalized array constants. Change these users to use xla_client.ops.Constant instead; no need for the fancy utility in these cases.

PiperOrigin-RevId: 404270649
 2021-10-19 08:40:51 -07:00
Peter Hawkins
714e19a794 Remove xla_bridge.make_computation_builder(). 
This is a vestigal wrapper around xla_client.XlaBuilder whose purpose is long gone.

Also rename uses of XlaComputationBuilder to XlaBuilder. XlaComputationBuilder was an older name that is gone in most places.
 2021-10-18 13:20:34 -04:00
Jake VanderPlas
198cc5ee4f chore: update jupytext to v0.1.13 & re-sync notebooks 2021-10-05 14:30:16 -07:00
Peter Hawkins
2c2f4033cc Move contents of jax.lib to jax._src.lib. 
Add shim libraries for functions exported from jax.lib that other code seems to use in practice.

PiperOrigin-RevId: 398471863
 2021-09-23 06:33:55 -07:00
Sharad Vikram
0fa70084a3 Pass x into transpose in autodidax 2021-08-24 17:46:34 -07:00
Matthew Johnson
24de3e992c small autodidax tweaks 2021-08-05 05:04:08 -07:00
elliotwaite
7392a57b75 DOC: many small fixes 2021-08-04 16:55:13 -07:00
Matthew Johnson
fa274f3ace small tweaks, type annotations 2021-07-22 21:09:58 -07:00
George Necula
6a48c60a72 Rename master to main in embedded links. 
Tried to avoid the change on external links to repos that
have not yet renamed master.
 2021-06-18 10:00:01 +03:00
Matthew Johnson
d88acd8b8c autodidax: delete while_loop for now 2021-05-06 12:43:02 -07:00
Matthew Johnson
83cd42271b autodidax: add cond and start while_loop 2021-05-05 12:32:39 -07:00
jax authors
2fea627cbc Merge pull request #6401 from dan-zheng:fix-doc-typos 
PiperOrigin-RevId: 367755029
 2021-04-09 22:10:01 -07:00
Dan Zheng
0b41ff93ea DOC: fix typos. 2021-04-09 23:53:55 -04:00
Matthew Johnson
5f6bce4bfe add 'open in colab' button, add numpy<1.18 compat 
Co-authored-by: Edward Loper <edloper@google.com>
 2021-04-09 20:09:00 -07:00
James Knighton
8772b2c114 Allodidactax... 2021-03-29 04:57:02 -07:00
James Knighton
0d19b7c082 Mirror the minor spelling fixes over the {.ipynb, .md, .py}. 2021-03-27 21:47:07 -07:00
Matthew Johnson
5a055b1dad add WIP disclaimer to autodidax, fix some typos 2021-03-27 19:34:42 -07:00
Matthew Johnson
3595af2ed2 add linearize, vjp, grad. fix bugs. 2021-03-26 21:27:37 -07:00
Matthew Johnson
3457696e80 add linearize code (needs text) 2021-03-12 14:49:13 -08:00
Matthew Johnson
6b5133d71f autodidax: jit, multi-output, pytrees, DeviceArrays 2021-03-05 19:38:52 -08:00
Matthew Johnson
e4f79fb13e add license 2021-02-24 20:57:15 -08:00
Matthew Johnson
9789677e85 reviewer comments, delint 2021-02-24 20:26:48 -08:00