[jax2tf] Added a flag and environment variable to control the serialization version

[gnecula]

This allows us to control the serialization version to be compatible with
the deployed version of tf.XlaCallModule. In particular, we can run
most tests with the maximum available version, while keeping the
default lower.

[gnecula]

@junwhanahn @rxsang PTAL

[junwhanahn]

The changes look good to me.

Somewhat related question: do you plan to maintain min/max supported version in jax2tf (if we ever need one when we clean up some old jax2tf logic) separately from TF2XLA, or in sync? I would guess that the former makes more sense, but not sure if it will make the testing more complex.

[gnecula]

The changes look good to me.

Somewhat related question: do you plan to maintain min/max supported version in jax2tf (if we ever need one when we clean up some old jax2tf logic) separately from TF2XLA, or in sync? I would guess that the former makes more sense, but not sure if it will make the testing more complex.

I don't quite see how to keep the versioning perfectly in sync. My plan is to support in jax2tf the versions that have been active in google3 over the last month. This means 1 or 2 versions, the most recent. The tests in google3 will verify that we can use the most recent version and the one that is marked as default. There is no test per se that we can use the version that has been active a month ago; we rely on people not updating the default serialization version too early.

Also in OSS people will have different versions of JAX and TF. I do assume that people will be willing to use tf-nightly if they get an error; this has been true for jax2tf so far because it relies always on some very recent details for the TF ops.

[junwhanahn]

The changes look good to me.

Somewhat related question: do you plan to maintain min/max supported version in jax2tf (if we ever need one when we clean up some old jax2tf logic) separately from TF2XLA, or in sync? I would guess that the former makes more sense, but not sure if it will make the testing more complex.

I don't quite see how to keep the versioning perfectly in sync. My plan is to support in jax2tf the versions that have been active in google3 over the last month. This means 1 or 2 versions, the most recent. The tests in google3 will verify that we can use the most recent version and the one that is marked as default. There is no test per se that we can use the version that has been active a month ago; we rely on people not updating the default serialization version too early.

Also in OSS people will have different versions of JAX and TF. I do assume that people will be willing to use tf-nightly if they get an error; this has been true for jax2tf so far because it relies always on some very recent details for the TF ops.

Thanks. Sounds good to me.

[gnecula]

This was already merged, but copybara did not close it.
See 603eeb1