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
7,970 Commits 506 Branches 263 Tags
Commit Graph

740 Commits

Author SHA1 Message Date
Bogdan Damoc
db39a67cca Update documentation for custom_jvp handling of nondiff_argnums as arguments of _fwd and _bwd rules. 2021-05-12 16:44:52 +01:00
Peter Hawkins
d005e38f78 Promote the x.at[idx].set(y) operators as the preferred way to do indexed updates. 
Mark the index_update() etc. operators as deprecated in the documentation.

Add new .divide and .power operators. Fixes #2694.
Add .multiply as an alias for .mul. To be more numpy-like we should probably prefer the longer names.
 2021-05-10 20:32:00 -04:00
Jake VanderPlas
26f74e64a6 DOC: add documentation of DeviceArray object properties & methods 2021-05-10 12:37:59 -07:00
jax authors
a946824732 Merge pull request #6682 from bdamoc:updated_doc 
PiperOrigin-RevId: 372583046
 2021-05-07 10:32:12 -07:00
Bogdan Damoc
33ca40a9db Update code example in custom VJP documentation. 2021-05-07 16:34:25 +01:00
Parker Schuh
92246017d2
Revert "Use convert_element_type instead of device_put_raw." 2021-05-06 20:19:23 -07:00
jax authors
d0aa875b4d Merge pull request #6400 from pschuh:convert_element 
PiperOrigin-RevId: 372431314
 2021-05-06 14:57:12 -07:00
jax authors
4bbb24d582 Merge pull request #6453 from google:autodidax 
PiperOrigin-RevId: 372422829
 2021-05-06 14:14:19 -07:00
Parker Schuh
9d3e535ad2
Merge branch 'master' into convert_element 2021-05-06 13:18:01 -07:00
Matthew Johnson
d88acd8b8c autodidax: delete while_loop for now 2021-05-06 12:43:02 -07:00
Adam Paszke
ee93ee221c Add a docstring for pjit 
Co-authored-by: Skye Wanderman-Milne <skyewm@google.com>
PiperOrigin-RevId: 372393461
 2021-05-06 12:00:49 -07:00
Matthew Johnson
83cd42271b autodidax: add cond and start while_loop 2021-05-05 12:32:39 -07:00
jax authors
0bf732be15 Merge pull request #6656 from skye:profiling_docs 
PiperOrigin-RevId: 372174753
 2021-05-05 11:49:32 -07:00
Skye Wanderman-Milne
22ae98ac50 Update profiling instructions. 2021-05-05 11:40:22 -07:00
Jake VanderPlas
382506ba1a DOC: add bicgstab to doc page 2021-05-04 12:24:50 -07:00
Matthew Johnson
7ec0b40173 Roll-forward of #6584, which broke internal tests. 
PiperOrigin-RevId: 371839298
 2021-05-03 21:41:23 -07:00
jax authors
ffecef5555 Merge pull request #6611 from njunge94:pip-install-fork 
PiperOrigin-RevId: 371724836
 2021-05-03 10:26:33 -07:00
jax authors
75b00a1235 Copybara import of the project: 
--
3c400a3e588abf9e2259119c50343cba6f3477f1 by Matthew Johnson <mattjj@google.com>:

add 'inline' option to xla_call for jaxpr inlining

--
fe297e39ca37896b75d7943b9b77c0b53fad13ee by Matthew Johnson <mattjj@google.com>:

add 'inline' to jit docstring

--
ff6866c4b3757cde66fe659c2f27d8aeff024e8f by Matthew Johnson <mattjj@google.com>:

new_sublevel in jax2tf

PiperOrigin-RevId: 371542778
 2021-05-01 22:18:39 -07:00
jax authors
e8f209c775 Merge pull request #6584 from google:jit-inline-2 
PiperOrigin-RevId: 371541392
 2021-05-01 21:53:34 -07:00
Ian Horn
74b6762745 Fixed error in autodidax broadasting batching 2021-05-01 19:07:58 -07:00
Nicholas Junge
f559a9bfb4 Fix pip install instruction for forks in contributing section 2021-05-01 19:32:19 +02:00
Nicholas Junge
e985e86808 Implement jnp.r_ and jnp.c_ 2021-05-01 01:05:22 +02:00
Lukas Geiger
a613ce12a3 Export jnp.broadcast_shapes as user facing function 2021-04-30 19:32:51 +02:00
Roy Frostig
b8f9dd6269 unify tree_map and tree_multimap 2021-04-28 19:59:31 -07:00
Matthew Johnson
3c400a3e58 add 'inline' option to xla_call for jaxpr inlining 2021-04-28 19:38:15 -07:00
Jake VanderPlas
ba83ab5d3b Add expanded contributing documentation 2021-04-27 15:00:21 -07:00
Matthew Johnson
ba9233b9b6 prune trivial convert_element_types from jaxprs 
also add a test for not performing H2D transfers while tracing jnp.array
 2021-04-22 12:46:26 -07:00
