4 ms·
That first call is from tf2xla which JAX doesn't use. But I think those calls just set up the global buffers used by the launched program inputs/temporaries.
by alevskaya 5y ago
That first call is from tf2xla which JAX doesn't use. But I think those calls just set up the global buffers used by the launched program inputs/temporaries. The entire point of XLA's static shape discipline is allowing for a highly optimized memory schedule inside compiled programs.
- tubby12345 5y agoI believe you but I'm wondering what that memory scheduling actually looks like (e.g. solving the full MIP?). do you have a code pointer to that scheduling routine?