Sergei Lebedev ec745f48c8 Use the current minimum jaxlib version for type checking on the CI

2024-10-10 12:46:15 +01:00

..

clusters

Merge pull request #24032 from Pixee-Bot-Python:main

2024-10-02 01:25:50 -07:00

cudnn

fix bwd batcher for unsupported dbias

2024-09-23 17:43:25 +00:00

debugger

Fix error in debugger tests that is showing up in CI.

2024-09-18 15:44:45 +00:00

export

Activate FFI implementation of symmetric Eigendecomposition.

2024-10-04 12:38:26 -07:00

extend

Deprecate the vectorized argument to pure_callback and ffi_call.

2024-10-02 11:33:51 -04:00

image

Run pyupgrade --py310-plus.

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

internal_test_util

Reverts 6cf09f8c24c67ff650b95d174501fff3cb59db0d

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

interpreters

Error out in donation if:

2024-10-08 11:02:17 -07:00

lax

[sharding_in_types] Add support for nary ops to propagate sharding when 1 input is sharded and all others are replicated.

2024-10-09 21:24:37 -07:00

lib

Merge pull request #24139 from hartikainen:fix-cuda_path

2024-10-07 12:02:29 -07:00

nn

Merge pull request #23740 from kaixih:dbias_bwd_batcher

2024-10-02 14:04:19 -07:00

numpy

Merge pull request #24167 from rajasekharporeddy:testbranch1

2024-10-09 10:23:44 -07:00

ops

Reverts 6cf09f8c24c67ff650b95d174501fff3cb59db0d

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

pallas

Use the current minimum jaxlib version for type checking on the CI

2024-10-10 12:46:15 +01:00

scipy

Fix dtype bug in jax.scipy.fft.idct

2024-09-25 12:55:43 -04:00

state

[Pallas/MGPU] Add support for tiled and swizzled loads/stores + support slices

2024-10-02 02:44:10 -07: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

Removed various ununsed functions

2024-08-01 11:18:19 +01:00

ad_checkpoint.py

Fix printing of saved_residual for jit by looking for pjit as the primitive instead of xla_call which was removed 2 years ago

2024-09-24 19:01:19 -07:00

ad_util.py

Add a utility function to create a tangent zero value from a primal value.

2024-09-19 09:42:12 -07:00

api_util.py

Consolidate handling of input argument resolution in custom_* APIs.

2024-08-13 00:30:23 -07:00

api.py

Clear caches on jax exit.

2024-10-04 05:55:30 -07:00

array.py

Make make_array_from_process_local_data go via device_put if there is only 1 process.

2024-09-21 10:23:20 -07:00

basearray.py

Improve the docstring for jax.Array.copy_to_host_async.

2024-09-09 14:04:17 -07:00

basearray.pyi

Fix pytype errors and args for jax.Array methods

2024-09-03 17:06:45 +00:00

blocked_sampler.py

Run pyupgrade --py310-plus.

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

cache_key.py

Squashed commit of the following:

2024-09-10 17:30:08 -07:00

callback.py

Deprecate the vectorized argument to pure_callback and ffi_call.

2024-10-02 11:33:51 -04:00

checkify.py

add checkify rule for remat

2024-10-01 02:01:18 +00: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

Fix lint at HEAD.

2024-10-02 16:18:03 -04:00

compute_on.py

Add compute_on_context_manager to thread local jit state. This is to avoid getting false cache hits

2024-09-05 14:16:13 -07:00

config.py

Fix lint at HEAD.

2024-10-02 16:18:03 -04:00

core.py

Fix mutable array effects not being tracked properly

2024-09-30 18:55:15 -07:00

custom_api_util.py

…

custom_batching.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

custom_derivatives.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

custom_partitioning.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

custom_transpose.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04: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

Added jax.experimental.pallas.mosaic_gpu

2024-10-07 04:05:08 -07:00

dispatch.py

Cache the iteration over jaxpr equation when extracting shardings because majority of the time, it's the same jaxpr so we don't need to evaluate it again and again.

2024-10-03 20:47:59 -07:00

distributed.py

Clean up JAX backend for all backends to avoid dangling PyClient references.

2024-09-10 14:19:00 -07:00

dlpack.py

Deprecated calling `jax.dlpack.from_dlpack` with a DLPack tensor

2024-09-03 15:16:02 -07:00

dtypes.py

simplify conversion logic involving extended dtypes

2024-09-25 00:10:01 +00:00

earray.py

If input layouts are specified via in_shardings to jit and the array that the jitted function is called with is uncommitted, reshard the input array to the layout specified by the user.

2024-08-19 15:10:32 -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

Run pyupgrade --py310-plus.

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

lax_reference.py

…

layout.py

If input layouts are specified via in_shardings to jit and the array that the jitted function is called with is uncommitted, reshard the input array to the layout specified by the user.

2024-08-19 15:10:32 -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

Bumped mypy to 1.11.0 and jaxlib to 0.4.31 on the CI

2024-08-01 22:30:24 +01:00

logging_config.py

…

lru_cache.py

Bumped jaxlib version mypy uses on the CI

2024-07-26 11:22:39 +01:00

mesh_utils.py

Create optimal order for v5e:8 devices which is [0, 1, 2, 3, 7, 6, 5, 4]

2024-09-09 13:24:18 -07:00

mesh.py

Add a private _extremely_unsafe_enter_tracing_context to enter abstractMesh into tracing context. This is a temporary workaround for internal use cases.

2024-10-06 14:50:24 -07: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

Use the current minimum jaxlib version for type checking on the CI

2024-10-10 12:46:15 +01: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

simplify conversion logic involving extended dtypes

2024-09-25 00:10:01 +00: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

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

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

[sharding_in_types] Add support for nary ops to propagate sharding when 1 input is sharded and all others are replicated.

2024-10-09 21:24:37 -07:00

sharding_specs.py

…

sharding.py

Add num_devices to Sharding interface so that it works with NamedSharding containing AbstractMesh too.

2024-08-14 09:03:17 -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

Move DeviceAssignmentMismatchError exception catching code to def lower method of Traced so that all libraries calling traced.lower() see a better error message

2024-09-12 19:04:07 -07:00

test_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

tpu_custom_call.py

[Mosaic] Parameterize the number of lanes and sublanes in TPU dialects.

2024-10-10 04:28:36 -07:00

traceback_util.py

Run pyupgrade --py310-plus.

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

tree_util.py

`jax.tree_util.register_dataclass now validates data_fields and meta_fields`

2024-08-30 02:01:50 -07: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

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

xla_metadata.py

Add frontend attributes to Jax. This allows Jax users to annotate Jax code with frontend_attributes which can be traced down to the HLO level, to be used for numerical debugging purposes.

2024-09-06 16:44:56 -07:00