Skye Wanderman-Milne 9128ba0c74 Replace host_id with process_index terminology, take 2. ...

We're switching to the new terminology to avoid confusion in cases
where multiple jax processes are running on a single host, and each
process has a unique process_index/host_id.

This keeps aliases for the old `host_id` APIs for now, but these will
eventually be removed.

This was originally commited in
b77ef5138b631378e6a8ceb8bafc94fe91239bae, but reverted in
14acd070c2afb11c81fc91f43790577cd48cbf67 due to Google-internal test
failures from renaming the local_devices argument name. This change is
identical except it also adds staging for the argument name change.

2021-04-20 18:13:34 -07:00

..

image

DOC: fix some minor formatting issues

2021-01-28 15:20:02 -08:00

lax

Move vtile to batching.py, make it possible to add new BatchTraces

2021-04-20 08:33:01 -07:00

nn

[JAX] Move contents of jax/dtypes.py to jax/_src/dtypes.py.

2021-04-07 19:35:51 -07:00

numpy

Merge pull request #6515 from jakevdp:unique-complex

2021-04-20 16:13:11 -07:00

ops

Implement segment_prod, segment_max, segment_min

2021-04-09 12:06:51 -07:00

scipy

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

third_party

Move jax.third_party to jax._src.third_party.

2020-10-17 11:43:07 -07:00

__init__.py

Move most of the implementation of jax.scipy into jax._src.scipy.

2020-10-16 17:04:25 -04:00

api.py

Replace host_id with process_index terminology, take 2.

2021-04-20 18:13:34 -07:00

cloud_tpu_init.py

Add v3-64 config to automatic Cloud TPU pod slice initialization.

2021-04-19 15:49:40 -07:00

config.py

Add optional distributed debugging logging.

2021-04-20 13:34:45 -07:00

custom_derivatives.py

Move vtile to batching.py, make it possible to add new BatchTraces

2021-04-20 08:33:01 -07:00

dlpack.py

Remove the JAX lazy sublanguage.

2021-03-09 21:40:46 -05:00

dtypes.py

Move jax.config to jax._src.config.

2021-04-19 08:53:12 -07:00

errors.py

remove disable_omnistaging mechanism

2021-03-29 15:26:57 -07:00

lax_reference.py

[JAX] Make lax_reference private by moving it to jax._src.lax_reference.

2021-04-08 09:06:14 -07:00

pprint_util.py

[JAX] Move pprint_util into jax._src.

2020-11-19 06:42:19 -08:00

profiler.py

Remove backward compatibility code paths for jaxlib < 0.1.65.

2021-04-09 15:39:38 -04:00

random.py

Merge pull request #6461 from apaszke:xmap-awn

2021-04-19 06:36:34 -07:00

source_info_util.py

Revert part of #5319: don't use traceback_util.path_starts in source_info_util.user_frame.

2021-01-19 15:01:52 -08:00

traceback_util.py

tiny change for source sync

2021-03-05 16:32:28 -08:00

tree_util.py

Move jax.tree_util implementation to jax._src.tree_util.

2021-03-24 12:00:38 -07:00

util.py

Add optional distributed debugging logging.

2021-04-20 13:34:45 -07:00