mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-20 05:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
13,235 Commits 506 Branches 263 Tags
Commit Graph

1085 Commits

Author SHA1 Message Date
jbushago
2038988783
Fix typo in faq.rst. 
Fixed a small typo in the FAQ: "inthe" -> "in the".
 2022-09-30 14:14:05 -04:00
jax authors
d2fcfb6b83 Merge pull request #12407 from hirwa-nshuti:docs-fix 
PiperOrigin-RevId: 476467728
 2022-09-23 14:51:11 -07:00
Felix Hirwa Nshuti
820efab6fa removed repeated nan_to_num in docs 2022-09-23 06:23:09 +00:00
Sharad Vikram
99d4d8b89a Update debugging docs to have sharding visualization 2022-09-22 19:42:36 -07:00
jax authors
bc08381da3 Merge pull request #12152 from nvcastet:add_slurm_orchestrator_support 
PiperOrigin-RevId: 476179963
 2022-09-22 13:18:25 -07: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
Nicolas Castet
412a5379c1 Add generic interface for auto initialization of distributed JAX service 
* Also add slurm cluster support
 2022-09-22 14:15:38 -05:00
Jake VanderPlas
fce1099997 Update JEP-12049 implementation discussion 2022-09-20 09:44:29 -07:00
Jake VanderPlas
be65694ac6 docs: avoid deprecated matplotlib axis creation 2022-09-19 12:55:18 -07:00
Jake VanderPlas
5829c6ae9d Change case of typing.Dtype -> typing.DType 
This follows the convention used in numpy.typing.DType.
 2022-09-14 15:03:55 -07:00
Jake VanderPlas
358363e17f JEP 12049: Type Annotation Roadmap 2022-09-13 09:14:48 -07:00
Jake VanderPlas
0fb462efd7 Add jax.print_environment_info() 2022-09-12 15:39:33 -07:00
Peter Hawkins
57b5acf1b6 Roll forward: Upgrade logistic into a primitive. 
Unlike the previous attempt, we don't try to use mhlo.logistic as the lowering of the new primitive yet. Instead, we lower to the old implementation of `expit`. This means that this change should be a no-op numerically and we can work on changing its implementation in a subsequent change.

PiperOrigin-RevId: 472705623
 2022-09-07 06:06:56 -07:00
Filippo Vicentini
52236adeed
Show link to GitHub repo in navbar 2022-09-05 13:41:49 +02:00
Roy Frostig
bb68fbeefa write in-process AOT walkthrough doc 2022-09-02 13:02:25 -07:00
Roy Frostig
43db06491c write and generate package API documentation for jax.stages 2022-09-01 19:26:53 -07:00
jax authors
9c16c83234 Rollback of upgrade logistic (sigmoid) function into a lax primitive. 
PiperOrigin-RevId: 471105650
 2022-08-30 15:30:43 -07:00
jax authors
cdd16e3bbd Fix typo in jax.debug.print() documentation. 
PiperOrigin-RevId: 471096544
 2022-08-30 14:51:12 -07:00
jax authors
bff98655a1 Merge pull request #12112 from sharadmv:docs-typo 
PiperOrigin-RevId: 470762998
 2022-08-29 10:58:40 -07:00
Sharad Vikram
27b313b287 Fix typo in debugging docs 2022-08-29 10:32:08 -07:00
Peter Hawkins
f68f1c0cd0 Upgrade logistic (sigmoid) function into a lax primitive. 
This allows us to lower it to `mhlo.logistic`, which allows XLA to generate more efficient code.

PiperOrigin-RevId: 470300985
 2022-08-26 11:58:28 -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
3e3542b0d6 Upgrade logistic (sigmoid) function into a lax primitive. 
This allows us to lower it to `mhlo.logistic`, which allows XLA to generate more efficient code.

PiperOrigin-RevId: 469841487
 2022-08-24 15:39:37 -07:00
Sharad Vikram
175823e066 Add callbacks to docs 2022-08-24 14:12:47 -07:00
Peter Hawkins
6276194e1c Upgrade logistic (sigmoid) function into a lax primitive. 
This allows us to lower it to `mhlo.logistic`, which allows XLA to generate more efficient code.

