The pickles drive me crazy. Tensorflow saved models are a great way to solve the problem... Save the computation graph and weights, and drop all the crusty cod

The pickles drive me crazy. Tensorflow saved models are a great way to solve the problem... Save the computation graph and weights, and drop all the crusty code dependencies. I think ONNX models are similar. I expect there should be a Jax equivalent at some point, as Jax is basically perfectly designed for this (everything is expressed in lax operations, which allows changing implementations for cpu/gpu/tpu freely... So just save the list of lax ops).