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

..

_src

Checkify: misc improvements.

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

example_libraries

Merge pull request #11542 from mattjj:remove-resnet50-example

2022-07-19 16:56:56 -07:00

experimental

Implement pjit fast path in cpp for jax.Array inputs

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

image

CI: fix flake8 ignore declarations

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

interpreters

Only copy_to_device if the indices match. Otherwise reshard the array if its uncommitted. This is important where you have 1 process per device.

2022-09-19 16:59:14 -07:00

lax

Roll forward: Upgrade logistic into a primitive.

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

lib

Reexport jaxlib.__version as jax.lib.__version__.

2022-04-28 10:25:06 -07:00

nn

Annotate nn.initializers

2022-08-04 23:17:40 -04:00

numpy

DOC: ensure that _wraps() generates correct links to wrapped functions

2022-07-21 11:12:35 -07:00

ops

CI: fix flake8 ignore declarations

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

scipy

JAX implementation of scipy.stats.multinomial pmf & logpmf

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

tools

Internal change

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

__init__.py

Delete soft_pmap as it has no users. Please use pjit or xmap if you do want soft_pmap.

2022-09-13 15:52:10 -07:00

abstract_arrays.py

CI: fix flake8 ignore declarations

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

ad_checkpoint.py

CI: fix flake8 ignore declarations

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

api_util.py

CI: fix flake8 ignore declarations

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

BUILD

Remove deprecated jax.experimental.optimizers

2022-08-09 08:50:59 -07:00

cloud_tpu_init.py

CI: fix flake8 ignore declarations

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

collect_profile.py

Add option to generate perfetto trace without generating link

2022-08-25 13:55:53 -07:00

config.py

CI: fix flake8 ignore declarations

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

core.py

Tracer: add missing __round__ and __reversed__ methods

2022-09-20 09:09:23 -07:00

custom_batching.py

CI: fix flake8 ignore declarations

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

custom_derivatives.py

Copybara import of the project:

2022-08-17 22:40:58 -07:00

custom_transpose.py

CI: fix flake8 ignore declarations

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

debug.py

Add function to visualize Shardings

2022-09-19 13:27:08 -07:00

distributed.py

Unbreak jax.distributed initialization.

2022-06-07 11:10:17 -07:00

dlpack.py

CI: fix flake8 ignore declarations

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

dtypes.py

CI: fix flake8 ignore declarations

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

errors.py

CI: fix flake8 ignore declarations

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

flatten_util.py

CI: fix flake8 ignore declarations

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

jaxpr_util.py

Fix or ignore some pytype errors.

2022-05-31 21:25:55 -07:00

linear_util.py

[dynamic-shapes] fix minor bint bugs

2022-07-19 16:38:40 -07:00

prng.py

CI: fix flake8 ignore declarations

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

profiler.py

jax.profiler: remove deprecated functions

2022-07-19 08:13:44 -07:00

py.typed

Improve type annotations

2021-01-13 10:26:35 -05:00

random.py

add key_data to jax.random for key array unwrapping

2022-08-31 09:23:11 -07:00

stages.py

write in-process AOT walkthrough doc

2022-09-02 13:02:25 -07:00

test_util.py

Internal change

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

tree_util.py

Remove deprecated function jax.tree_util.tree_multimap

2022-07-26 09:37:27 -07:00

util.py

CI: fix flake8 ignore declarations

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

version.py

Bump the minimum jaxlib version to 0.3.15.

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