Jake VanderPlas 093b7032a8 Implement jnp.from* array creation functions

2022-03-29 10:52:47 -07:00

..

image

Documentation improvements.

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

lax

allow for recursive uses of custom_transpose

2022-03-26 12:09:15 -07:00

lib

Fix for hipsparse in ROCm.

2022-03-25 17:41:42 +00:00

nn

Rename jax.nn.normalize to standardize. Add normalize alias with DeprecationWarning.

2022-03-23 20:55:22 +00:00

numpy

Implement jnp.from* array creation functions

2022-03-29 10:52:47 -07:00

ops

lax_numpy: factor out indexing tricks

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

scipy

Fix a bug in fft helper appears when nperseg=1.

2022-03-28 14:51:54 +09: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 separate mechanism for threading name stacks to the lowering

2022-02-23 09:59:09 -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

Merge pull request #9967 from froystig:aot-arg-trees

2022-03-21 12:15:38 -07:00

cloud_tpu_init.py

Add TPU v4 support to cloud_tpu_init

2022-03-15 16:00:16 -07:00

config.py

add custom_vjp upgrade flag to jax module and fix doc rendering for upgrade flags

2022-03-28 20:36:35 -07: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

refine occurrences of the "forward-mode disallowed" error under new custom_vjp

2022-03-28 16:50:47 -07:00

custom_transpose.py

allow for recursive uses of custom_transpose

2022-03-26 12:09:15 -07:00

device_array.py

Fix core.pytype_aval_mappings for _DeviceArray

2022-03-10 13:18:34 -06:00

dispatch.py

protocols for the lowering and executable underlying Lowered and Compiled

2022-03-22 19:06:36 -07:00

distributed.py

Add jax.distributed and jax.dlpack to the docs.

2022-02-17 16:10:07 -05:00

dlpack.py

Documentation improvements.

2022-03-08 09:37:33 -05: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

Documentation improvements.

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

iree.py

Merge pull request #9623 from rsuderman:FixIreeJax

2022-02-24 17:51:09 -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

remove _broadcasting_shape_rule from public jax.lax module

2022-03-08 16:34:26 -08:00

profiler.py

Add profiler annotations

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

random.py

add random.loggamma and improve dirichlet & beta implementation

2022-03-21 08:33:11 -07:00

source_info_util.py

Increase minimum jaxlib version to 0.3.0.

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

stages.py

remove unreachable code in Lowered

2022-03-22 21:31:23 -07:00

test_util.py

Fix bug in jax.test_util._dtype

2022-03-04 13:22:52 -08:00

traceback_util.py

Increase minimum jaxlib version to 0.3.0.

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

tree_util.py

improve pjit in/out_axis_resources pytree errors

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

util.py

Update util.py

2022-03-23 12:26:09 -07:00