mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src
History
Yash Katariya 8bf3a86860 [roll forward 2] Remove the `jax_require_devices_during_lowering flag since it was temporary. Added the semi-breaking change to Changelog.md. 
Reverts b52bcc1639368069075284eefc763f824ca155f1

PiperOrigin-RevId: 590959383
2023-12-14 09:14:25 -08:00
..
clusters
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
debugger
jax.debug.breakpoint now supports being DCE'd. Drive-by: fix #16186
2023-07-31 15:36:43 -07:00
extend
Upgrade most .py sources to 3.9
2023-12-08 12:23:15 +00:00
image
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
internal_test_util
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
interpreters
[roll forward 2] Remove the `jax_require_devices_during_lowering flag since it was temporary. Added the semi-breaking change to Changelog.md.
2023-12-14 09:14:25 -08:00
lax
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
lib
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
nn
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
numpy
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
ops
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
pallas
Import new version of Triton
2023-12-13 01:24:24 -08:00
scipy
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
state
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
third_party
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00: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
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
ad_util.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
api_util.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
api.py
Upgrade most .py sources to 3.9
2023-12-08 12:23:15 +00:00
array.py
Deprecate arr.device_buffer and arr.device_buffers
2023-12-06 10:20:29 -08:00
basearray.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
basearray.pyi
Ensure sharding-related array properties are documented
2023-11-03 09:56:33 -07:00
cache_key.py
Remove the old cache-key generation code.
2023-12-12 16:34:32 -08:00
callback.py
Add a registry for primitives that require device_assignment during lowering
2023-12-08 16:31:41 -08:00
checkify.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
cloud_tpu_init.py
Remove calling configure_library_path during jax import and get libtpu path from libtpu_module.get_library_path().
2023-10-10 10:59:37 -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
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
compiler.py
Remove the old cache-key generation code.
2023-12-12 16:34:32 -08:00
config.py
Merge pull request #18949 from froystig:seed-offset
2023-12-13 10:18:40 -08:00
core.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00: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
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
custom_derivatives.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
custom_transpose.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
debugging.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
deprecations.py
Make deprecation warnings warn at appropriate stacklevel
2023-05-23 14:43:38 -07:00
dispatch.py
Add a flag jax_require_devices_during_lowering to control if physical devices are passed during lowering to stablehlo. This is temporary to unblock nvidia.
2023-12-12 13:34:09 -08:00
distributed.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
dlpack.py
Increase minimum jaxlib version to 0.4.19.
2023-11-17 09:38:31 -08:00
dtypes.py
Upgrade most .py sources to 3.9
2023-12-08 12:23:15 +00:00
effects.py
Effects is now an immutable set
2023-12-11 08:45:52 -08:00
environment_info.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
errors.py
Fix typos across the package
2023-09-22 14:54:31 -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
jaxpr_util.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
lax_reference.py
add an exp2 primitive and lax.exp2
2023-07-28 12:33:49 -07:00
layout.py
Make the SpecifiedLayout class opaque.
2023-11-18 15:17:16 -08:00
lazy_loader.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
linear_util.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
logging_config.py
Add jax_debug_log_modules config option.
2023-07-28 18:11:12 +00:00
maps.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
mesh.py
Upgrade most .py sources to 3.9
2023-12-08 12:23:15 +00:00
monitoring.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
op_shardings.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
partition_spec.py
Improve the sharding documentation.
2023-08-03 11:12:02 -04: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
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
pretty_printer.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
prng.py
introduce a config flag to control a random seed offset
2023-12-12 18:31:07 -08:00
profiler.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
public_test_util.py
Properly pack and unpack int4 arrays on CPU in PJRT.
2023-11-01 17:39:24 -07:00
random.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
sharding_impls.py
Add a registry for primitives that require device_assignment during lowering
2023-12-08 16:31:41 -08:00
sharding_specs.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
sharding.py
Make is_fully_addressable an abstract method and implement it on each concrete Sharding.
2023-09-05 17:28:22 -07:00
source_info_util.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
stages.py
Fix indexing bug when querying _input_layouts
2023-11-27 21:12:07 -08:00
test_util.py
Add experimental static key reuse checking
2023-12-11 12:03:48 -08:00
tpu_custom_call.py
Make lowering oblivious to real physical devices. Instead cache lowering on HloSharding only (which is based on logical device numbers)
2023-12-08 14:36:09 -08:00
traceback_util.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
tree_util.py
Merge pull request #18881 from superbobry:pyupgrade
2023-12-08 11:20:50 -08:00
typing.py
[shape_poly] Simplify the API for processing polymorphic_shape specifications
2023-11-28 12:45:59 +02:00
util.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
xla_bridge.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00