Matthew Johnson 4236eb2b59

omnistaging, under a flag and disabled by default (#3370) ...

This change, when enabled, stages out all primitive calls in the dynamic
scope of a jitted, pmapped, or control flow function, rather than only
staging out based on data dependence. One improvement is that jitted
functions can consume less memory, by avoiding instantiating large
constants at trace time, and cause less memory fragmentation as well. It
also simplifies several internals.

See https://github.com/google/jax/pull/3370 fo more information.

2020-07-30 12:59:36 -07:00

..

__init__.py

omnistaging, under a flag and disabled by default (#3370)

2020-07-30 12:59:36 -07:00

lax_control_flow.py

omnistaging, under a flag and disabled by default (#3370)

2020-07-30 12:59:36 -07:00

lax_fft.py

Cleanup: convert uses of import numpy as onp in library code (#3754)

2020-07-14 13:05:31 -07:00

lax_parallel.py

omnistaging, under a flag and disabled by default (#3370)

2020-07-30 12:59:36 -07:00

lax.py

omnistaging, under a flag and disabled by default (#3370)

2020-07-30 12:59:36 -07:00