Peter Hawkins 885e8a2311 Don't recompute abstract eval rules when inlining a jit jaxpr. ...

The current implementation of jit inlining uses core.eval_jaxpr() and retraces the subjaxpr. This ends up performing abstract evaluation a second time. Instead, write a direct implementation of inlining that doesn't use the tracing machinery.

PiperOrigin-RevId: 607418006

2024-02-15 12:28:48 -08:00

..

clusters

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

cudnn

Merge pull request #18814 from Cjkkkk:spda

2024-02-12 16:11:37 -08:00

debugger

jax.debug.breakpoint now supports being DCE'd. Drive-by: fix #16186

2023-07-31 15:36:43 -07:00

extend

Upgrade most .py sources to 3.9

2023-12-08 12:23:15 +00:00

image

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

internal_test_util

Roll-back 1ae054b003088d873902fa62cfa8099260471e16 to re-enable nextafter tests

2024-01-26 11:08:43 -08:00

interpreters

Don't recompute abstract eval rules when inlining a jit jaxpr.

2024-02-15 12:28:48 -08:00

lax

Fix JVP rule for convert_element_type to integer types.

2024-02-15 14:09:04 +01:00

lib

Check if the Triton dialect bindings are available in lib/triton.py

2024-02-14 20:49:08 -08:00

nn

nn.softmax: use double-where when where is specified

2024-01-26 09:45:31 -08:00

numpy

[shape_poly] Add support for jnp.split

2024-02-15 14:43:41 +01:00

ops

jax.numpy: implement scalar boolean indexing

2024-02-09 11:00:00 -08:00

pallas

[jax:pallas:triton] Fix Triton lowering rule for broadcast_to with constant input.

2024-02-15 12:02:53 -08:00

scipy

Rename _wraps to implements

2024-01-24 14:14:19 -08:00

state

Slice.from_slice now works for slices with a negative start index

2024-01-29 13:13:06 +00:00

third_party

Rename _wraps to implements

2024-01-24 14:14:19 -08:00

__init__.py

…

abstract_arrays.py

Remove code that guards against int4 not being available.

2024-02-15 08:41:17 -08:00

ad_checkpoint.py

Don't compute the log message for checkpointing residuals unless it is going to be logged.

2024-02-13 18:29:29 -08:00

ad_util.py

partial rollback of #19096 due to internal breakage (relying on jax internals)

2023-12-22 15:54:26 -08:00

api_util.py

[attrs] allow passing a jax-attrs object to jit functions

2024-02-13 16:53:46 -08:00

api.py

document that under disable_jit, individual primitives are still compiled

2024-02-05 12:01:33 -08:00

array.py

Merge pull request #19737 from jakevdp:devices-dep

2024-02-09 13:09:41 -08:00

basearray.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

basearray.pyi

[typing] add where paramter in any/all type declaration

2024-01-16 08:29:14 -08:00

cache_key.py

Remove xla_gpu_cuda_data_dir compile option from persistent cache key.

2024-01-24 17:26:59 -08:00

callback.py

Add a registry for primitives that require device_assignment during lowering

2023-12-08 16:31:41 -08:00

checkify.py

integrate attrs in jax.jit

2024-01-27 17:44:43 -08:00

cloud_tpu_init.py

making sure enhanced barrier only turns on when there is a supported TPU available.

2024-01-12 23:47:37 +00:00

compilation_cache_interface.py

Replaced most usages of abc.ABC with util.StrictABC

2024-01-29 12:40:43 +00:00

compilation_cache.py

Check for errors when initializing the compilation cache.

2024-02-13 14:04:58 -08:00

compiler.py

To avoid the inconsistency between process_index and process_id, replace backend.process_index with distributed.global_state.process_id in Jax compilation _cache_write function.

2024-02-15 10:48:05 -08:00

config.py

Make config.value_holders private to avoid misuse.

2024-02-12 18:00:18 -08:00

core.py

Don't recompute abstract eval rules when inlining a jit jaxpr.

2024-02-15 12:28:48 -08:00

custom_api_util.py

…

custom_batching.py

integrate attrs in jax.jit

2024-01-27 17:44:43 -08:00

custom_derivatives.py

integrate attrs in jax.jit

2024-01-27 17:44:43 -08:00

custom_transpose.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

debugging.py

Cache ClosedJaxpr creation to minimize cache_misses. ClosedJaxpr should always be created under a cache.

