rocm_jax

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

[shape_poly] Remove caching for the symbolic shape evaluator

2024-11-09 11:13:48 +02:00

extend

Add typing overloads for jax.extend.ffi.ffi_call() to aid type checkers

2024-11-08 14:23:01 -08:00

image

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

internal_test_util

[export] Cleanup in the export module.

2024-11-01 22:56:44 -07:00

interpreters

Put the set of current spmd axis names in the axis env instead of spelunking

2024-11-08 18:15:26 -08:00

lax

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

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 argmin docstring to not say "maximum"

2024-11-08 11:19:56 -05:00

ops

Reverts 6cf09f8c24c67ff650b95d174501fff3cb59db0d

2024-10-04 13:56:27 -07:00

pallas

Plumb dot dimension numbers into TPU matmul op.

2024-11-07 15:32:02 -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

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

array.py

Fix pre-commit to run on all files in CI.

2024-11-08 13:47:27 -05: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

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

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

Fix pre-commit to run on all files in CI.

2024-11-08 13:47:27 -05:00

core.py

Put the set of current spmd axis names in the axis env instead of spelunking

2024-11-08 18:15:26 -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

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

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

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

…

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

[key reuse] rename flag to jax_debug_key_reuse

2024-03-22 05:37:30 -07: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

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

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

partition_spec.py

Make UNCONSTRAINED's __str__ and __repr__ so it prints nicely in all cases

2024-05-30 17:41:32 -07: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 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

Merge pull request #24481 from jakevdp:key-array-error

2024-11-08 13:47:05 -08:00

profiler.py

Small fix for the jax trace dumping path

2024-08-26 10:51:08 -07:00

public_test_util.py

Added basic support for int2/uint2 dtypes to JAX

2024-07-04 04:13:24 -07:00

random.py

[shape_poly] Expands support for random.choice

2024-10-24 17:20:09 +03: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

Add support for tpu v5e to jax.make_mesh

2024-11-06 08:13:46 -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

[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

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

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

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

tree_util.py

Merge pull request #24700 from jakevdp:register-dataclass

2024-11-07 13:14:58 -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

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