George Necula c70de6deed [better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo ...

Previously, we had two almost identical classes: `TracingDebugInfo` and
`JaxprDebugInfo`. The only difference was that `TracingDebugInfo` had
a thunk to return the result paths, while `JaxprDebugInfo` had the
result paths resolved to a tuple. The separation of these types
provided some clarity, but also led to code duplication and
required conversions as the debugging info goes from `WrappedFun`
to a `Jaxpr` and then to `WrappedFun` again.

2025-02-02 06:23:03 +02:00

..

clusters

Adds a flag to control proxy env checking.

2024-11-16 00:00:59 -08:00

cudnn

support head size of 256

2024-12-19 18:38:06 +00:00

debugger

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07:00

export

Internal: avoid adding _DimExpr to dtypes._weak_types

2025-01-29 09:11:02 -08:00

extend

Move jex.ffi to jax.ffi.

2024-12-29 13:06:19 +00:00

image

Run pyupgrade --py310-plus.

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

internal_test_util

Activate FFI implementation of tridiagonal reduction on GPU.

2025-01-10 09:28:15 -08:00

interpreters

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02:00

lax

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02:00

lib

[mosaic_gpu] Fixed mosaic_gpu-serde pass registration

2025-01-24 06:35:54 -08:00

nn

Merge pull request #25766 from carlosgmartin:nn_initializers_variance_scaling_mode_fan_geo_avg

2025-01-31 15:41:50 -08:00

numpy

delete unnecessary line

2025-01-31 12:44:14 -08:00

ops

fix grad(logsumexp) to produce 0s where where is False

2025-01-08 23:38:06 +00:00

pallas

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02:00

scipy

enable partitionable threefry by default

2025-01-13 22:46:24 -08:00

state

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02: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

Internal: use a single registry for abstractify APIs

2024-12-23 08:44:35 -08:00

ad_checkpoint.py

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02:00

ad_util.py

Cleanup: remove uses of no-op raise_to_shaped

2024-12-12 09:49:06 -08:00

api_util.py

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02:00

api.py

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02:00

array.py

Reverts bb951136e9b91a584bb422119ada76cc69c86024

2025-01-31 14:42:22 -08:00

basearray.py

Make Tracers non-hashable

2024-11-05 09:08:33 -08:00

basearray.pyi

[sharding_in_types] Rename .at[...].get(out_spec) to .at[...].get(out_sharding).

2025-01-16 18:56:52 -08: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

Revert https://github.com/jax-ml/jax/pull/25982 since callbacks can now use JAX functions.

2025-01-29 11:12:32 -05:00

checkify.py

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02: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

Fix cache init when JAX Array is created early (#25768)

2025-01-22 03:44:29 +00:00

compiler.py

Remove code that supported jaxlib < 0.5.

2025-01-17 14:22:27 -05:00

compute_on.py

Optimize implementation of the compute_on context manager.

2025-01-26 09:24:33 -08:00

config.py

Make JAX_CPU_COLLECTIVES_IMPLEMENTATION and JAX_NUM_CPU_DEVICES env vars

2025-01-28 17:17:56 -08:00

core.py

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02:00

custom_api_util.py

…

custom_batching.py

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02:00

custom_dce.py

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02:00

custom_derivatives.py

Add DCE rules for custom_jvp and custom_vjp.

2025-01-23 15:22:43 -05:00

custom_partitioning_sharding_rule.py

[jax:custom_partitioning] Support SdyShardingRule with multiple leading

2025-01-23 08:20:38 -08:00

custom_partitioning.py

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02: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

Revert https://github.com/jax-ml/jax/pull/25982 since callbacks can now use JAX functions.

2025-01-29 11:12:32 -05:00

deprecations.py

Add sph_harm_y to jax.scipy.special and deprecate sph_harm.

2025-01-09 12:53:00 -05:00

dispatch.py

[JAX] Optimize array shard reordering

2025-01-23 13:45:59 -08:00

distributed.py

improve runtime error message

2025-01-29 21:03:26 +00:00

dlpack.py

Pulled DLDeviceType to XLA backend mapping into a global

2025-01-29 11:38:50 -08:00

dtypes.py

Internal: avoid adding _DimExpr to dtypes._weak_types

2025-01-29 09:11:02 -08:00

earray.py

Internal: use a single registry for abstractify APIs

2024-12-23 08:44:35 -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

Consolidate material on PRNGs and add a short summary to Key Concepts.

2024-11-15 14:44:57 -08:00

ffi.py

Add a register_custom_type_id function to the GPU plugins.

2025-01-07 07:29:38 -08:00

flatten_util.py

Avoid calls to warnings.catch_warnings in JAX core code.

2024-12-20 15:43:03 -05: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

[better_errors] Expand the tests for debug_info

2025-01-22 16:49:16 +01:00

lax_reference.py

Add int4/uint4 support to bitcast_convert_type

2024-12-20 12:45:24 -08: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

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02: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

Fix type annotations for NumPy 2.2

2024-12-11 14:24:58 -08:00

mesh.py

[sharding_in_types] Upstream changes from defaulting sharding_in_types config to True experiment. There aren't a lot of failures in TGP but we can atleast upstream these changes until we work on the failures.

2025-01-28 11:04:42 -08:00

monitoring.py

Add JAX events that have time spans, not only durations.

2025-01-07 23:08:30 -08:00

op_shardings.py

Fix type annotations for NumPy 2.2

2024-12-11 14:24:58 -08:00

partition_spec.py

Allow multiple meshes for avals but in that case, just use empty_abstract_mesh instead of enabling computation follows data only for **Auto mode**.

2025-01-29 20:47:34 -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

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

2025-02-02 06:23:03 +02:00

pretty_printer.py

Make api_test.py work when test cases are run using multiple threads.

2025-01-08 14:09:07 -08:00

prng.py

Fix pmap with sharded typed prng key

2025-01-10 18:20:09 -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 docstrings for check_vjp and check_jvp

2025-01-23 17:31:13 -08:00

random.py

Fix dtype issue in stirling approximation

2025-01-31 14:13:02 -08:00

shard_alike.py

Remove internal uses of api_util.shaped_abstractify

2024-12-19 07:06:36 -08:00

sharding_impls.py

Allow multiple meshes for avals but in that case, just use empty_abstract_mesh instead of enabling computation follows data only for **Auto mode**.

2025-01-29 20:47:34 -08:00

sharding_specs.py

Fix type annotations for NumPy 2.2

2024-12-11 14:24:58 -08:00

sharding.py

[sharding_in_types] Add vmap + explicit sharding support. The main changes are:

2025-01-29 09:34:27 -08:00

source_info_util.py

Speed up name stack printing.

2025-01-22 16:24:16 -08:00

sourcemap.py

Run pyupgrade --py310-plus.

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

stages.py

Allow multiple meshes for avals but in that case, just use empty_abstract_mesh instead of enabling computation follows data only for **Auto mode**.

2025-01-29 20:47:34 -08:00

test_util.py

Merge pull request #26174 from skye:cpu_configs

2025-02-01 16:20:16 -08:00

test_warning_util.py

Switch to a new thread-safe utility for catching warnings.

2025-01-09 11:58:34 -05:00

tpu_custom_call.py

Remove all MHLO uses, replace it with StableHLO

2025-01-23 21:35:36 -08:00

traceback_util.py

Run pyupgrade --py310-plus.

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

tree_util.py

Add an option to simplify keystr output and use a custom separator.

2025-01-21 10:18:42 -08:00

tree.py

Add jax.tree shortcuts for .*_with_path calls, for convenience of users.

2024-12-12 15:13:32 -08: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

Remove device_context from trace_context because we don't need it there. We can get compilation cache misses (and tracing/lowering cache hit) naturally without putting concrete devices into trace_context.

2025-01-21 16:21:36 -08:00

xla_bridge.py

Make JAX_CPU_COLLECTIVES_IMPLEMENTATION and JAX_NUM_CPU_DEVICES env vars

2025-01-28 17:17:56 -08:00

xla_metadata.py

Fix CI failures in xla_metadata_test.

2025-01-30 17:28:54 -08:00