mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 04:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src
History
Jake VanderPlas bbfd4f2c26 jax.numpy: implement scalar boolean indexing
2024-02-09 11:00:00 -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
Roll-back 1ae054b003088d873902fa62cfa8099260471e16 to re-enable nextafter tests
2024-01-26 11:08:43 -08:00
interpreters
Support activation offloading to host in JAX!
2024-02-05 17:29:47 -08:00
lax
[shape_poly] Add limited support for equality explicit constraints.
2024-02-08 10:09:47 +01:00
lib
Do not load Triton bindings eagerly in jax/lib/__init__.py
2024-01-13 03:01:02 -08:00
nn
nn.softmax: use double-where when where is specified
2024-01-26 09:45:31 -08:00
numpy
jax.numpy: implement scalar boolean indexing
2024-02-09 11:00:00 -08:00
ops
jax.numpy: implement scalar boolean indexing
2024-02-09 11:00:00 -08:00
pallas
[Pallas] Add tpu bitcast op to pallas primitives.
2024-02-09 10:44:33 -08:00
scipy
Rename _wraps to implements
2024-01-24 14:14:19 -08:00
state
Slice.from_slice now works for slices with a negative start index
2024-01-29 13:13:06 +00:00
third_party
Rename _wraps to implements
2024-01-24 14:14:19 -08: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
del add_any_p and zeros_like_p, replace aval-dispatched traceable
2023-12-21 17:04:21 -08:00
ad_checkpoint.py
Support activation offloading to host in JAX!
2024-02-05 17:29:47 -08:00
ad_util.py
partial rollback of #19096 due to internal breakage (relying on jax internals)
2023-12-22 15:54:26 -08:00
api_util.py
Always flatten args and kwargs together i.e. tree_flatten((args, kwargs)) so that we have a uniform in_tree structure everywhere.
2023-12-19 17:32:07 -08:00
api.py
document that under disable_jit, individual primitives are still compiled
2024-02-05 12:01:33 -08:00
array.py
jax.numpy: implement scalar boolean indexing
2024-02-09 11:00:00 -08:00
basearray.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
basearray.pyi
[typing] add where paramter in any/all type declaration
2024-01-16 08:29:14 -08:00
cache_key.py
Remove xla_gpu_cuda_data_dir compile option from persistent cache key.
2024-01-24 17:26:59 -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
integrate attrs in jax.jit
2024-01-27 17:44:43 -08:00
cloud_tpu_init.py
making sure enhanced barrier only turns on when there is a supported TPU available.
2024-01-12 23:47:37 +00:00
compilation_cache_interface.py
Replaced most usages of abc.ABC with util.StrictABC
2024-01-29 12:40:43 +00:00
compilation_cache.py
Cleanup deprecated compilation cache APIs.
2024-01-12 22:44:48 -08:00
compiler.py
Share autotune config between hosts.
2024-02-06 01:28:18 -08:00
config.py
Share autotune config between hosts.
2024-02-06 01:28:18 -08:00
core.py
[shape_poly] Add limited support for equality explicit constraints.
2024-02-08 10:09:47 +01: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
integrate attrs in jax.jit
2024-01-27 17:44:43 -08:00
custom_derivatives.py
integrate attrs in jax.jit
2024-01-27 17:44:43 -08:00
custom_transpose.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
debugging.py
Cache ClosedJaxpr creation to minimize cache_misses. ClosedJaxpr should always be created under a cache.
2023-12-21 22:15:52 -08:00
deprecations.py
Add (private) mechanism for registering and accelerating deprecations
2024-02-07 12:28:03 -08:00
dispatch.py
Remove indices and devices from shard_arg_handlers and shard_args.
2024-01-05 14:17:14 -08:00
distributed.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
dlpack.py
Add support for bool dlpack values.
2024-01-17 09:30:42 -08:00
dtypes.py
Replaced most usages of abc.ABC with util.StrictABC
2024-01-29 12:40:43 +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
Fix duplicate word occurrences
2023-12-19 06:15:30 +01:00
gfile_cache.py
Fix FileExistsError for tests/gfile_cache_test.py::FileSystemCacheTest::test_threads
2023-06-14 08:07:22 +08:00
hardware_utils.py
making sure enhanced barrier only turns on when there is a supported TPU available.
2024-01-12 23:47:37 +00:00
jaxpr_util.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
lax_reference.py
test: fix dtype in test of lax.logistic
2024-01-16 11:59:30 -08: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
add jax.explain_cache_misses tracing cache miss explanations
2023-12-26 21:54:27 -08:00
logging_config.py
Add jax_debug_log_modules config option.
2023-07-28 18:11:12 +00:00
maps.py
integrate attrs in jax.jit
2024-01-27 17:44:43 -08:00
mesh.py
Add __str__ to Mesh so that in jaxprs the mesh doesn't print all the device ids.
2024-01-18 11:23:25 -08: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
Merge pull request #19566 from mattjj:attrs-aqt
2024-01-30 15:51:00 -08:00
pretty_printer.py
Replaced most usages of abc.ABC with util.StrictABC
2024-01-29 12:40:43 +00:00
prng.py
remove PRNGKeyArray ABC
2024-01-29 12:41:26 -08:00
profiler.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
public_test_util.py
Support float8 in reduce_min & reduce_max
2023-12-18 13:37:45 -08:00
random.py
jax.random: deprecate passing of batched keys to APIs
2024-01-17 12:53:24 -08:00
shard_alike.py
Guard shard_alike usage on xla_extension_version
2024-01-19 13:02:29 -08:00
sharding_impls.py
Add missing super().__init__() involving types wrapped in xla/python/sharding.cc
2024-01-30 17:17:12 -08:00
sharding_specs.py
Add a new experimental option jax_pmap_no_rank_reduction.
2024-01-19 03:53:37 -08:00
sharding.py
Remove indices and devices from shard_arg_handlers and shard_args.
2024-01-05 14:17:14 -08:00
source_info_util.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
stages.py
Deprecate support for the mhlo dialect.
2024-01-15 02:13:40 -08:00
test_util.py
test_util: avoid overflow errors in NumPy 2.0
2024-02-07 12:39:31 -08:00
tpu_custom_call.py
Import submodules from jax._src explicitly, instead of relying on import side-effects. It will lead to the missing x-refs in code search according to go/pywald-sawmill-analysis.
2024-02-06 15:47:16 -08:00
traceback_util.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
tree_util.py
Remove a number of deprecated APIs
2023-12-18 10:08:47 -08:00
typing.py
[shape_poly] Simplify the API for processing polymorphic_shape specifications
2023-11-28 12:45:59 +02:00
util.py
Replaced most usages of abc.ABC with util.StrictABC
2024-01-29 12:40:43 +00:00
xla_bridge.py
Only register profiler when the plugin is dynamically loaded.
2024-02-02 13:48:26 -08:00