rocm_jax/_src at 1526c3e20cd8852ed16f3e8640a47f09b0de5c70 - rocm_jax - Gitea For EOELAB

mirrors/rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00

History

Yash Katariya 1526c3e20c Improve the error message which is raised from _get_and_check_device_assignment.

Before:

```
ValueError: Devices of all `Array` inputs and outputs should be the same. Got array device ids [0] on platform CPU and another array's device ids [0, 1, 2, 3] on platform CPU
```

After:

```
ValueError: Received incompatible devices for jitted computation. Got argument inp of ArrayPjitTest.test_jit_with_sharding_constraint_committed_inp_error.<locals>.sharded_inp with bfloat16[8,2] and device ids [0] on platform CPU and with_sharding_constraint or nested pjit or shard_map with device ids [0, 1, 2, 3] on platform CPU at jax/tests/pjit_test.py:2509 (sharded_inp)
```
PiperOrigin-RevId: 508746961

2023-02-10 13:54:15 -08:00

..

Don't set TPU topology env vars in cloud_tpu_init.py

2023-01-13 22:02:16 +00:00

Remove typing_extensions dependency

2022-12-05 15:42:26 -08:00

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

Improve the error message which is raised from _get_and_check_device_assignment.

2023-02-10 13:54:15 -08:00

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

[PJRT C API] Register a backend factory for every PJRT plugin set in PJRT_NAMES_AND_LIBRARY_PATHS.

2023-02-09 14:33:46 -08:00

Added shape error checking for compute_fans

2023-01-18 20:59:11 -08:00

Merge pull request #14364 from jakevdp:fix-tril-indices

2023-02-10 12:25:06 -08:00

Remove typing_extensions dependency

2022-12-05 15:42:26 -08:00

Added scipy.stats.rankdata

2023-02-07 12:07:00 -05:00

simpler pretty-print for pjit, tweak custom pp rule signature

2023-02-09 12:45:51 -08:00

Address drastic slowdown in mypy runtime

2022-12-16 14:48:26 -08:00

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

abstract_arrays.py

Expose fp8 in jax dtypes and mlir builder.

2023-01-13 18:12:12 -08:00

ad_checkpoint.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

ad_util.py

Merge pull request #13693 from jakevdp:typing-ad-util

2022-12-20 16:50:11 -08:00

api_util.py

Make error_test a jax_test so that we can test other configs and fix it with jit/pjit merge.

2023-01-17 18:43:05 -08:00

api.py

Improve the error message which is raised from _get_and_check_device_assignment.

2023-02-10 13:54:15 -08:00

array.py

Add infrastructure for managing deprecations.

2023-02-09 05:48:40 -08:00

basearray.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

basearray.pyi

[typing] define additional methods & properties on jax.Array

2023-02-10 09:42:32 -08:00

callback.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

checkify.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

cloud_tpu_init.py

Add environment variable check in additional to libtpu check for cloud tpu vm

2023-01-25 09:51:37 -08:00

config.py

Remove jax_experimental_subjaxpr_lowering_cache since it was only for jit and was False by default. Now that jit/pjit are merged, this cache is not needed since pjit does the caching and we get it for free.

2023-02-08 14:55:56 -08:00

core.py

[typing] define additional methods & properties on jax.Array

2023-02-10 09:42:32 -08:00

custom_api_util.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

custom_batching.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

custom_derivatives.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

custom_transpose.py

migrate internal dependencies from jax.interpreters.ad to jax._src.interpreters.ad

2023-02-06 22:52:36 -08:00

debugging.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

deprecations.py

Deprecate jax.interpreters.xla.Device and jax.interpreters.xla.DeviceArray.

2023-02-09 16:11:48 -08:00

device_array.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

dispatch.py

Improve the error message which is raised from _get_and_check_device_assignment.

2023-02-10 13:54:15 -08:00

distributed.py

Add Open MPI automatic distributed initialization

2023-01-11 17:08:09 -06:00

dlpack.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

dtypes.py

Expose fp8 in jax dtypes and mlir builder.

2023-01-13 18:12:12 -08:00

environment_info.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

errors.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

flatten_util.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

iree.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

lax_reference.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

lazy_loader.py

Avoid loading scipy eagerly.

2022-10-12 19:51:09 +00:00

linear_util.py

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -08:00

monitoring.py

Move contents of jax.monitoring to jax._src.monitoring.

2023-02-08 15:03:22 -08:00

path.py

Make the etils[epath] dependency of JAX optional.

2022-10-12 21:16:48 +00:00

pjit.py

Improve the error message which is raised from _get_and_check_device_assignment.

2023-02-10 13:54:15 -08:00

pretty_printer.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

prng.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

profiler.py

Add a number of missing function cross-references in the docs.

2022-11-07 12:00:26 -05:00

public_test_util.py

Roll-back https://github.com/google/jax/pull/14144 due to downstream test failures

2023-01-25 12:15:36 -08:00

random.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

sharding.py

Prune accidental exports from jax.interpreters.pxla.

2023-02-07 11:16:42 -08:00

source_info_util.py

Let is_user_frame ignore frames from stdlib.

2022-11-08 10:50:08 -08:00

stages.py

Move jax.interpreters.xla to jax._src.interpreters.xla.

2023-02-07 15:01:32 -08:00

test_util.py

Move jax.interpreters.mlir to jax._src.interpreters.mlir.

2023-02-08 14:39:01 -08:00

traceback_util.py

add blank line, mainly to trigger/test source sync

2023-02-01 13:56:29 -08:00

tree_util.py

Use jnp.arange to break the pjit cache (when jit and pjit are merged) because pytest runs tests non-hermetically.

2023-02-08 10:17:37 -08:00

typing.py

Remove typing_extensions dependency

2022-12-05 15:42:26 -08:00

util.py

shard_map (shmap) prototype and JEP

2023-02-02 23:01:30 -08:00