add a way to print or save summaries within an @jit function

[mattjj]

c.f. #196 and this comment in particular, maybe we could add something like this:

y, python_io = jax.py_call(some_py_fun, x, python_io)

[jonasrauber]

Sounds great. If I understand it correctly, wouldn’t a nojit decorator also make sense? E.g. something like nojit = lambda f: partial(py_call, f) (ignoring python_io for the moment; can you explain its purpose?)

I guess the called function‘s output needs to have fixed dtype and shape and they will be auto-detected during jit and an error will be thrown if they change?

[mattjj]

Interesting idea for a nojit decorator! We'll definitely have to keep that idea in mind, and see if we can implement it.

The purpose of python_io there is to model the state of the world (and in particular the stateful Python interpreter) and thread it through our computation, or at least through the parts that have side effects. That lets us model side-effects in a functionally pure way, since we can give semantics to these programs in terms of producing a new state of the world rather than modifying the old one (and the fact that we don't actually copy the whole world every time, and instead rely on side-effects, is an optimization in the implementation). The use of python_io would have some special limitations though, like linearity in the sense of linear logic (e.g. one python_io value can't be passed into two different py_call calls, and instead must be linearly threaded through the computation) because we really can't copy the world.

Another way to think about it would be like we're passing around the GIL (without acquiring it), and that py_call needs to acquire the GIL when it's called and release it when it returns.

That's basically how functional languages like Haskell model effectful operations like IO (also worth looking at the writer monad and the state monad), though there's a lot of variation and detail.

[mattjj]

cc @gnecula thinking about recent things recently. (Just noticed this because of issue 2796.)

[gnecula]

We now have a new experimental feature (experimental/host_callback.py:id_tap) that should work to save summaries from compiled code.

[gnecula]

I believe that #3006 addresses this.

[jonasrauber]

Nice work, thanks!