rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-17 04:16:07 +00:00

History

Yash Katariya cc5af7ed98 Rename ReshapeableDevicesSharding to PositionalSharding and add an alias NamedSharding for MeshPspecSharding.

`MeshPspecSharding` name will be replaced with `NamedSharding` in 3 months.

PiperOrigin-RevId: 485753078

2022-11-02 19:13:13 -07:00

clusters

Migrate JAX internals to builtin Python logging

2022-10-13 21:32:44 +02:00

debugger

Cleanup: remove unused imports in private modules

2022-10-20 14:37:21 -07:00

image

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

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

lax

Enable state effect in cond_p (except in grad and vmap)

2022-11-02 16:07:58 -07:00

lib

Increase the minimum jaxlib version to 0.3.22.

2022-10-27 10:24:11 -04:00

Avoid loading scipy eagerly.

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

numpy

Merge pull request #13035 from jakevdp:jnp-put

2022-10-31 09:45:17 -07:00

ops

[typing] annotate jax.numpy reduction operations

2022-10-26 13:33:15 -07:00

scipy

[typing] annotate jax.numpy reduction operations

2022-10-26 13:33:15 -07:00

state

Enable state effect in cond_p (except in grad and vmap)

2022-11-02 16:07:58 -07:00

third_party

Call _check_arraylike for jnp.linalg & jnp.fft functions

2022-10-31 09:19:53 -07: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

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

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

ad_checkpoint.py

Annotate tree_util

2022-10-26 13:38:38 -04:00

ad_util.py

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

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

api_util.py

Annotate tree_util

2022-10-26 13:38:38 -04:00

api.py

Make device_put accept a prefix tree with Sharding leaves as the second argument

2022-11-01 14:32:55 -07:00

array.py

Add global_shards to jax.Array as it exists on GDA and is being used in various places.

2022-10-31 09:08:03 -07:00

basearray.py

Add initial jax.Array base class for instance checks & annotation

2022-09-26 07:48:43 -07:00

basearray.pyi

[typing] overloads for jnp.linalg.svd & jnp.linalg.qr

2022-10-05 14:29:00 -07:00

callback.py

Increase the minimum jaxlib version to 0.3.22.

2022-10-27 10:24:11 -04:00

checkify.py

Add lax.cumlogsumexp for cumulative logsumexp operations.

2022-10-31 15:08:52 -07:00

cloud_tpu_init.py

Set JAX_PLATFORMS=tpu,cpu on TPUs

2022-10-05 04:39:04 +00:00

config.py

Merge pull request #13039 from skye:cache_compile_time_heuristic

2022-11-02 11:13:52 -07: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

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

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

custom_derivatives.py

Cleanup: remove unused imports in private modules

2022-10-20 14:37:21 -07:00

custom_transpose.py

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

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

debugging.py

Merge pull request #13065 from sharadmv:vis-colors

2022-11-01 18:11:06 -07:00

device_array.py

Add addressable_shards to SDA and DA as a compatibility API to match with jax.Array. This will aid in transition to jax.Array.

2022-10-05 12:35:17 -07:00

dispatch.py

Merge pull request #13039 from skye:cache_compile_time_heuristic

2022-11-02 11:13:52 -07:00

distributed.py

Migrate JAX internals to builtin Python logging

2022-10-13 21:32:44 +02:00

dlpack.py

canonicalize dtypes when loading arrays via dlpack

2022-10-20 15:20:34 -07:00

dtypes.py

[typing] annotate jax._src.dtypes

2022-10-17 13:49:26 -07: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

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

2022-09-22 12:27:19 -07: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

path.py

Make the etils[epath] dependency of JAX optional.

2022-10-12 21:16:48 +00: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

integrate new (partitionable, count-space-exhaustive) counts generation

2022-10-29 00:05:49 -07:00

profiler.py

Migrate JAX internals to builtin Python logging

2022-10-13 21:32:44 +02:00

public_test_util.py

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

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

random.py

Re-land https://github.com/google/jax/pull/12588 with minor fixes

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

sharding.py

Rename ReshapeableDevicesSharding to PositionalSharding and add an alias NamedSharding for MeshPspecSharding.

2022-11-02 19:13:13 -07:00

source_info_util.py

Include column information in Python locations under Python 3.11.

2022-10-14 19:14:35 +00:00

stages.py

Roll forward with fix: Remove the original python function fun_ from C++ PjitFunction, as the destroying fun_ may yield the thread in some cases, which causes error during deleting the python object of PjitFunction.

2022-10-18 10:05:53 -07:00

test_util.py

Remove error-prone most_recent_entry() support from lu.cache.

2022-10-27 16:41:44 -07:00

traceback_util.py

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

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

tree_util.py

Annotate tree_util

2022-10-26 13:38:38 -04:00

typing.py

add allow_opaque_dtype to dtypes.canonicalize_dtype utility

2022-10-17 13:47:42 -07:00

util.py

Increase the minimum jaxlib version to 0.3.22.

2022-10-27 10:24:11 -04:00