Sharad Vikram c34aa3933f Various debugger improvements ...

- disables globals
- can opt out of filtering frames
- can limit number of frames

2022-08-08 19:54:29 -07:00

..

debugger

Various debugger improvements

2022-08-08 19:54:29 -07:00

image

[x64] make image_test compatible with strict dtype promotion

2022-06-15 11:41:23 -07:00

lax

[mhlo] Add result type inference for mhlo.fft.

2022-08-08 10:32:07 -07:00

lib

Merge pull request #11672 from hawkinsp:visible

2022-08-08 08:38:02 -07:00

nn

Annotate nn.initializers

2022-08-04 23:17:40 -04:00

numpy

jnp.linalg.matrix_rank: support stacks of matrices

2022-08-08 11:16:34 -07:00

ops

[x64] deprecate unsafe type casting in scatter-update operations

2022-06-09 15:21:49 -07:00

scipy

Merge pull request #11556 from RuffaloLavoisier:tYpO

2022-07-22 10:13:10 -07:00

state

Refactor state out of for_loop

2022-08-01 15:26:55 -07:00

third_party

[x64] make jax.scipy.signal compatible with strict dtype promotion.

2022-06-21 09:28:46 -07:00

__init__.py

Move most of the implementation of jax.scipy into jax._src.scipy.

2020-10-16 17:04:25 -04:00

abstract_arrays.py

[remove-units] prevent cond partial eval from introducing units

2022-04-24 14:28:33 -07:00

ad_checkpoint.py

add api_boundary, fix flax failure with new-remat-by-default

2022-08-05 20:55:27 -07:00

ad_util.py

Add stateful for loop primitives (#10982)

2022-06-15 15:55:38 -07:00

api_util.py

Rolling forward #11768 after test failures caused roll-back (from use of np.empty).

2022-08-05 22:19:33 -07:00

api.py

Rolling forward #11768 after test failures caused roll-back (from use of np.empty).

2022-08-05 22:19:33 -07:00

checkify.py

Canonicalize all shardings to OpShardingSharding throughout pjit. Places where pspec is needed, parsed_flatten_op_sharding function is used to retrieve the pspec. The major places are global_to_local and local_to_global. Rest of the changes are just threading through OpShardingSharding.

2022-08-05 12:18:10 -07:00

cloud_tpu_init.py

adding os env to track JAX platform

2022-06-14 21:21:34 +00:00

config.py

fix ad_checkpoint.checkpoint caching issue

2022-07-29 19:59:28 -07:00

custom_api_util.py

enable custom transformation "stacking"

2022-01-11 17:55:08 -08:00

custom_batching.py

[remove-units] remove units

2022-05-04 10:58:56 -07:00

custom_derivatives.py

Use equality rather than identity when checking for float0

2022-06-23 11:46:20 -07:00

custom_transpose.py

[dynamic-shapes] initial support for dynamic shape typechecks

2022-06-17 14:57:19 -07:00

debugging.py

Enable receives in TPU callbacks and add tests

2022-08-08 11:42:16 -07:00

device_array.py

Add dlpack support to device_array and jax.numpy

2022-07-15 17:31:11 -07:00

dispatch.py

Rolling forward #11768 after test failures caused roll-back (from use of np.empty).

2022-08-05 22:19:33 -07:00

distributed.py

Update minimum jaxlib version to 0.3.14.

2022-07-08 00:36:02 +00:00

dlpack.py

[JAX] Split the "gpu" platform in internal JAX usage into separate "cuda" and "rocm" platforms.

2022-05-05 09:33:06 -07:00

dtypes.py

Rolling forward #11768 after test failures caused roll-back (from use of np.empty).

2022-08-05 22:19:33 -07:00

errors.py

prepare to switch to new remat

2022-08-04 12:25:03 -07:00

flatten_util.py

Documentation improvements.

2022-03-08 09:37:33 -05:00

iree.py

[dynamic-shapes] revive basic bounded int machinery, add tests

2022-07-06 22:31:26 -07:00

lax_reference.py

Use util.unzipN() in more places instead of zip(*args).

2022-06-01 16:28:31 -04:00

pretty_printer.py

pretty_printing: handle case where stdout is patched by a logger

2022-07-15 14:50:17 -07:00

prng.py

Add dtype arg to jnp.concatenate and update tests

2022-08-01 15:48:40 -07:00

profiler.py

jax.profiler: remove deprecated functions

2022-07-19 08:13:44 -07:00

public_test_util.py

Separate jax.test_util implementations into public and private sources.

2022-04-04 14:43:39 -07:00

random.py

random.choice: make compatible with strict promotion

2022-06-28 10:41:30 -07:00

source_info_util.py

Only wrap one level of the name-stack with transform names.

2022-06-10 03:50:26 -07:00

stages.py

[dynamic-shapes] add basic abstracted_axes support to jit(f, ...).lower(...)

2022-07-07 12:48:29 -07:00

test_util.py

prepare to switch to new remat

2022-08-04 12:25:03 -07:00

traceback_util.py

Increase minimum jaxlib version to 0.3.0.

2022-03-04 10:33:03 -05:00

tree_util.py

promise to flatten trees in left-to-right order

2022-07-28 19:28:20 -07:00

util.py

remove dead code: jax._src.util.taggedtuple

2022-07-25 15:14:25 -07:00