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

32 Commits

Author SHA1 Message Date
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
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