Как читать этот пост? Если вы уже знаете что такое jax и haiku/flax/optax - можно сразу прыгать до секции "почему плохо для продакшена". Если интересно про опыт использования библиотек, то в секцию "про опыт использования". Если хочется TL;DR, то jax быстрый, на нём приятно писать, у него круто хэндлится рандом, но в нём легко набаговать, если вы не сталкивались с jit-компиляцией и функциональными языками. Если дать jax джуну, то он сможет найти миллион неожиданных способов выстрелить себе в ногу. Документация…