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

157 Commits

Author SHA1 Message Date
Dan Foreman-Mackey
60bf5b7727 Add a jax.process_indices function. 
The `jax.host_ids` function has be long deprecated, but the suggested alternative of `list(range(jax.process_count()))` relies on the current behavior that the list of process indices is always dense. In the future we may want to allow dynamic addition and removal of processes in which case `jax.process_count` and `jax.process_indices` would need to be updated, and it is useful for users to be able to use this forward-compatible interface.

PiperOrigin-RevId: 662142636
 2024-08-12 10:30:41 -07:00
Yash Katariya
0d5dae09ff Delete xmap and the jax.experimental.maps module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period). 
PiperOrigin-RevId: 655614395
 2024-07-24 10:24:09 -07:00
rdyro
c6d6207170 Unifying persistent cache messages 
and moving them to WARNING logging when explain_cache_misses is true.
 2024-07-16 00:47:53 +00:00
Matthew Johnson
bd166e1d99 add more info to xla_computation deprecation warning 2024-07-02 13:31:07 +00:00
Yash Katariya
e6f26ff256 Deprecate jax.xla_computation. Use JAX AOT APIs to get the equivalent of jax.xla_computation functionality. 
PiperOrigin-RevId: 644107276
 2024-06-17 13:02:35 -07:00
Jake VanderPlas
0a86e9a929 Deprecate hashing of tracers 2024-06-13 13:14:27 -07:00
Jake VanderPlas
aa1452375b Register beta args deprecation 
PiperOrigin-RevId: 642427224
 2024-06-11 16:19:14 -07:00
Dan Foreman-Mackey
1e206880d3 Move jax.ffi submodule to jax.extend.ffi 2024-05-31 12:34:59 -04:00
Dan Foreman-Mackey
88790711e8 Package XLA FFI headers with jaxlib wheel 
The new "typed" API that XLA provides for foreign function calls is
header-only and packaging it as part of jaxlib could simplify the open
source workflow for building custom calls.

It's not completely obvious that we need to include this, because jaxlib
isn't strictly required as a _build_ dependency for FFI calls, although
it typically will be required as a _run time_ dependency. Also, it
probably wouldn't be too painful for external projects to use the
headers directly from the openxla/xla repo.

All that being said, I wanted to figure out how to do this, and it has
been requested a few times.
 2024-05-22 12:28:38 -04:00
Jake VanderPlas
d33a5689de Refactor & test internal deprecation APIs 
The names and APIs were previously too similar and therefore somewhat confusing; this will be more clear I think.

PiperOrigin-RevId: 635615163
 2024-05-20 17:16:31 -07:00
Jake VanderPlas
4bac10e750 Finalize deprecation of the config module. 
To configure JAX, use `import jax` and reference the config object via `jax.config`.

PiperOrigin-RevId: 635430169
 2024-05-20 05:49:31 -07:00
Mark Sandler
8f045cafd2 Add jax.make_array_from_process_local_data to create a distributed tensor from host data and supporting scaffolding in sharding to be able to figure out dimensions of host data required. 
PiperOrigin-RevId: 634205261
 2024-05-15 22:06:45 -07:00
Jake VanderPlas
2daaf49541 Remove extraneous pure_callback_api wrapper 2024-04-25 10:21:49 -07:00
Matthew Johnson
8588d4b747 alias jax.sharding.NamedSharding -> jax.NamedSharding 2024-04-11 16:23:59 -07:00
Jake VanderPlas
8949a63ce1 [key reuse] rename flag to jax_debug_key_reuse 2024-03-22 05:37:30 -07:00
Jake VanderPlas
d8662886d7 Register maps module deprecation outside of module 
PiperOrigin-RevId: 617194807
 2024-03-19 09:21:48 -07:00
Yue Sheng
147c363ea6 Deprecate jax.clear_backends. 
`jax.clear_backends` does not necessarily do what its name suggests and can lead to unexpected consequences, e.g., it will not destroy existing backends and release corresponding owned resources. Use `jax.clear_caches` if you only want to clean up compilation caches. For backward compatibilty or you really need to switch/reinitialize the default backend, use `jax.extend.backend.clear_backends`.

PiperOrigin-RevId: 616946337
 2024-03-18 14:23:18 -07:00
