George Necula f960c287c4 [pallas] Improve error messages for input_output_aliasing ...

Currently, we get MLIR verification errors when the inputs
and outputs declared to be aliased do not have matching
shapes and dtypes. We add a nicer error message that localizes
the inputs and outputs in the corresponding PyTrees.

Interestingly, if the output index is out of bounds, there
is no MLIR verification error. This seems to be a bug
in the StableHLO verification code.

Currently, in interpreter mode we get a mix of internal
assertion errors when there are errors in input_output_aliasing.

2024-07-08 15:18:27 +03:00

..

_src

[pallas] Improve error messages for input_output_aliasing

2024-07-08 15:18:27 +03:00

example_libraries

Run pyupgrade --py310-plus.

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

experimental

[pallas] Move the import of tpu_custom_call from tests to a more central place.

2024-07-08 04:27:41 -07:00

extend

Add ffi_call function with a similar signature to pure_callback.

2024-07-01 09:40:31 -04:00

image

…

interpreters

Merge pull request #22211 from hawkinsp:singletons

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

lax

Finalize deprecations of several APIs

2024-05-14 10:40:40 -07:00

lib

…

nn

Finalize deprecations of several APIs

2024-05-14 10:40:40 -07:00

numpy

jnp.arange: fix incorrect type annotation

2024-07-07 22:25:39 -07:00

ops

…

scipy

Copy nn.{softmax,log_softmax} to scipy.special

2024-06-22 09:32:14 -04:00

tools

Added basic support for int2/uint2 dtypes to JAX

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

__init__.py

add more info to xla_computation deprecation warning

2024-07-02 13:31:07 +00:00

ad_checkpoint.py

…

api_util.py

…

BUILD

Clean up

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

cloud_tpu_init.py

…

collect_profile.py

…

core.py

Finalize deprecation of several previously-deprecated jax.core functions:

2024-06-28 07:28:28 -07:00

custom_batching.py

…

custom_derivatives.py

…

custom_transpose.py

…

debug.py

…

distributed.py

…

dlpack.py

…

dtypes.py

…

errors.py

…

export.py

[export] Rename some API entry points

2024-06-13 06:44:13 +02:00

flatten_util.py

…

monitoring.py

…

profiler.py

…

py.typed

…

random.py

…

sharding.py

Deprecate XLACompatibleSharding in favor of jax.sharding.Sharding.

2024-06-05 09:07:27 -07:00

stages.py

rename Specialized to Traced (and specialize to trace)

2024-06-06 17:43:08 -07:00

test_util.py

…

tree_util.py

Add builtin cc dataclass pytree node for performance.

2024-04-23 14:14:49 -07:00

tree.py

…

typing.py

…

util.py

Add discharge rules for scan with mutable arrays. Move mutable array tests to separate file.

2024-05-02 14:36:16 -04:00

version.py

Bump minimum jaxlib version to v0.4.30.

2024-06-18 12:35:08 -04:00