Peter Hawkins 7ba36fc178 Change implementation of jax.scipy.linalg.polar() and jax._src.scipy.eigh to use the QDWH decomposition from jax._src.lax.qdwh. ...

Remove jax._src.lax.polar.

PiperOrigin-RevId: 448241206

2022-05-12 07:20:52 -07:00

..

image

Merge pull request #9689 from fehiepsi:image

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

lax

Change implementation of jax.scipy.linalg.polar() and jax._src.scipy.eigh to use the QDWH decomposition from jax._src.lax.qdwh.

2022-05-12 07:20:52 -07:00

lib

Merge remaining CUDA and ROCM Python code.

2022-05-06 09:35:01 -07:00

nn

Use lax.broadcasted_iota in jax.nn.one_hot.

2022-04-20 09:09:19 -04:00

numpy

Change non-array arguments to jax.lax.linalg functions to be keyword-only arguments.

2022-05-11 13:06:54 -07:00

ops

lax_numpy: factor out indexing tricks

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

scipy

Change implementation of jax.scipy.linalg.polar() and jax._src.scipy.eigh to use the QDWH decomposition from jax._src.lax.qdwh.

2022-05-12 07:20:52 -07: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

[remove-units] remove units

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

api_util.py

fix grad(..., argnums=-1), regressed in #10453

2022-05-11 11:19:22 -07:00

api.py

Add an option when lowering to not remove unused arguments.

2022-05-04 01:22:14 -07:00

checkify.py

Internal change

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

cloud_tpu_init.py

Fixed a typo in cloud_tpu_init.py

2022-05-11 21:19:28 +00:00

config.py

use upgrade option in defining our latest upgrade flag

2022-05-10 11:25:23 -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

[remove-units] remove units

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

custom_transpose.py

[MHLO] Remove most XLA translation rules.

2022-04-18 08:28:35 -07:00

debugging.py

Enable AD rules for debug_print

2022-05-05 10:11:58 -07:00

device_array.py

Re-land changes from https://github.com/google/jax/pull/10069

2022-04-04 12:18:43 -07:00

dispatch.py

[remove-units] remove units

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

distributed.py

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

2022-05-05 09:33:06 -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

dtypes.result_type: add optional return_weak_type_flag argument

2022-04-29 12:54:40 -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

add some simple iree tests

2022-04-14 10:55:00 -07:00

lax_reference.py

Improve implementation of cbrt() in JAX.

2021-07-22 14:01:28 -07:00

pretty_printer.py

add flag for jaxpr colorful syntax highlighting

2022-02-12 14:15:28 -08:00

prng.py

Consolidate hip_prng and cuda_prng.

2022-05-05 10:55:29 -07:00

profiler.py

Add profiler annotations

2021-12-06 15:13:35 -08: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

Added random.orthogonal.

2022-04-29 14:20:50 -04:00

source_info_util.py

Increase minimum jaxlib version to 0.3.0.

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

stages.py

Expose the input and output sharding on the compiled object.

2022-04-13 10:18:25 -07:00

test_util.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

traceback_util.py

Increase minimum jaxlib version to 0.3.0.

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

tree_util.py

Merge pull request #9569 from GJBoth:tree_flatten_docs

2022-04-14 16:09:47 -07:00

util.py

Comment on implementation of unzip2 & unzip3

2022-04-14 13:41:05 -07:00