George Necula 5983d385da [dynamic-shapes] Expand the handling of dynamic shapes for reshape and iota. ...

Also add more tests.

2022-07-05 12:14:15 +03:00

..

debugger

Add JAX debugger

2022-06-29 10:08:58 -07:00

image

[x64] make image_test compatible with strict dtype promotion

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

lax

[dynamic-shapes] Expand the handling of dynamic shapes for reshape and iota.

2022-07-05 12:14:15 +03:00

lib

Fix ModuleNotFoundError for phawkins only with version

2022-06-28 22:42:45 +00:00

nn

[x64] make nn_test pass with strict dtype promotion

2022-06-16 10:56:49 -07:00

numpy

[dynamic-shapes] Expand the handling of dynamic shapes for reshape and iota.

2022-07-05 12:14:15 +03:00

ops

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

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

scipy

jax.scipy.optimize: fix type inconsistency

2022-06-29 09:47:53 -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

[dynamic-shapes] basic linearize and grad working

2022-06-30 14:30:22 -07:00

ad_util.py

Add stateful for loop primitives (#10982)

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

api_util.py

feat: validate jit args

2022-05-18 21:54:47 +01:00

api.py

vmap: better errors for mismatched axis in keyword arguments

2022-06-29 14:31:03 -07:00

checkify.py

Checkify: add (checkify-of-)vmap-of-check.

2022-06-27 10:34:26 +01:00

cloud_tpu_init.py

adding os env to track JAX platform

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

config.py

[dynamic-shapes] Expand the handling of dynamic shapes for reshape and iota.

2022-07-05 12:14:15 +03: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 python callback on GPU

2022-06-01 12:15:07 -07:00

device_array.py

DeviceArray: Improve support for copy, deepcopy, and pickle

2022-05-19 12:00:58 -07:00

dispatch.py

[dynamic-shapes] Expand the handling of dynamic shapes for reshape and iota.

2022-07-05 12:14:15 +03:00

distributed.py

Save an on-demand checkpoint when any worker receives a preemption signal.

2022-07-01 12:45:30 -07: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

[x64] mention flag value in strict type promotion error

2022-06-21 10:55:24 -07:00

errors.py

Use imports relative to the jax package consistently, rather than .-relative imports.

2021-11-24 07:48:29 -08:00

flatten_util.py

Documentation improvements.

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

iree.py

[IREE] Allow backend selection via a flag.

2022-06-02 21:56:10 -04: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

Enable colors when we are using a terminal or IPython

2022-05-14 14:54:11 -07:00

prng.py

Fix overflow of large prng computation

2022-06-20 10:48:15 +02:00

profiler.py

Remove spurious single quote in profiler.py

2022-06-17 16:39:21 -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

handle unimplemented hlo_modules() on XLA executables

2022-07-01 23:40:52 -07:00

test_util.py

jax.test_util: remove deprecated test classes.

2022-06-27 11:04:50 -07:00

traceback_util.py

Increase minimum jaxlib version to 0.3.0.

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

tree_util.py

[JAX] Drop the private _process_pytree method from tree_util.

2022-06-17 13:50:50 -07:00

util.py

[dynamic-shapes] basic linearize and grad working

2022-06-30 14:30:22 -07:00