lenamartens 7078f81dd0 Checkify: misc improvements. ...

- err.throw == check_error(err) -> meaning they have the same behavior
  under checkify now
- "divided by zero" -> "division by zero"
- add validation that check_error only takes args of type Error

2022-09-23 14:33:06 +01:00

..

debugger

Internal change

2022-08-18 11:39:07 -07:00

image

[x64] make image_test compatible with strict dtype promotion

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

lax

[typing] overloaded type declaration for variadic lax.sort

2022-09-19 13:40:28 -07:00

lib

Bump the minimum jaxlib version to 0.3.15.

2022-09-08 16:43:46 -04:00

nn

add paper link about grad-relu-at-zero

2022-09-14 14:16:01 -07:00

numpy

Add a jit decorator to comparison operators.

2022-09-14 09:19:22 -07:00

ops

[typing] clear up logic in scatter_update

2022-09-14 15:32:52 -07:00

scipy

JAX implementation of scipy.stats.multinomial pmf & logpmf

2022-09-15 13:21:44 -07:00

state

Add support for closing over Refs in nested for loops

2022-09-13 13:32:44 -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

Fix Forward. The fix is on the user's end. Original PR: https://github.com/google/jax/pull/12217

2022-09-08 08:49:40 -07:00

ad_util.py

Add stateful for loop primitives (#10982)

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

api_util.py

ensure_index: raise better error for traced inputs

2022-09-12 12:28:14 -07:00

api.py

Implement pjit fast path in cpp for jax.Array inputs

2022-09-21 20:18:18 -07:00

callback.py

Use lowering as impl rule for pure_callback

2022-09-01 15:29:31 -07:00

checkify.py

Checkify: misc improvements.

2022-09-23 14:33:06 +01:00

cloud_tpu_init.py

Internal change

2022-08-18 11:39:07 -07:00

config.py

Flip default value of jax_unique_mhlo_module_names to False.

2022-09-21 09:48:01 -07:00

custom_api_util.py

enable custom transformation "stacking"

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

custom_batching.py

Add spmd_axis_name to vmap to allow constraining mapped PartitionSpecs.

2022-08-08 19:41:42 -07:00

custom_derivatives.py

Fix Forward. The fix is on the user's end. Original PR: https://github.com/google/jax/pull/12217

2022-09-08 08:49:40 -07:00

custom_transpose.py

[dynamic-shapes] initial support for dynamic shape typechecks

2022-06-17 14:57:19 -07:00

debugging.py

Add support for pmap sharding

2022-09-19 19:29:44 -07:00

device_array.py

[JAX] Deprecate .to_py() property on arrays. Implement __array__ instead.

2022-08-25 07:28:27 -07:00

dispatch.py

Don't tuple arguments passed to XLA:CPU

2022-09-22 01:29:14 -07:00

distributed.py

Bump the minimum jaxlib version to 0.3.15.

2022-09-08 16:43:46 -04:00

dlpack.py

Avoid import-time dependency on jax.experimental

2022-08-19 11:30:25 -07:00

dtypes.py

Silence some pytype errors.

2022-09-16 14:11:58 -07:00

environment_info.py

Add jax.print_environment_info()

2022-09-12 15:39:33 -07:00

errors.py

prepare to switch to new remat

2022-08-04 12:25:03 -07:00

flatten_util.py

Documentation improvements.

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

iree.py

[dynamic-shapes] add basic vmap-of-indexing support

2022-09-08 17:52:12 -07:00

lax_reference.py

Roll forward: Upgrade logistic into a primitive.

2022-09-07 06:06:56 -07:00

pretty_printer.py

Internal change

2022-08-18 11:39:07 -07:00

prng.py

Merge pull request #12297 from mattjj:computation-follows-data-prng

2022-09-08 14:57:31 -07:00

profiler.py

Add option to generate perfetto trace without generating link

2022-08-25 13:55:53 -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

require scalar key arrays when seeding, splitting, and folding

2022-09-01 13:02:47 -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

docstring for jax.stages.Wrapped

2022-09-01 18:31:38 -07:00

test_util.py

tests: fix is_valid_shape() function

2022-09-15 15:28:55 -07:00

traceback_util.py

Fix Python 3.11 compatibility problems.

2022-08-10 19:45:24 +00:00

tree_util.py

Bump the minimum jaxlib version to 0.3.15.

2022-09-08 16:43:46 -04:00

typing.py

Change case of typing.Dtype -> typing.DType

2022-09-14 15:03:55 -07:00

util.py

Use c++ weakref LRU cache implementation as a drop in replacement for jax's

2022-08-18 14:36:08 -07:00