AutoComposite vs ExtensionTypes

[Jonas Eschle]

Dear all,

TFP features since a while (experimentally) the possibility to use automatically composite tensors, e.g. by sticking to a contract, a complex object will be able to pass through a tf.function.

TensorFlow also introduced something similar now, the ExtensionType.

Since we would like to use this features ourselves in a library and were wondering:
what is the future of these? Is the auto composite function of TFP meant to stay? Or should we rather rely on the ExtensionTypes?
What is about JAX, will it adopt one or the other as well? (in case you know)

Many thanks in advance,
Best,
Jonas

[Brian Patton 🚀]

For your own Distribution subclass, it seems like you could subclass both Distribution and ExtensionType. Insofar as that is a public TF API, it's probably going to have better long term support than the AutoCompositeTensorDistribution TFP provides.

I don't see a lot of novelty to ExtensionType relative to CompositeTensor [apart from "doesn't use internal TF APIs"], so I guess it would take a strong push to make us switch over the TFP-built-in distributions.

JAX as an ML frontend is fairly independent of TF, especially this far into the Python layer. I'd expect JAX pytree functionality to remain the preferred way of doing this for JAX, but I'm not an authority on it.

Brian Patton | Software Engineer | b...@google.com

--
You received this message because you are subscribed to the Google Groups "TensorFlow Probability" group.
To unsubscribe from this group and stop receiving emails from it, send an email to tfprobabilit...@tensorflow.org.
To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/tfprobability/17b6cc02-4572-4f64-aaf6-3c3c33a8f422n%40tensorflow.org.