George Necula 6d53aaf7d0 [pallas] Improve the error localization ...

* Add the source location information for the index map function to
    `BlockMapping`.
  * Removed the `compute_index` wrapper around the index_map, so that
    we can get the location information for the index_map, not the wrapper.
  * Added source location to the errors related to index map functions.
  * Added an error if the index map returns something other than integer
    scalars.
  * Construct BlockSpec origins for arguments using JAX helper functions
    to get argument names
  * Removed redundant API error tests from tpu_pallas_test.py

2024-07-30 14:11:57 +02:00

..

clusters

Merge pull request #20174 from coreyjadams:main

2024-07-08 12:19:18 -07:00

cudnn

add sm86/sm89

2024-07-24 21:12:37 +00:00

debugger

Merge pull request #20858 from rajasekharporeddy:doc_typos

2024-04-25 06:58:27 -07:00

export

Bumped jaxlib version mypy uses on the CI

2024-07-26 11:22:39 +01:00

extend

Export jaxlib.xla_client.register_custom_call_target as jax.extend.ffi.register_ffi_target.

2024-07-19 08:12:25 -04:00

image

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

internal_test_util

Activate Cholesky Factorization Kernel to XLA's FFI

2024-07-25 09:59:28 -07:00

interpreters

[pallas] Improve the error localization

2024-07-30 14:11:57 +02:00

lax

Merge pull request #22080 from vfdev-5:add-device-kwarg-linspace-array

2024-07-26 11:18:24 -07:00

lib

#sdy Initial set of changes to allow for lowering to the Shardy dialect.

2024-07-23 05:32:06 -07:00

nn

remove named_shapes (since xmap is now gone)

2024-07-25 00:54:50 +00:00

numpy

[array API] add device argument to fftfreq & rfftfreq

2024-07-29 13:23:54 -07:00

ops

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

pallas

[pallas] Improve the error localization

2024-07-30 14:11:57 +02:00

scipy

Add canonical arg to Rotation.as_quat() and switch .inv() to use the quaternion conjugate.

2024-07-24 15:29:19 -07:00

state

remove named_shapes (since xmap is now gone)

2024-07-25 00:54:50 +00:00

third_party

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

__init__.py

…

abstract_arrays.py

Added basic support for int2/uint2 dtypes to JAX

2024-07-04 04:13:24 -07:00

ad_checkpoint.py

Merge pull request #22263 from hawkinsp:tuples

2024-07-17 09:56:18 -07:00

ad_util.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

api_util.py

remove named_shapes (since xmap is now gone)

2024-07-25 00:54:50 +00:00

api.py

remove inlined jax.nn.initializers definitions, resolving TODO of levskaya et al

2024-07-24 20:55:36 -07:00

array.py

Bumped jaxlib version mypy uses on the CI

2024-07-26 11:22:39 +01:00

basearray.py

Document copy_to_host_async method of jax.Array

2024-07-26 08:21:01 -07:00

basearray.pyi

[array API] add device argument to fftfreq & rfftfreq

2024-07-29 13:23:54 -07:00

blocked_sampler.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

cache_key.py

Bumped mypy to 1.10.0 and ruff to 0.4.4

2024-05-16 23:16:32 +01:00

callback.py

Merge pull request #22211 from hawkinsp:singletons

2024-07-03 11:07:00 -07:00

checkify.py

Remove linear parameter from lax.cond_p.

2024-07-01 10:25:42 -04:00

cloud_tpu_init.py

Set the value of config.jax_platforms directly instead of setting the env variable JAX_PLATFORMS. Setting JAX_PLATFORMS doesn’t do anything because config.jax_platforms is already initialized at that point.

2024-07-22 15:09:15 -07:00

compilation_cache_interface.py

Reverts f18739900c615a85e8d182bcf3217f704cf7aa0d

2024-04-16 20:28:24 -07:00

compilation_cache.py

Merge pull request #22556 from cool-RR:log-cache-key

2024-07-26 05:32:11 -07:00

compiler.py

[xla] Replace debug option xla_use_shardy with execution option

2024-07-29 16:11:36 -07:00

compute_on.py

Move throwing the error for invalid compute_on into the try-finally context.

2024-05-28 15:20:26 -07:00

config.py

#sdy Initial set of changes to allow for lowering to the Shardy dialect.

2024-07-23 05:32:06 -07:00

core.py

[pallas] More simplification of grid mapping and calling convention

2024-07-29 15:53:47 +02:00

custom_api_util.py

…

custom_batching.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

custom_derivatives.py

remove named_shapes (since xmap is now gone)

2024-07-25 00:54:50 +00:00

custom_partitioning.py

Bumped jaxlib version mypy uses on the CI

