XLA and MultiWorkerMirroredStrategy

The release notes of tensorflow 2.5.0 mentions that: “XLA can now compile MirroredStrategy: the step function passed tostrategy.run can now be annoted with jit_compile=True.”

Is the same support also available for MultiWorkerMirroredStrategy? I tried adding jit_compile=True to the step function in this colab tutorial:

And I got: “RuntimeError: merge_call called while defining a new graph or a tf.function …”

Thanks!

1 Like

Let’s see if they approve:

https://github.com/tensorflow/tensorflow/pull/55963

1 Like