Alexey Radul f97db31ead Fix type errors caught by pytype.

2023-07-12 15:11:14 -04:00

..

control_flow

Add discharge rules for scan/while

2023-07-06 22:30:35 +00:00

__init__.py

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

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

ann.py

[shape_poly] Shape polymorphism support for approx_top_k

2023-06-26 22:02:41 -07:00

convolution.py

[shape_poly] Cleanups for the shape polymorphism APIs.

2023-06-30 15:56:57 +03:00

eigh.py

Improve the bisection strategy in eigh.

2023-05-04 18:01:23 -07:00

fft.py

Add support for dynamic shapes to jax.fft.

2023-06-17 04:50:54 -07:00

lax.py

Fix type errors caught by pytype.

2023-07-12 15:11:14 -04:00

linalg.py

schur: fix broken jvp rule

2023-06-30 02:30:25 -07:00

other.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

parallel.py

always lower all_to_all to AllToAll

2023-04-11 18:31:17 -07:00

qdwh.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

slicing.py

Test and implement slicing not dropping raggedness information.

2023-07-11 15:21:55 -04:00

special.py

Split dtype argument from other arguments in special functions.

2023-03-25 11:41:14 -07:00

stack.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

svd.py

Fix corner cases in JAX SVD: a) Clamp negative singular values to zero. b) Return all NaN for matrices with non-finite values.

2023-06-13 11:06:49 -07:00

utils.py

KeyArray: improve errors for unimplemented primitives

2023-04-24 16:54:25 -07:00

windowed_reductions.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00