jax authors 54a6e4dad3 Merge pull request #9422 from yotarok:signal_stft ...

PiperOrigin-RevId: 429377655

2022-02-17 12:46:12 -08:00

..

image

Added workaround for bug in XLA

2021-11-18 11:01:50 +02:00

lax

err on repeated axes to expand_dims, as numpy does

2022-02-17 11:27:20 -08:00

lib

Implement the JAX transfer guard API

2022-02-14 13:11:49 -08:00

nn

Add batch_axis to variance scaling initializers

2022-02-04 17:02:11 -08:00

numpy

Merge pull request #9609 from froystig:prng-array-stack

2022-02-17 10:25:29 -08:00

ops

[x64] ensure scatter functionality preserves weak_type

2021-11-30 15:43:06 -08:00

scipy

Merge pull request #9422 from yotarok:signal_stft

2022-02-17 12:46:12 -08:00

third_party

Add some functions for spectral analysis.

2022-02-17 15:59:24 +09: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

Move dtype canonicalization out of core.AbstractValue subclasses.

2021-12-07 06:13:07 -08:00

ad_checkpoint.py

Add a partial_eval_jaxpr_custom_rule for xmap

2022-02-15 03:44:19 -08:00

ad_util.py

Move jax.ad_util to jax._src.ad_util.

2021-06-07 14:51:34 -07:00

api_util.py

[JAX] Require exact type equality using is for static arguments.

2022-01-25 14:36:02 -08:00

api.py

Implement the JAX transfer guard API

2022-02-14 13:11:49 -08:00

cloud_tpu_init.py

Changes to make jax[tpu] work better in a docker container.

2021-07-12 17:42:46 -07:00

config.py

Merge branch 'main' into opt-barrier

2022-02-15 14:16:20 -08:00

custom_api_util.py

enable custom transformation "stacking"

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

custom_batching.py

introduce custom_batching.sequential_vmap

2022-01-20 21:37:58 -08:00

custom_derivatives.py

Merge pull request #9175 from froystig:custom-xform-wrappers-forward-attrs

2022-01-12 19:11:22 -08:00

custom_transpose.py

enable custom transformation "stacking"

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

device_array.py

Move dtype canonicalization out of core.AbstractValue subclasses.

2021-12-07 06:13:07 -08:00

dispatch.py

Try adding support for nesting with_sharding_constraint in MANUAL xmaps

2022-02-15 10:13:11 -08:00

distributed.py

Add jax.distributed.initialize for multi-host GPU.

2021-10-26 14:37:54 -07:00

dlpack.py

Add complex number DLPack support to JAX and TensorFlow.

2022-02-09 14:58:00 -08:00

dtypes.py

Only use config.x64_enabled as the memo cache key for canonicalize_dtype, not any other fields.

2022-01-18 13:24:15 -08: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

make ravel_pytree unraveler dtype-polymorphic

2021-12-14 14:47:35 -08:00

iree.py

[JAX:IREE] Fix bitrot in iree plugin for JAX.

2022-01-25 15:08:49 -08: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

implement jnp.expand_dims and jnp.stack for PRNGKeyArrays

2022-02-16 20:47:27 -08:00

profiler.py

Add profiler annotations

2021-12-06 15:13:35 -08:00

random.py

Add a warning to random.choice to notify users of the ill-defined behaviour when requesting more samples than non-zero probabilities and replace=False

2022-02-14 21:41:30 +01:00

source_info_util.py

Speed up source location computation when lowering a jaxpr to HLO/MHLO.

2022-02-08 16:17:40 -08:00

test_util.py

Run all tests with jax_traceback_filtering=off.

2022-02-15 02:42:58 -08:00

traceback_util.py

[jax] completely truncate trivial filtered tracebacks

2022-01-19 19:42:02 -08:00

tree_util.py

improve pjit in/out_axis_resources pytree errors

2022-02-08 16:23:15 -08:00

util.py

[remove units] make JaxprTrace.process_call not introduce units

2022-02-12 13:48:12 -08:00