Jake VanderPlas
236275ebe1 Deprecate jax.tree_map for jax v0.4.26 
Reverts f4045dceb206be1ea10ee651ccc6151809f2d9f3

PiperOrigin-RevId: 611230367
 2024-02-28 14:29:01 -08:00
Yash Katariya
f4045dceb2 Remove the deprecation of jax.tree_map for the release of 0.4.25 
PiperOrigin-RevId: 610014256
 2024-02-24 09:30:06 -08:00
Jake VanderPlas
e59a0506fe Deprecate jax.tree_map in favor of jax.tree.map 2024-02-22 11:35:39 -08:00
Jake VanderPlas
cf80f574b5 Register jax.config module deprecation 
PiperOrigin-RevId: 609352291
 2024-02-22 06:38:56 -08:00
Sergei Lebedev
57e59eb6c3 Removed deprecated jax.config methods and jax.config.config 
Reverts dcc65e621ea3a68fdc79fa9f2c995743a7b3faf7

PiperOrigin-RevId: 608676645
 2024-02-20 11:25:16 -08:00
Thomas Köppe
dcc65e621e Reverts b506fee9e389391efb1336bc7575dba913e75cdf 
PiperOrigin-RevId: 608319964
 2024-02-19 06:23:00 -08:00
Sergei Lebedev
b506fee9e3 Removed deprecated jax.config methods and jax.config.config 
Reverts eb0343683547b6e2d29245f3ab6c91037c0cff81

PiperOrigin-RevId: 607803834
 2024-02-19 06:21:15 -08:00
jax authors
eb03436835 Reverts 318a19a89387caebd116168c4e47592e7d71ca65 
PiperOrigin-RevId: 607708463
 2024-02-16 09:11:05 -08:00
Sergei Lebedev
318a19a893 Removed deprecated jax.config methods 
PiperOrigin-RevId: 607675571
 2024-02-16 06:49:13 -08:00
