Anish Tondwalkar bf416a8b5c geqrf_p and householder_product_p directly call custom_calls ...

This replaces the xla_fallback path, which just used the Client HLO API to
generate custom_calls.

PiperOrigin-RevId: 518060025

2023-03-20 13:29:29 -07:00

..

control_flow

[jax2tf] Minor improvement in an error message

2023-03-18 11:00:46 +02: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

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

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

convolution.py

Implement lowering for convolutions with dynamic padding

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

eigh.py

internal: avoid unused imports in lax_numpy

2023-03-08 10:29:04 -08:00

fft.py

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07:00

lax.py

[LAX:RBG] Allow any type to RngBitGenerator. BF16 values are heavily quantized for long distributions which leads to failing the distribution test but in reality the distributions match.

2023-03-17 22:39:43 -07:00

linalg.py

geqrf_p and householder_product_p directly call custom_calls

2023-03-20 13:29:29 -07:00

other.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

parallel.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

qdwh.py

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

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

slicing.py

Remove references to jax.config.jax_array, which is always True at head.

2023-03-15 17:09:11 -07:00

stack.py

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

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

svd.py

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

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

utils.py

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

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

windowed_reductions.py

[shape_poly] Add support for jnp.cum{sum,prod,max,min} with shape polymorphism

2023-03-03 10:22:45 +01:00