mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 20:06:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src
History
Peter Hawkins 76cda0ae07 Update flags to use the ABSL typed flag API. 
Change flags to use the newer definition style where the flag is read via a typed FlagHolder object returned by the DEFINE_... function. The advantage of doing this is that `flag.value` has a type known to the type checker, rather than reading it as an attr out of a gigantic config dictionary.

For jax.config flags, define a typed FlagHolder object that is returned when defining a flag, matching the ABSL API.

Move a number of flags into the file that consumes them. There's no reason we're defining every flag in `config.py`.

This PR does not change the similar "state" objects in `jax.config`. Changing those is for a future PR.

PiperOrigin-RevId: 551604974
2023-07-27 12:15:58 -07:00
..
clusters
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
debugger
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
image
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
internal_test_util
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
interpreters
Change mhlo.is_same_data_across_replicas from unit attr to bool attr
2023-07-24 19:50:33 -07:00
lax
Merge pull request #16826 from mattjj:issue16805
2023-07-26 11:20:31 -07:00
lib
Enable building jaxlib w/ Mosaic
2023-07-26 03:59:30 -07:00
nn
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
numpy
Add implementation of jnp.put
2023-07-26 08:54:54 -07:00
ops
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
scipy
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
state
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
third_party
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
__init__.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
abstract_arrays.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
ad_checkpoint.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
ad_util.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
api_util.py
Copybara import of the project:
2023-07-24 14:38:20 -07:00
api.py
Copybara import of the project:
2023-07-24 14:38:20 -07:00
array.py
Copybara import of the project:
2023-07-24 14:38:20 -07:00
basearray.py
Split basearray into separate Bazel module.
2023-03-13 11:14:41 -07:00
basearray.pyi
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
callback.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
checkify.py
custom prng: introduce mechanism to identify key arrays by dtype
2023-07-21 12:27:32 -07:00
cloud_tpu_init.py
Issue warning if JAX_USE_PJRT_C_API_ON_TPU=false
2023-06-20 14:27:08 -07:00
compilation_cache_interface.py
Move implementation of compilation cache out of jax/experimental and into jax/_src.
2023-04-17 08:35:53 -07:00
compilation_cache.py
Add support for int fields to compiler_options.
2023-07-20 17:37:19 -07:00
config.py
Update flags to use the ABSL typed flag API.
2023-07-27 12:15:58 -07:00
core.py
Update flags to use the ABSL typed flag API.
2023-07-27 12:15:58 -07:00
custom_api_util.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
custom_batching.py
Fix type failures under an upcoming pytype change.
2023-04-07 07:10:09 -07:00
custom_derivatives.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
custom_transpose.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
debugging.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
deprecations.py
Make deprecation warnings warn at appropriate stacklevel
2023-05-23 14:43:38 -07:00
dispatch.py
Update flags to use the ABSL typed flag API.
2023-07-27 12:15:58 -07:00
distributed.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
dlpack.py
[JAX] Delete _DeviceArray and DeviceArray.
2023-03-29 15:07:14 -07:00
dtypes.py
Update flags to use the ABSL typed flag API.
2023-07-27 12:15:58 -07:00
effects.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
environment_info.py
Split _src files custom_api_util, deprecations, effects and environment_info into separate Bazel targets.
2023-03-10 12:26:05 -08:00
errors.py
errors: create TracerBoolConversionError for more targeted debugging tips
2023-06-21 01:41:45 -07:00
flatten_util.py
Avoid imports from the public jax.* namespace in more places internally.
2023-04-04 11:41:40 -07:00
gfile_cache.py
Fix FileExistsError for tests/gfile_cache_test.py::FileSystemCacheTest::test_threads
2023-06-14 08:07:22 +08:00
lax_reference.py
Add padding option "SAME_LOWER" for ticket https://github.com/google/jax/pull/14990
2023-03-23 15:50:16 -07:00
lazy_loader.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
linear_util.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
maps.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
mesh.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
monitoring.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
op_shardings.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
partition_spec.py
Move PartitionSpec into its own file (jax/_src/partition_spec.py).
2023-04-06 11:43:25 -07:00
path.py
Make the etils[epath] dependency of JAX optional.
2022-10-12 21:16:48 +00:00
pickle_util.py
Use cloudpickle for serializing host callbacks
2023-06-09 13:34:43 -07:00
pjit.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
pretty_printer.py
Update flags to use the ABSL typed flag API.
2023-07-27 12:15:58 -07:00
prng.py
Copybara import of the project:
2023-07-24 14:38:20 -07:00
profiler.py
Split _src/profiler into a separate BUILD target.
2023-03-10 11:38:46 -08:00
public_test_util.py
Update flags to use the ABSL typed flag API.
2023-07-27 12:15:58 -07:00
random.py
custom prng: introduce mechanism to identify key arrays by dtype
2023-07-21 12:27:32 -07:00
sharding_impls.py
Make _remake internal and add return type hints
2023-07-24 17:36:36 -07:00
sharding_specs.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
sharding.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
source_info_util.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
stages.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
test_util.py
Update flags to use the ABSL typed flag API.
2023-07-27 12:15:58 -07:00
tpu_custom_call.py
Enable building jaxlib w/ Mosaic
2023-07-26 03:59:30 -07:00
traceback_util.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
tree_util.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
typing.py
Copybara import of the project:
2023-07-24 14:38:20 -07:00
util.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
xla_bridge.py
Update flags to use the ABSL typed flag API.
2023-07-27 12:15:58 -07:00