mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-23 06:56:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
6541a62099
rocm_jax/jax/_src
History
Jake VanderPlas 6541a62099 jax.core: deprecate a number of APIs
2024-12-10 11:11:32 -08:00
..
clusters
Adds a flag to control proxy env checking.
2024-11-16 00:00:59 -08:00
cudnn
jax.core: deprecate a number of APIs
2024-12-10 11:11:32 -08:00
debugger
Cleanup: fix unused imports & mark exported names
2024-10-16 17:42:41 -07:00
export
Merge pull request from gnecula:export_doc_call
2024-12-10 00:45:20 -08:00
extend
Allow 64-bit output types from ffi_call regardless of enable_x64 flag.
2024-11-11 15:01:53 -08:00
image
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04:00
internal_test_util
Activate Tridiagonal Reduction to XLA's FFI
2024-12-09 04:36:59 -08:00
interpreters
Merge pull request from jax-ml:mutable-arrays-ad
2024-12-10 06:36:02 -08:00
lax
Merge pull request from jax-ml:mutable-arrays-ad
2024-12-10 06:36:02 -08:00
lib
[Mosaic GPU] Ensure that the dialect module can be loaded successfully.
2024-11-05 00:47:21 -08:00
nn
Remove internal KeyArray alias
2024-11-20 10:30:12 -08:00
numpy
Fix type annotation for numpy.linalg.matrix_norm argument 'ord'.
2024-12-08 20:11:06 -05:00
ops
Reverts 6cf09f8c24c67ff650b95d174501fff3cb59db0d
2024-10-04 13:56:27 -07:00
pallas
jax.core: deprecate a number of APIs
2024-12-10 11:11:32 -08:00
scipy
Update the render documentation job to use the new self-hosted runners
2024-11-26 21:01:50 -08:00
state
Merge pull request from jax-ml:mutable-arrays-ad
2024-12-10 06:36:02 -08:00
third_party
feat(lib): add real-valued implementation of jax.scipy.special.fresnel
2024-09-03 09:50:19 +02:00
__init__.py
…
abstract_arrays.py
Remove ConcreteArray from JAX. It's easy to do trace-time concretization without it.
2024-10-31 14:06:54 -07:00
ad_checkpoint.py
remove vestigial ad.reducing_transposes table
2024-12-05 05:44:40 +00:00
ad_util.py
Reverts f281c6f46475270a57a02416469226315377592c
2024-11-05 07:17:14 -08:00
api_util.py
Use private names for args in api_util to avoid shadowing kwargs keys.
2024-12-09 11:24:11 -05:00
api.py
Remove _pjit_lower_cached cache. We can simplify the caching of jit as we have downstream caches and a cpp cache too.
2024-11-25 11:40:50 -08:00
array.py
Remove dead code after minimum jaxlib version bump to v0.4.36.
2024-12-09 07:35:05 -08:00
basearray.py
Make Tracers non-hashable
2024-11-05 09:08:33 -08:00
basearray.pyi
Make jnp.subtract a ufunc
2024-10-21 10:11:51 -07:00
blocked_sampler.py
Remove internal KeyArray alias
2024-11-20 10:30:12 -08:00
cache_key.py
Remove dead code after minimum jaxlib version bump to v0.4.36.
2024-12-09 07:35:05 -08:00
callback.py
Fix typos: Change 'arugments' to 'arguments'.
2024-11-13 15:58:45 -05:00
checkify.py
Stop using generators for linear_util transformations.
2024-11-13 13:47:07 -08:00
cloud_tpu_init.py
Fix KeyError recently introduced in cloud_tpu_init.py
2024-11-20 17:46:06 +00:00
compilation_cache_interface.py
Import etils.epath lazily
2024-08-13 14:48:38 -07:00
compilation_cache.py
[AutoPGLE] Explicitly ignore host callback pointers
2024-11-26 04:06:15 -08:00
compiler.py
Remove dead code after minimum jaxlib version bump to v0.4.36.
2024-12-09 07:35:05 -08:00
compute_on.py
Reverts a913fbf2fddc5b8c1b6c85b159d0eeb1bf65d461
2024-11-05 08:32:25 -08:00
config.py
Remove dead code after minimum jaxlib version bump to v0.4.36.
2024-12-09 07:35:05 -08:00
core.py
Add a freeze primitive to delimit ref lifetimes for AD.
2024-12-09 20:57:07 -05:00
custom_api_util.py
…
custom_batching.py
Add error message when using custom_vmap with reverse-mode AD, and add docstrings.
2024-12-09 11:17:44 -08:00
custom_derivatives.py
Stop using generators for linear_util transformations.
2024-11-13 13:47:07 -08:00
custom_partitioning_sharding_rule.py
[jax:custom_partitioning] Implement SdyShardingRule to support
2024-12-05 11:33:23 -08:00
custom_partitioning.py
Make trace dispatch purely a function of context rather than a function of both context and data. This lets us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind and so on.
2024-10-29 11:04:31 -07:00
custom_transpose.py
Make trace dispatch purely a function of context rather than a function of both context and data. This lets us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind and so on.
2024-10-29 11:04:31 -07:00
debugging.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
deprecations.py
Merge pull request from jakevdp:reduction-where
2024-12-05 11:17:15 -08:00
dispatch.py
Use JAX's default device instead of jax.devices()[0], if set.
2024-12-03 16:53:58 -08:00
distributed.py
Add an option to deactivate automatic cluster detection in jax.distributed.initialize().
2024-12-02 17:42:12 -08:00
dlpack.py
Deprecated calling `jax.dlpack.from_dlpack` with a DLPack tensor
2024-09-03 15:16:02 -07:00
dtypes.py
Merge pull request from apivovarov:float8_e4m3
2024-11-18 14:34:59 -08:00
earray.py
Propagate CopySemantics from python to C++ transfer APIs so that device_put works correctly in presence of copy/donate options that user specified.
2024-11-07 15:51:54 -08:00
effects.py
Effects is now an immutable set
2023-12-11 08:45:52 -08:00
environment_info.py
Update device info printed in print_environment_info to print the TPU kind also along with global and local device count.
2024-10-04 09:20:27 -07:00
errors.py
Consolidate material on PRNGs and add a short summary to Key Concepts.
2024-11-15 14:44:57 -08:00
flatten_util.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
hardware_utils.py
Enable the enhanced TPU launch barrier on all TPU generations.
2024-09-11 15:28:14 -07:00
jaxpr_util.py
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04:00
lax_reference.py
Implements Ragged Dot API
2024-05-11 06:40:18 -07:00
layout.py
sharding cleanup: use inline checks for unimplemented and auto
2024-10-25 04:22:40 -07:00
lazy_loader.py
lazy_loader.attach now only imports the submodule once on first access
2024-08-28 15:37:28 -07:00
linear_util.py
Stop using generators for linear_util transformations.
2024-11-13 13:47:07 -08:00
logging_config.py
Adding JAX_LOGGING_LEVEL configuration option
2024-11-05 09:56:46 -08:00
lru_cache.py
Bumped jaxlib version mypy uses on the CI
2024-07-26 11:22:39 +01:00
mesh_utils.py
[Take 2] Raise a better error message if anything other than a sequence of ints is passed to make_mesh or create_device_mesh
2024-11-28 09:24:20 -08:00
mesh.py
context manager methods for AbstractMesh to appease type checker.
2024-12-04 15:58:03 -08:00
monitoring.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
op_shardings.py
Removed noop # type: ignore comments
2024-05-19 21:01:29 +01:00
partition_spec.py
[sharding_in_types] Support shard_map with sharding in types. Right now only full manual mode is supported.
2024-11-14 09:58:03 -08:00
path.py
Import etils.epath lazily
2024-08-13 14:48:38 -07:00
pickle_util.py
Removed noop # type: ignore comments
2024-05-19 21:01:29 +01:00
pjit.py
Add a freeze primitive to delimit ref lifetimes for AD.
2024-12-09 20:57:07 -05:00
pretty_printer.py
Bumped mypy to 1.11.0 and jaxlib to 0.4.31 on the CI
2024-08-01 22:30:24 +01:00
prng.py
Lower threefry as an out-of-line MLIR function on TPU.
2024-11-15 08:49:35 -08:00
profiler.py
Small fix for the jax trace dumping path
2024-08-26 10:51:08 -07:00
public_test_util.py
Add float8_e4m3 and float8_e3m4 types support
2024-11-08 18:58:31 +00:00
random.py
Remove internal KeyArray alias
2024-11-20 10:30:12 -08:00
shard_alike.py
[SDY] add JAX lowering to Shardy ShardingGroupOp for shard_alike.
2024-11-08 11:02:50 -08:00
sharding_impls.py
[Take 2] Raise a better error message if anything other than a sequence of ints is passed to make_mesh or create_device_mesh
2024-11-28 09:24:20 -08:00
sharding_specs.py
Remove mesh_shape from ShardingSpec since it's not in use anymore.
2024-04-07 21:13:41 -07:00
sharding.py
[jax] Typing on common_devices_indices_map
2024-12-02 12:40:36 -08:00
source_info_util.py
Use a class with __slots__ instead of a NamedTuple in JaxprEqn and SourceInfo, which are two tuples we build frequently.
2024-07-01 19:18:58 -07:00
sourcemap.py
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04:00
stages.py
Simply abstract_mesh and device_context context managers and handle everything via their corresponding configs in config.py
2024-12-04 14:04:25 -08:00
test_util.py
[sharding_in_types] Use set_mesh API to trigger sharding_in_types instead of the config option.
2024-12-04 12:12:29 -08:00
tpu_custom_call.py
jax.core: deprecate a number of APIs
2024-12-10 11:11:32 -08:00
traceback_util.py
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04:00
tree_util.py
Remove dead code after minimum jaxlib version bump to v0.4.36.
2024-12-09 07:35:05 -08:00
tree.py
Fix annotation for jax.tree.transpose
2024-07-09 08:06:16 -07:00
typing.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
util.py
Add numpy.put_along_axis.
2024-11-14 15:23:26 -05:00
xla_bridge.py
Add an option to specify mock GPU topology
2024-11-12 08:36:27 -08:00
xla_metadata.py
Reverts a913fbf2fddc5b8c1b6c85b159d0eeb1bf65d461
2024-11-05 08:32:25 -08:00