carlosgmartin 57b89ba7cb Added scipy.stats.gennorm.

2022-06-14 13:38:24 -04:00

..

_src

Added scipy.stats.gennorm.

2022-06-14 13:38:24 -04:00

example_libraries

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

experimental

Make Shard a normal class making index and replica_id properties on that class. Raise an error if an indices cannot be calculated from a sharding.

2022-06-14 10:37:19 -07:00

image

CI: fix flake8 ignore declarations

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

interpreters

Allow optional out_axis_resources for jax.Array. I'll look into making this enabled for all types in a follow up CL. If out_axis_resources is _UNSPECIFIED, then we let classic spmd choose the sharding for us. If you want to use the AUTO spmd partitioner, then pass pjit.AUTO to out_axis_resources explicitly.

2022-06-10 19:12:32 -07:00

lax

CI: fix flake8 ignore declarations

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

lib

Reexport jaxlib.__version as jax.lib.__version__.

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

nn

CI: fix flake8 ignore declarations

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

numpy

CI: fix flake8 ignore declarations

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

ops

CI: fix flake8 ignore declarations

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

scipy

Added scipy.stats.gennorm.

2022-06-14 13:38:24 -04:00

tools

Drop Bazel visibility of //jaxlib.

2022-04-28 13:16:54 -07:00

__init__.py

Add a public facing named_scope function to allow adding to the name stack.

2022-06-08 17:23:57 -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

cloud_tpu_init.py

CI: fix flake8 ignore declarations

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

collect_profile.py

Add collect_profile script

2022-06-03 17:56:17 -07:00

config.py

CI: fix flake8 ignore declarations

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

core.py

Initialize the thread-local compilation context when undefined in new threads.

2022-05-31 12:57:48 -07:00

custom_batching.py

CI: fix flake8 ignore declarations

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

custom_derivatives.py

CI: fix flake8 ignore declarations

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

custom_transpose.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -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

Add an undocumented method on jit() functions to clear the function cache.

2022-06-10 18:36:18 -07:00

prng.py

CI: fix flake8 ignore declarations

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

profiler.py

Store profiler server as a global variable and add a stop_server function

2022-06-01 17:50:06 -07:00

py.typed

Improve type annotations

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

random.py

Added random.generalized_normal and random.ball.

2022-06-03 15:11:29 -04:00

stages.py

CI: fix flake8 ignore declarations

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

test_util.py

CI: fix flake8 ignore declarations

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

tree_util.py

CI: fix flake8 ignore declarations

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

util.py

CI: fix flake8 ignore declarations

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

version.py

Update version + CHANGELOG for jax 0.3.13 release

2022-05-16 12:17:07 -07:00