Marc van Zee 9d18f43a01 Do not normalize FFT by a constant "1" if no normalization is provided (i.e., norm is None). ...

Without this, the compiled graph will still contain a node multipying a complex number with a constant 1+0j (1 is cast to complex because the other term is complex as well). This is problematic when converting to TFLite using jax2tf, because multiplying complex numbers is not supported in TFLite. With this change, the multiplication is removed from the graph all together.

PiperOrigin-RevId: 459566727

2022-07-07 11:54:39 -07:00

..

__init__.py

Move jax.numpy internals into jax._src.numpy.

2020-10-16 20:35:19 -04:00

fft.py

Do not normalize FFT by a constant "1" if no normalization is provided (i.e., norm is None).

2022-07-07 11:54:39 -07:00

index_tricks.py

[x64] make jnp.mgrid compatible with strict dtype promotion

2022-06-01 10:21:49 -07:00

lax_numpy.py

[dynamic-shapes] revive basic bounded int machinery, add tests

2022-07-06 22:31:26 -07:00

linalg.py

[x64] make linalg functions & tests compatible with strict dtype promotion

2022-06-16 10:32:20 -07:00

ndarray.py

Make jnp.array return jax.Array. Add input and result handlers for jax.Array. Also added tests for add under jit.

2022-06-24 10:05:06 -07:00

polynomial.py

jnp.roots: better support for computation under JIT

2022-06-23 14:48:53 -07:00

reductions.py

jnp.average: support keepdims argument

2022-06-28 10:55:55 -07:00

setops.py

jnp.unique: improve error when run under JIT

2022-06-08 15:57:41 -07:00

ufuncs.py

[x64] make jnp.modf() compatible with strict dtype promotion

2022-06-03 15:23:06 -07:00

util.py

[dynamic shapes] revive iree

2022-07-06 15:01:16 -07:00

vectorize.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00