Dan Foreman-Mackey
782138fb6f
Add custom_dce to changelogs and API docs.
2025-01-27 13:03:34 -05:00
jax authors
726abc9c31
Merge pull request #26082 from Saransh-cpp:index-update-syntax-link
...
PiperOrigin-RevId: 719356167
2025-01-24 10:37:07 -08:00
Saransh Chopra
9e89ae7a19
docs: outdated link for index update syntax
2025-01-24 17:36:35 +00:00
Stephan Hoyer
b2afb5bf4f
add docstrings for check_vjp and check_jvp
2025-01-23 17:31:13 -08:00
Stephan Hoyer
458f6a6efe
Add jax.test_util to public API docs
2025-01-23 16:04:35 -08:00
Jake VanderPlas
423be16ecc
DOC: mention adoption of EffVer in JEP
2025-01-22 15:44:11 -08:00
Jakob Roth
4f8699c8a1
Update docs of callbacks
...
Callback functions should not call into JAX. This information was
missing in the docs of the callbacks. This commit adds this information
to the docs.
See: #25861 , #24255
2025-01-19 11:33:20 +01:00
Jake VanderPlas
3141453b1b
Add performance note at the top of sparse docs
2025-01-17 10:17:16 -08:00
Robert Dyro
aa9cea0a55
Add quotes to pip commands in docs around option install for zsh
...
PiperOrigin-RevId: 716358130
2025-01-16 13:30:43 -08:00
Zac Mustin
2d72e8de84
Jax: Stop returning a list of cost-analyses.
...
As it stands, there is only ever one element in this list (see b/384741132) and only the 0th element is ever used so we can simplify.
This is a potentially breaking change for external users, but (as stated in the [documentation](https://jax.readthedocs.io/en/latest/aot.html#debug-information-and-analyses-when-available )) no guarantees are made on this type, which is intended for debugging purposes and not intended to be a reliable public API.
PiperOrigin-RevId: 715837855
2025-01-15 09:53:59 -08:00
Benjamin Chetioui
57a259f447
[Docs] Remove --xla_gpu_enable_triton_softmax_fusion from docs
...
This flag has been a no-op for a while.
PiperOrigin-RevId: 715491248
2025-01-14 12:50:59 -08:00
Roy Frostig
a60ead6fd1
enable partitionable threefry by default
...
PiperOrigin-RevId: 715242560
2025-01-13 22:46:24 -08:00
Ayaka
9ba1fd2801
[Pallas TPU] Add vector support to pl.debug_print
...
PiperOrigin-RevId: 715085454
2025-01-13 13:22:21 -08:00
jax authors
1fe72ee880
Merge pull request #25771 from dfm:custom-partition-ffi-tutorial
...
PiperOrigin-RevId: 714048619
2025-01-10 07:43:10 -08:00
George Necula
dd0447a7c6
[aot] Add support for as_text(debug_info=True).
...
This exposes an easier way to get StableHLO and HLO
with more debugging information (source locations
for StableHLO and metadata for HLO).
2025-01-10 07:59:56 +02:00
Dan Foreman-Mackey
729418094e
Add a discussion of sharding to the FFI tutorial.
2025-01-09 11:24:40 -05:00
Sergei Lebedev
90201ce2b7
Removed leftover mentions of xmap from the code
...
PiperOrigin-RevId: 713202387
2025-01-08 01:39:38 -08:00
Jane Liu
21fb171ef9
Merge branch 'jax-ml:main' into activation-offloading-doc
2025-01-07 21:20:29 -08:00
jax authors
57c2afe7a8
Merge pull request #25441 from Exferro:fixed_advanced_autodiff_doc
...
PiperOrigin-RevId: 712929769
2025-01-07 08:56:05 -08:00
Dan Foreman-Mackey
a7f384cc6e
Add a register_custom_type_id function to the GPU plugins.
...
This enables dynamic registration of custom FFI types on the appropriate platform via PJRT.
PiperOrigin-RevId: 712904085
2025-01-07 07:29:38 -08:00
Aleksei Malyshev
f881f507d6
Update the advanced autodiff tutorial and replace some vmap with grad
2025-01-07 15:56:23 +01:00
jax authors
56f0f9534d
Merge pull request #25633 from dfm:move-ffi
...
PiperOrigin-RevId: 712863350
2025-01-07 04:40:21 -08:00
jax authors
7997f080f2
Merge pull request #25728 from zhenying-liu:scipy.misc
...
PiperOrigin-RevId: 712707311
2025-01-06 17:53:10 -08:00
Jake VanderPlas
c7b0d681bd
Remove deprecated jax.experimental.array_api
2025-01-06 15:19:02 -08:00
Jane Liu
77c6947a59
fix the doc error: module 'scipy.misc' has no attribute 'face'
2025-01-06 12:00:02 -08:00
Jane Liu
5e3a692d36
Merge branch 'jax-ml:main' into activation-offloading-doc
2025-01-04 16:36:27 -08:00
jax authors
800f903f9b
Merge pull request #25686 from Mikcl:docs/working-with-pytrees-formatting
...
PiperOrigin-RevId: 711448394
2025-01-02 09:40:17 -08:00
Robert Dyro
213e1782ac
tbp nightly instructions
2025-01-02 09:49:31 +01:00
Dan Foreman-Mackey
cb4d97aa1f
Move jex.ffi to jax.ffi.
2024-12-29 13:06:19 +00:00
mikcl
008c25a369
Fix formatting in the docs for transposing pytrees
2024-12-26 01:18:10 +00:00
Jane Liu
522a8fd792
consolidate the code example
2024-12-19 14:47:29 -08:00
Jane Liu
f124232f53
Add activation offloading examples in a subsection
2024-12-19 14:47:29 -08:00
Jake VanderPlas
26930366a3
DOC: clarify API compatibility discussion & mention EffVer
2024-12-19 11:23:43 -08:00
jax authors
bad6f0f0e9
Merge pull request #25440 from ROCm:jax_pypi-upstream
...
PiperOrigin-RevId: 707934347
2024-12-19 09:17:46 -08:00
jax authors
49c92460a4
Merge pull request #25516 from jakevdp:jep-effver
...
PiperOrigin-RevId: 707923040
2024-12-19 08:33:52 -08:00
Ruturaj4
fefebebea4
Update documentation and add setup.py pypi bindings
2024-12-19 10:21:28 -06:00
Matthew Johnson
c6576827c6
remove whitespace from distributed arrays doc
2024-12-18 23:55:20 +00:00
Jake VanderPlas
88c8941991
JEP: effort-based versioning (EffVer)
2024-12-18 09:25:49 -08:00
Yash Katariya
83ca7c33d2
Don't create a mesh inside jit in sharded-computation tutorial
...
PiperOrigin-RevId: 707288222
2024-12-17 16:11:00 -08:00
Peter Hawkins
7de9eb20df
Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0
...
PiperOrigin-RevId: 707146329
2024-12-17 10:12:34 -08:00
jax authors
cd7109e6b5
Minor typo fixin doc.
...
PiperOrigin-RevId: 706774033
2024-12-16 11:22:45 -08:00
George Necula
afcb62ea20
[export] Expand exporting to work with AbstractMesh.
...
This is a follow up from #25640 that enabled lowering with
AbstractMesh.
This required adding `num_devices` to `lowering.compiler_args`
because in presence of an AbstractMesh the device_assignment
is not accurate.
2024-12-16 10:30:46 +02:00
Jake VanderPlas
c73f306099
Finalize deprecation of jnp.round_
...
PiperOrigin-RevId: 705998500
2024-12-13 14:13:44 -08:00
jax authors
99b390ce96
Merge pull request #25444 from justinjfu:gpu_docs_update
...
PiperOrigin-RevId: 705938411
2024-12-13 11:05:31 -08:00
Sevin Varoglu
c563e47314
Add flag desc to gpu_performance_tips.md
2024-12-12 15:39:07 -08:00
Ivy Zheng
26c40fadfd
Add jax.tree shortcuts for .*_with_path calls, for convenience of users.
...
PiperOrigin-RevId: 705645570
2024-12-12 15:13:32 -08:00
Justin Fu
1021603f85
Remove deprecated XLA GPU flags.
2024-12-12 15:10:43 -08:00
Jake VanderPlas
f858a71461
Finalize some deprecations in jax.core, jax.lib.xla_bridge, and jax.lib.xla_client.
2024-12-11 09:50:33 -08:00
Jake VanderPlas
f6d58761d1
jax.numpy: implement matvec & vecmat
2024-12-10 16:03:19 -08:00
Jake VanderPlas
6541a62099
jax.core: deprecate a number of APIs
2024-12-10 11:11:32 -08:00
