Change TFP code to point to tfxla.variadic_reduce_v2.

PiperOrigin-RevId: 385176212

Author: gnecula

tensorflow_probability/python/internal/variadic_reduce.py

@@ -123,11 +123,20 @@ def make_variadic_reduce(reducer, vjp_bwd, tangents_fn):
   def _xla_reduce(operands, inits, axis):
     """JIT-ed wrapper for TF `xla.variadic_reduce(..., reducer)`."""
     from tensorflow.compiler.tf2xla.python import xla  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
-    result = xla.variadic_reduce(
-        operands,
-        init_value=inits,
-        dimensions_to_reduce=axis,
-        reducer=tf.function(reducer).get_concrete_function(inits, inits))
+    try:
+      result = xla.variadic_reduce(
+          operands,
+          init_values=inits,
+          dimensions_to_reduce=axis,
+          reducer=tf.function(reducer).get_concrete_function(inits, inits))
+    except TypeError:
+      # TODO(necula): This is needed only temporarily, until cl/384088502
+      # makes it in tf-nightly.
+      result = xla.variadic_reduce_v2(
+          operands,
+          init_values=inits,
+          dimensions_to_reduce=axis,
+          reducer=tf.function(reducer).get_concrete_function(inits, inits))
     # Graph mode: variadic reduce doesn't specify output shapes. Patch that.
     shp = operands[0].shape
     for arg in operands: