you can't store a function as a tensor. the tensors are the inputs/outputs that flow along the edges of the graph. TF stores things directly as a graph: https:/

you can't store a function as a tensor. the tensors are the inputs/outputs that flow along the edges of the graph. TF stores things directly as a graph: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/node_def.proto https://github.com/tensorflow/tensorflow/blob/master/tensorf... while Jax represents computations in an intermediate representation: https://jax.readthedocs.io/en/latest/jaxpr.html https://jax.readthedocs.io/en/latest/jaxpr.html