2023-12-21 22:15:52 -08:00

deprecations.py

Add (private) mechanism for registering and accelerating deprecations

2024-02-07 12:28:03 -08:00

dispatch.py

Summarize equations in sharding mismatch errors lazily.

2024-02-13 18:29:47 -08:00

distributed.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

dlpack.py

Add support for bool dlpack values.

2024-01-17 09:30:42 -08:00

dtypes.py

Remove code that guards against int4 not being available.

2024-02-15 08:41:17 -08:00

effects.py

Effects is now an immutable set

2023-12-11 08:45:52 -08:00

environment_info.py

jax.print_environment_info: add platform.uname()

2024-02-13 11:23:56 -08:00

errors.py

Fix typos across the package

2023-09-22 14:54:31 -07:00

flatten_util.py

Fix duplicate word occurrences

2023-12-19 06:15:30 +01:00

gfile_cache.py

Fix FileExistsError for tests/gfile_cache_test.py::FileSystemCacheTest::test_threads

2023-06-14 08:07:22 +08:00

hardware_utils.py

making sure enhanced barrier only turns on when there is a supported TPU available.

2024-01-12 23:47:37 +00:00

jaxpr_util.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

lax_reference.py

test: fix dtype in test of lax.logistic

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

layout.py

Make the SpecifiedLayout class opaque.

2023-11-18 15:17:16 -08:00

lazy_loader.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

linear_util.py

Small optimization to trace_context.

2024-02-12 14:30:11 -08:00

logging_config.py

Add jax_debug_log_modules config option.

2023-07-28 18:11:12 +00:00

maps.py

integrate attrs in jax.jit

2024-01-27 17:44:43 -08:00

mesh.py

Add __str__ to Mesh so that in jaxprs the mesh doesn't print all the device ids.

2024-01-18 11:23:25 -08:00

monitoring.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

op_shardings.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

partition_spec.py

Improve the sharding documentation.

2023-08-03 11:12:02 -04:00

path.py

…

pickle_util.py

Use cloudpickle for serializing host callbacks

2023-06-09 13:34:43 -07:00

pjit.py

Don't recompute abstract eval rules when inlining a jit jaxpr.

2024-02-15 12:28:48 -08:00

pretty_printer.py

Replaced most usages of abc.ABC with util.StrictABC

2024-01-29 12:40:43 +00:00

prng.py

Define reuse_key primitive in jax._src.prng

2024-02-14 14:01:08 -08:00

profiler.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

public_test_util.py

Support float8 in reduce_min & reduce_max

2023-12-18 13:37:45 -08:00

random.py

jax.random: deprecate passing of batched keys to APIs

2024-01-17 12:53:24 -08:00

shard_alike.py

Guard shard_alike usage on xla_extension_version

2024-01-19 13:02:29 -08:00

sharding_impls.py

Add missing super().__init__() involving types wrapped in xla/python/sharding.cc

2024-01-30 17:17:12 -08:00

sharding_specs.py

Add a new experimental option jax_pmap_no_rank_reduction.

2024-01-19 03:53:37 -08:00

sharding.py

Remove indices and devices from shard_arg_handlers and shard_args.

2024-01-05 14:17:14 -08:00

source_info_util.py

Speed up user filename identification.

2024-02-13 11:04:47 -08:00

stages.py

Deprecate support for the mhlo dialect.

2024-01-15 02:13:40 -08:00

test_util.py

test_util: avoid overflow errors in NumPy 2.0

2024-02-07 12:39:31 -08:00

tpu_custom_call.py

Import submodules from jax._src explicitly, instead of relying on import side-effects. It will lead to the missing x-refs in code search according to go/pywald-sawmill-analysis.

2024-02-06 15:47:16 -08:00

traceback_util.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

tree_util.py

Add jax.tree module with aliases of jax.tree_util

2024-02-12 13:07:59 -08:00

tree.py

Add jax.tree module with aliases of jax.tree_util

2024-02-12 13:07:59 -08:00

typing.py

[shape_poly] Simplify the API for processing polymorphic_shape specifications

2023-11-28 12:45:59 +02:00

util.py

Small optimization to trace_context.

2024-02-12 14:30:11 -08:00

xla_bridge.py

Only register profiler when the plugin is dynamically loaded.

2024-02-02 13:48:26 -08:00