Sharad Vikram 426c7356fb Guard has_explicit_device with xla_client version

2022-06-02 12:03:27 -07:00

..

image

Merge pull request #9689 from fehiepsi:image

2022-03-30 06:05:30 -07:00

lax

[mhlo] Make variadic Scatter legal

2022-06-02 09:50:22 -07:00

lib

Merge pull request #10910 from hawkinsp:macarm

2022-06-01 09:04:52 -07:00

nn

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

numpy

[x64] make jax.numpy reductions respect input dtypes

2022-06-01 16:24:36 -07:00

ops

lax_numpy: factor out indexing tricks

2022-03-17 11:05:45 -07:00

scipy

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

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

third_party

Add jax.scipy.linalg.funm

2022-05-02 21:46:41 +01: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

improve partial_eval_jaxpr_custom

2022-05-11 13:20:23 -07:00

ad_util.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

api_util.py

feat: validate jit args

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

api.py

Guard has_explicit_device with xla_client version

2022-06-02 12:03:27 -07:00

checkify.py

Internal change

2022-05-09 13:26:30 -07:00

cloud_tpu_init.py

* Make all arguments to distributed.initialize equal to None.

2022-05-17 10:53:54 -07:00

config.py

Merge pull request #9118 from skye:device_context_manager

2022-06-02 10:33:53 -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

Merge pull request #10756 from mattjj:10750

2022-05-18 15:56:13 -07:00

custom_transpose.py

[MHLO] Remove most XLA translation rules.

2022-04-18 08:28:35 -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

Merge pull request #9118 from skye:device_context_manager

2022-06-02 10:33:53 -07:00

distributed.py

Fix distributed system initialization.

2022-05-25 07:31:19 -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] Add to_complex_dtype utility function

2022-05-31 09:26:08 -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

2nd attempt at fixing IREE integration.

2022-06-01 21:54:56 -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

[x64] make jax.random compatible with jax_numpy_dtype_promotion=strict

2022-05-27 11:12:39 -07:00

profiler.py

Store profiler server as a global variable and add a stop_server function

2022-06-01 17:50:06 -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

[x64] make jax.random compatible with jax_numpy_dtype_promotion=strict

2022-05-27 11:12:39 -07:00

source_info_util.py

Increase minimum jaxlib version to 0.3.0.

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

stages.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

test_util.py

feat: validate jit args

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

traceback_util.py

Increase minimum jaxlib version to 0.3.0.

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

tree_util.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

util.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00