Peter Hawkins
aafe8870ae Document that JAX follows the NEP-29 deprecation policy. 
Remove the "experimental" disclaimer from the concurrency documentation.
 2021-04-21 11:12:41 -04:00
Skye Wanderman-Milne
9128ba0c74 Replace host_id with process_index terminology, take 2. 
We're switching to the new terminology to avoid confusion in cases
where multiple jax processes are running on a single host, and each
process has a unique process_index/host_id.

This keeps aliases for the old `host_id` APIs for now, but these will
eventually be removed.

This was originally commited in
b77ef5138b631378e6a8ceb8bafc94fe91239bae, but reverted in
14acd070c2afb11c81fc91f43790577cd48cbf67 due to Google-internal test
failures from renaming the local_devices argument name. This change is
identical except it also adds staging for the argument name change.
 2021-04-20 18:13:34 -07:00
jax authors
14acd070c2 Internal change 
PiperOrigin-RevId: 369345279
 2021-04-19 18:23:07 -07:00
Skye Wanderman-Milne
b77ef5138b Replace host_id with process_index terminology. 
We're switching to the new terminology to avoid confusion in cases
where multiple jax processes are running on a single host, and each
process has a unique process_index/host_id.

This keeps aliases for the old `host_id` APIs for now, but these will
eventually be removed.
 2021-04-19 14:09:19 -07:00
Jake VanderPlas
89dc691da7 DOC: fix issue in pytrees.md 2021-04-15 09:57:53 -07:00
Skye Wanderman-Milne
cec132e32c Pin sphinx-autodoc-typehints version to 1.11.1 
Our docs build fails with 1.12.1: https://readthedocs.org/projects/jax/builds/13514541/
```
reading sources... [ 50%] _autosummary/jax.numpy.linalg.pinv

Traceback (most recent call last):
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.8/site-packages/sphinx/events.py", line 111, in emit
    results.append(listener.handler(self.app, *args))
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.8/site-packages/sphinx_autodoc_typehints.py", line 417, in process_docstring
    type_hints = get_all_type_hints(obj, name)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.8/site-packages/sphinx_autodoc_typehints.py", line 271, in get_all_type_hints
    rv = backfill_type_hints(obj, name)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.8/site-packages/sphinx_autodoc_typehints.py", line 312, in backfill_type_hints
    obj_ast = ast.parse(textwrap.dedent(
  File "/home/docs/.pyenv/versions/3.8.6/lib/python3.8/ast.py", line 47, in parse
    return compile(source, filename, mode, flags,
  File "<unknown>", line 12
    max_rows_cols = max(a.shape[-2:])
                                    ^
IndentationError: expected an indented block
```

This may be caused by https://github.com/agronholm/sphinx-autodoc-typehints/pull/166
 2021-04-14 17:58:27 -07:00
jax authors
cb881bb4f0 Merge pull request #6425 from apaszke:notebook-fix 
PiperOrigin-RevId: 368451960
 2021-04-14 10:00:28 -07:00
Adam Paszke
62ac3e058b Fix inline image embedding in the xmap tutorial 2021-04-14 10:14:50 +00:00
Peter Hawkins
26e9ebcdae Move jax.api to jax._src.api. 
PiperOrigin-RevId: 368233837
 2021-04-13 09:43:24 -07:00
Sharad Vikram
577c601191 Update references in custom interpreter notebook 2021-04-12 19:00:28 -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
Jake VanderPlas
cbb7052379 Implement segment_prod, segment_max, segment_min 2021-04-09 12:06:51 -07:00
jax authors
a27b9005c5 Merge pull request #6100 from apaszke:xmap-tutorial 
PiperOrigin-RevId: 367613958
 2021-04-09 05:10:12 -07:00
Jake VanderPlas
ffd185612c DOC: fix tree_util.rst 2021-04-08 15:59:55 -07:00
Adam Paszke
40253f471d Make xmap tutorial public 
Note that xmap is still experimental, so treat this more as a preview
than as documentation.
 2021-04-08 15:56:58 +00:00
Jake VanderPlas
011c5a203e DOC: add register_pytree_node_class example 2021-04-08 08:55:18 -07:00
George Necula
d9468c7513 Cleanup the API, and more documentation 2021-04-08 11:25:32 +03:00
Peter Hawkins
6a6f13e1b0 [JAX] Move contents of jax/dtypes.py to jax/_src/dtypes.py. 
PiperOrigin-RevId: 367345623
 2021-04-07 19:35:51 -07:00
jax authors
75c898e4f0 Merge pull request #6342 from minoring:impl-ogrid 
PiperOrigin-RevId: 366907253
 2021-04-05 17:30:47 -07:00
minoring
4c67dd1f48 Implement jnp.ogrid 
Related to #5850
 2021-04-06 08:37:07 +09:00
Jake VanderPlas
6eb761576c Fix cross-reference in pytrees.md 2021-04-05 16:09:04 -07:00