Jake VanderPlas
6934a4b76b Add jax.tree module with aliases of jax.tree_util 2024-02-12 13:07:59 -08:00
Jake VanderPlas
e356d76913 Remove a number of deprecated APIs 
All of these were deprecated prior to the JAX 0.4.16 release, on Sept 18 2023.
As of Monday Dec 18, we have met the 3 month deprecation period specified by the [API Compatiblity Policy](https://jax.readthedocs.io/en/latest/api_compatibility.html).

PiperOrigin-RevId: 591933493
 2023-12-18 10:08:47 -08:00
Jake VanderPlas
a52d18781e Add experimental static key reuse checking 2023-12-11 12:03:48 -08:00
Roy Frostig
ed9a4c2939 add jax.threefry_partitionable context manager 2023-10-31 13:45:55 -07:00
Sergei Lebedev
f2ce5dbd01 MAINT Do not use str() and repr() in f-string replacement fields 
`str()` is called by default by the formatting machinery, and `repr()` only
needs `!r`.
 2023-10-23 15:12:04 +01:00
Jake VanderPlas
024b1f23d7 Remove deprecated submodule jax.abstract_arrays 2023-09-19 15:40:18 -07:00
Jake VanderPlas
1800015884 Import jax.version first 2023-09-12 12:27:20 -07:00
Jake VanderPlas
ca39457ea9 JEX: move jax.linear_util to jax.extend.linear_util 2023-08-30 18:32:12 -07:00
Muhammad Abdullah
599b35eeaa
Update __init__.py 2023-08-28 10:13:16 +05:00
Muhammad Abdullah
d09c55aa2d
Update __init__.py to include dlpack module 
import for dlpack module was missing in `__init__.py` file. just added that
 2023-08-28 10:01:17 +05:00
Jake VanderPlas
630a69f41b [random] add jax_legacy_prng_key flag 2023-08-22 15:08:51 -07:00
Skye Wanderman-Milne
8b58e38ec5 Add jax_debug_log_modules config option. 
This can be used to enable debug logging for specific files
(e.g. `JAX_DEBUG_LOG_MODULES="jax._src.xla_bridge,jax._src.dispatch"`)
or all jax (`JAX_DEBUG_LOG_MODULES="jax"`).

Example output:
```
$ JAX_DEBUG_LOG_MODULES=jax python3 -c "import jax; jax.numpy.add(1,1)"
DEBUG:2023-06-07 00:27:57,399:jax._src.xla_bridge:352: No jax_plugins namespace packages available
DEBUG:2023-06-07 00:27:57,488:jax._src.path:29: etils.epath found. Using etils.epath for file I/O.
DEBUG:2023-06-07 00:27:57,663:jax._src.dispatch:272: Finished tracing + transforming fn for pjit in 0.0005719661712646484 sec
DEBUG:2023-06-07 00:27:57,664:jax._src.xla_bridge:590: Initializing backend 'tpu'
DEBUG:2023-06-07 00:28:00,502:jax._src.xla_bridge:602: Backend 'tpu' initialized
DEBUG:2023-06-07 00:28:00,502:jax._src.xla_bridge:590: Initializing backend 'cpu'
DEBUG:2023-06-07 00:28:00,542:jax._src.xla_bridge:602: Backend 'cpu' initialized
DEBUG:2023-06-07 00:28:00,544:jax._src.interpreters.pxla:1890: Compiling fn for with global shapes and types [ShapedArray(int32[], weak_type=True), ShapedArray(int32[], weak_type=True)]. Argument mapping: (GSPMDSharding({replicated}), GSPMDSharding({replicated})).
DEBUG:2023-06-07 00:28:00,547:jax._src.dispatch:272: Finished jaxpr to MLIR module conversion jit(fn) in 0.0023522377014160156 sec
DEBUG:2023-06-07 00:28:00,547:jax._src.xla_bridge:140: get_compile_options: num_replicas=1 num_partitions=1 device_assignment=[[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)]]
DEBUG:2023-06-07 00:28:00,571:jax._src.dispatch:272: Finished XLA compilation of jit(fn) in 0.023587703704833984 sec
```
 2023-07-28 18:11:12 +00:00
Jake VanderPlas
2691d7edb9 Use standard framework for jax.tree* deprecation. 2023-07-20 12:58:17 -07:00
Peter Hawkins
651f87733b Remove jax_jit_pjit_api_merge. 
PiperOrigin-RevId: 548236671
 2023-07-14 15:25:00 -07:00
Jake VanderPlas
21f6736005 Remove several deprecated APIs 2023-07-11 12:42:32 -07:00
Jake VanderPlas
47ae5bddd7 Mark jax.abstract_arrays as deprecated 2023-06-07 23:36:40 -07:00
Matthew Johnson
e0d2736e37 add custom_jvp for jax.nn.softmax 
This avoids saving the jnp.exp(...) value.
 2023-04-22 11:28:03 -07:00
Peter Hawkins
a377caec3a Import jax.experimental.compilation_cache.compilation_cache by default. 
This is to fix users who were relying on this module being imported as part of 'import jax'.

PiperOrigin-RevId: 525151996
 2023-04-18 08:19:45 -07:00
Peter Hawkins
f8fe5d0542 Import jax.experimental.compilation_cache by default 
PiperOrigin-RevId: 525033643
 2023-04-17 21:28:35 -07:00
Matthew Johnson
26562a4382 [JAX] Add jax.clear_caches, plumb a way to clear pmap caches 
fixes #10828

Co-authored-by: Roy Frostig <frostig@google.com>
PiperOrigin-RevId: 522654093
 2023-04-07 12:19:00 -07:00
Peter Hawkins
c7b99e6ea9 Import jax.monitoring by default. 
A JAX refactoring meant this was no longer being imported by default. Restore the previous state.

PiperOrigin-RevId: 522474571
 2023-04-06 17:03:38 -07:00
Peter Hawkins
bf50551e0f Explicitly import jax.custom_{batching,derivatives,transpose}. 
https://github.com/google/jax/pull/15391 had the unintentional side effect of causing these names not to be imported by default. Restore the status quo by importing them.

PiperOrigin-RevId: 521898088
 2023-04-04 16:40:15 -07:00
Peter Hawkins
c1f65fc8b2 Avoid imports from the public jax.* namespace in more places internally. 
This change is in preparation for more cycle breaking in the Bazel dependency graph.

PiperOrigin-RevId: 521822756
 2023-04-04 11:41:40 -07:00
Peter Hawkins
abf1acf76c Replace references to jax.interpreters with jax._src.interpreters in JAX core. 
PiperOrigin-RevId: 520933067
 2023-03-31 08:58:00 -07:00