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

43 Commits

Author SHA1 Message Date
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
Jan Hrček
4da56dcdd7 Fix duplicate word occurrences 2023-12-19 06:15:30 +01: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
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
Neil Girdhar
ad38a6bb28 Fix common typo: Tuple[X] -> Tuple[X, ...] 2022-08-16 11:47:22 -04: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
29efc30c63 DOC: update myst-nb to v0.15.0 2022-05-09 14:32:48 -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
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
0d19b7c082 Mirror the minor spelling fixes over the {.ipynb, .md, .py}. 2021-03-27 21:47:07 -07:00
James Knighton
a5a93dc845
Speling 2021-03-27 21:20:34 -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
Jake VanderPlas
b2e226a8b9 Fix build warning in autodidax.md 2021-02-26 09:13:19 -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
Matthew Johnson
0fcecb2bd2 Add first two sections of Autodidax. 2021-02-24 01:03:52 -08:00