PiperOrigin-RevId: 469789339
 2022-08-24 12:04:01 -07:00
Peter Hawkins
335b2cfb26 [JAX] Prepare not to export jax._src by default. 
Currently
```
import jax
```
populates `jax._src` in the names exported from JAX. This change prepares for not exporting `jax._src` by default.

In particular, explicitly import modules from jax._src and refer to those imports rather than assuming jax._src contents will be around later. This is a common pattern in tests.

This change does not yet remove any exported names.

Issue https://github.com/google/jax/issues/11951

PiperOrigin-RevId: 469480816
 2022-08-23 09:36:47 -07:00
Matthew Johnson
d19e34fa4a delete old remat implementation 
moved lowering rule logic from remat_impl.py (now deleted) to ad_checkpoint.py
 2022-08-16 23:16:37 -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
Sharad Vikram
8b54298052 Fix typo in debugging docs 2022-08-12 11:59:00 -07:00
Roy Frostig
c6ab3a6a60 convert custom VJP update guide to a retroactive JEP 2022-08-10 13:45:57 -07:00
Matthew Johnson
be6f6bfe9f set new jax.remat / jax.checkpoint to be on-by-default 2022-08-10 10:29:38 -07:00
jax authors
d95b27ce1c Merge pull request #11803 from hawkinsp:multigpu 
PiperOrigin-RevId: 466374926
 2022-08-09 09:02:21 -07:00
Peter Hawkins
a2c21958a5 Document multiprocess GPU support. 
Fixes #2731
 2022-08-09 11:31:05 -04:00
Roy Frostig
4c18f1a580 link to both closed and open enhancement proposals 
PiperOrigin-RevId: 466212251
 2022-08-08 18:48:38 -07:00
jax authors
38ab3d88ae Merge pull request #11799 from hawkinsp:jep 
PiperOrigin-RevId: 466178951
 2022-08-08 16:16:18 -07:00
Peter Hawkins
71b29b1cc6 Create JAX Enhancement Proposals (JEPs). 
Migrate existing design documents to JEPs.
 2022-08-08 16:13:58 -04:00
jax authors
25dd62c388 Merge pull request #11793 from mbrukman:patch-1 
PiperOrigin-RevId: 466093897
 2022-08-08 11:12:16 -07:00
Misha Brukman
86a1a6b1ee
Fix sample code block formatting 
Add missing colon to format sample code as a code block rather than a quote.
 2022-08-08 12:41:01 -04:00
Peter Hawkins
3dc7ce2bfb Pin a minimum version of sphinx-book-theme to avoid bad pip dependency 
resolution.
 2022-08-08 16:11:19 +00:00
Matthew Johnson
7ed4fa550e in sharp bits notebook, add back iaml blog post link 2022-08-04 08:29:26 -07:00
Jake VanderPlas
37c249fc80 MAINT: unpin autodoc-typehints 2022-08-01 10:33:50 -07:00
Sharad Vikram
11b206a18a Enable debugging primitives in pjit on CPU/GPU 
PiperOrigin-RevId: 464208326
 2022-07-29 20:10:27 -07:00
jax authors
6cffa720e7 Merge pull request #11670 from sharadmv:debugging-docs 
PiperOrigin-RevId: 464145832
 2022-07-29 13:26:56 -07:00
Sharad Vikram
decdca60c8 Change jaxdb->jdb and add option to force a backend 2022-07-29 12:51:27 -07:00
jax authors
cc19c94d36 Merge pull request #11665 from LenaMartens:docs 
PiperOrigin-RevId: 464113785
 2022-07-29 10:56:58 -07:00
lenamartens
0256301617 tweak checkify docs. 2022-07-29 18:15:36 +01:00
jax authors
25538c8260 Merge pull request #11662 from sharadmv:debugging-docs 
PiperOrigin-RevId: 464052649
 2022-07-29 04:53:44 -07:00
Sharad Vikram
fb0cf668b8 Update debugging docs to mention pjit 2022-07-28 22:00:27 -07:00