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 83383fc717 Error on numpy array conversion of PRNG key array
2024-11-07 10:08:49 -08:00
..
clusters
Avoid querying metadata query to check if it's GCE if TPU_SKIP_MDS_QUERY is set.
2024-10-23 10:09:02 -07:00
cudnn
Simplified conditional in flash attention.
2024-10-31 16:28:11 -07:00
debugger
Cleanup: fix unused imports & mark exported names
2024-10-16 17:42:41 -07:00
export
[export] Cleanup in the export module.
2024-11-01 22:56:44 -07:00
extend
Add support for layouts and other advanced features in ffi_call.
2024-10-25 12:31:07 -04:00
image
…
internal_test_util
[export] Cleanup in the export module.
2024-11-01 22:56:44 -07:00
interpreters
Reverts f281c6f46475270a57a02416469226315377592c
2024-11-05 07:17:14 -08:00
lax
Reverts f281c6f46475270a57a02416469226315377592c
2024-11-05 07:17:14 -08:00
lib
[Mosaic GPU] Ensure that the dialect module can be loaded successfully.
2024-11-05 00:47:21 -08:00
nn
Remove the initial argument to jax.nn.softmax and jax.nn.log_softmax.
2024-11-04 10:55:21 -08:00
numpy
Fix debug_nans false positive in jnp.quantile
2024-11-05 15:36:14 -08:00
ops
Reverts 6cf09f8c24c67ff650b95d174501fff3cb59db0d
2024-10-04 13:56:27 -07:00
pallas
Reverts f281c6f46475270a57a02416469226315377592c
2024-11-05 07:17:14 -08:00
scipy
Merge pull request #24723 from jakevdp:beta-dep
2024-11-06 09:42:29 -08:00
state
Reverts f281c6f46475270a57a02416469226315377592c
2024-11-05 07:17:14 -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
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
ad_util.py
Reverts f281c6f46475270a57a02416469226315377592c
2024-11-05 07:17:14 -08:00
api_util.py
Consolidate handling of input argument resolution in custom_* APIs.
2024-08-13 00:30:23 -07:00
api.py
Reverts a913fbf2fddc5b8c1b6c85b159d0eeb1bf65d461
2024-11-05 08:32:25 -08:00
array.py
Remove ConcreteArray from JAX. It's easy to do trace-time concretization without it.
2024-10-31 14:06:54 -07: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
…
cache_key.py
Remove silent data corruption runtime flags from persistent cache key.
2024-10-24 10:59:44 -07:00
callback.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
checkify.py
Add compiler_options argument to jax.jit.
2024-11-01 14:01:19 -07:00
cloud_tpu_init.py
Fix TPU_ML_PLATFORM and TPU_ML_PLATFORM_VERSION environment variables.
2024-10-02 15:14:50 -07:00
compilation_cache_interface.py
Import etils.epath lazily
2024-08-13 14:48:38 -07:00
compilation_cache.py
declare buffer donation and compilation cache support for platform "neuron"
2024-09-18 17:07:46 -07:00
compiler.py
support exec_time_optimization_effort and memory_fitting_effort xla compilation
2024-11-01 16:25:50 -07:00
compute_on.py
Reverts a913fbf2fddc5b8c1b6c85b159d0eeb1bf65d461
2024-11-05 08:32:25 -08:00
config.py
Adding JAX_LOGGING_LEVEL configuration option
2024-11-05 09:56:46 -08:00
core.py
Make Tracers non-hashable
2024-11-05 09:08:33 -08:00
custom_api_util.py
…
custom_batching.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_derivatives.py
Reverts f281c6f46475270a57a02416469226315377592c
2024-11-05 07:17:14 -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
Remove some obsolete deprecation registrations
2024-11-06 11:10:04 -08:00
dispatch.py
[sharding_in_types] Allow device_put to reshard inputs. device_put is a good choice for resharding since it already handles transpose correctly because it tracks the src sharding too.
2024-11-01 13:25:08 -07:00
distributed.py
[jax.distributed] Enable grpc channel compression
2024-11-05 16:47:29 +00:00
dlpack.py
Deprecated calling `jax.dlpack.from_dlpack` with a DLPack tensor
2024-09-03 15:16:02 -07:00
dtypes.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
earray.py
Make committed a public property of jax.Array.
2024-10-15 19:46:10 -07:00
effects.py
…
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
…
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
…
lax_reference.py
…
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
Remove ConcreteArray from JAX. It's easy to do trace-time concretization without it.
2024-10-31 14:06:54 -07: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
Add support for tpu v5e to jax.make_mesh
2024-11-06 08:13:46 -08:00
mesh.py
Reverts a913fbf2fddc5b8c1b6c85b159d0eeb1bf65d461
2024-11-05 08:32:25 -08:00
monitoring.py
…
op_shardings.py
…
partition_spec.py
…
path.py
Import etils.epath lazily
2024-08-13 14:48:38 -07:00
pickle_util.py
…
pjit.py
Add compiler_options argument to jax.jit.
2024-11-01 14:01:19 -07: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
Error on numpy array conversion of PRNG key array
2024-11-07 10:08:49 -08:00
profiler.py
Small fix for the jax trace dumping path
2024-08-26 10:51:08 -07:00
public_test_util.py
…
random.py
[shape_poly] Expands support for random.choice
2024-10-24 17:20:09 +03:00
shard_alike.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
sharding_impls.py
Add support for tpu v5e to jax.make_mesh
2024-11-06 08:13:46 -08:00
sharding_specs.py
…
sharding.py
[Shardy] Inline meshes when using shardy and get rid of global meshes from the MLIR body.
2024-10-14 10:08:04 -07:00
source_info_util.py
…
sourcemap.py
…
stages.py
sharding cleanup: use inline checks for unimplemented and auto
2024-10-25 04:22:40 -07:00
test_util.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
tpu_custom_call.py
[Mosaic] Set TPU CustomCall device type based on the core_type attribute
2024-11-01 10:02:49 -07:00
traceback_util.py
…
tree_util.py
Fix errors when prefix does not match pytree
2024-09-16 22:25:28 -07:00
tree.py
…
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
Collapsed a few unnecessary `if TYPE_CHECKING` blocks
2024-08-12 13:08:55 +01:00
xla_bridge.py
Raise an error when registering a lowering for an unknown platform
2024-10-22 13:29:48 -04:00
xla_metadata.py
Reverts a913fbf2fddc5b8c1b6c85b159d0eeb1bf65d461
2024-11-05 08:32:25 -08:00