2024-07-26 11:22:39 +01:00

custom_transpose.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

debugging.py

Added a trivial discharge rule for debug_callback_p

2024-07-29 22:26:01 +01:00

deprecations.py

Remove the unaccelerate_deprecation utility

2024-07-23 05:07:49 -07:00

dispatch.py

Shmallas, a.k.a. allow lowering shard_map + run_state to a pallas_call.

2024-07-23 15:16:50 -07:00

distributed.py

Squashed commit of the following:

2024-07-02 13:18:05 -05:00

dlpack.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

dtypes.py

extend type annotation for lax.convert_element_type

2024-07-11 05:27:11 -07:00

earray.py

[array API] add device property & to_device method

2024-07-23 11:12:35 -07:00

effects.py

Effects is now an immutable set

2023-12-11 08:45:52 -08:00

environment_info.py

jax.print_environment_info: add platform.uname()

2024-02-13 11:23:56 -08:00

errors.py

[key reuse] rename flag to jax_debug_key_reuse

2024-03-22 05:37:30 -07:00

flatten_util.py

Fix duplicate word occurrences

2023-12-19 06:15:30 +01:00

hardware_utils.py

Skip cuda backend initialization if no nvidia GPUs are visible.

2024-07-29 15:12:58 +00:00

jaxpr_util.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

lax_reference.py

Implements Ragged Dot API

2024-05-11 06:40:18 -07:00

layout.py

Bump minimum jaxlib version to 0.4.31. The corresponding xla_extension_version is 279 and mlir_api_version is 57

2024-07-29 18:44:31 -07:00

lazy_loader.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

linear_util.py

`debug_info no longer requires non-None func_src_info`

2024-07-05 20:08:53 +01:00

logging_config.py

Add jax_debug_log_modules config option.

2023-07-28 18:11:12 +00:00

lru_cache.py

Bumped jaxlib version mypy uses on the CI

2024-07-26 11:22:39 +01:00

mesh.py

Delete mesh.Loop now that xmap has been deleted

2024-07-25 14:08:32 -07:00

monitoring.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

op_shardings.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

partition_spec.py

Make UNCONSTRAINED's __str__ and __repr__ so it prints nicely in all cases

2024-05-30 17:41:32 -07:00

path.py

Clean up

2024-07-04 17:16:32 +04:00

pickle_util.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

pjit.py

Bump minimum jaxlib version to 0.4.31. The corresponding xla_extension_version is 279 and mlir_api_version is 57

2024-07-29 18:44:31 -07:00

pretty_printer.py

Changed the naming of internal config APIs

2024-06-18 11:48:57 +01:00

prng.py

remove named_shapes (since xmap is now gone)

2024-07-25 00:54:50 +00:00

profiler.py

Support pathlib.Path input to start_trace

2024-07-19 23:41:02 +03:00

public_test_util.py

Added basic support for int2/uint2 dtypes to JAX

2024-07-04 04:13:24 -07:00

random.py

Finalize deprecation of jax.random.shuffle

2024-07-27 11:21:49 -07:00

shard_alike.py

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -07:00

sharding_impls.py

Bump minimum jaxlib version to 0.4.31. The corresponding xla_extension_version is 279 and mlir_api_version is 57

2024-07-29 18:44:31 -07:00

sharding_specs.py

Remove mesh_shape from ShardingSpec since it's not in use anymore.

2024-04-07 21:13:41 -07:00

sharding.py

#sdy Add dataclass that stores the data needed to build a Shardy MLIR sharding attribute.

2024-07-23 12:22:05 -07:00

source_info_util.py

Use a class with __slots__ instead of a NamedTuple in JaxprEqn and SourceInfo, which are two tuples we build frequently.

2024-07-01 19:18:58 -07:00

sourcemap.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

stages.py

require keyword arguments to Traced.lower

2024-07-10 17:59:26 -07:00

test_util.py

Remove the unaccelerate_deprecation utility

2024-07-23 05:07:49 -07:00

tpu_custom_call.py

Delete xmap and the jax.experimental.maps module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period).

2024-07-24 10:24:09 -07:00

traceback_util.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

tree_util.py

[pallas] Improve some error messages and add API tests.

2024-07-04 09:02:16 +02:00

tree.py

Fix annotation for jax.tree.transpose

2024-07-09 08:06:16 -07:00

typing.py

[jax] Make the SupportsDType protocol runtime checkable.

2024-07-16 10:57:15 -07:00

util.py

Prefer __qualname__ as a pjit_p name.

2024-07-02 14:53:50 -04:00

xla_bridge.py

Merge pull request #22703 from Rifur13:plugin-fix

2024-07-29 12:10:42 -07:00