Expose SetUpAlias to XLA Python client.

PiperOrigin-RevId: 306796534
Change-Id: I141c4e7ef52e96d26ad9aa0ec8d3094fb6f0ba42

tensorflow/compiler/xla/python/xla.cc

@@ -1318,7 +1318,15 @@ PYBIND11_MODULE(xla_extension, m) {
       .def("IsConstant", &XlaBuilder::IsConstant)
       .def("SetOpMetadata", &XlaBuilder::SetOpMetadata)
       .def("SetSharding", &XlaBuilder::SetSharding)
-      .def("ClearSharding", &XlaBuilder::ClearSharding);
+      .def("ClearSharding", &XlaBuilder::ClearSharding)
+      .def("SetUpAlias",
+           [](XlaBuilder& builder, const std::vector<int64>& output_index,
+              int64 param_number, const std::vector<int64>& param_index) {
+             builder.SetUpAlias(
+                 ShapeIndex(output_index.begin(), output_index.end()),
+                 param_number,
+                 ShapeIndex(param_index.begin(), param_index.end()));
+           });
 
   m.def("BufferToDLPackManagedTensor", BufferToDLPackManagedTensor);
   m.def("DLPackManagedTensorToBuffer",

tensorflow/compiler/xla/python/xla_client.py

@@ -719,6 +719,16 @@ def Build(self, root=None, backend=None):
     else:
       return Computation(self._builder.Build(), backend=backend)
 
+  def SetUpAlias(self, output_index, param_number, param_index):
+    """Adds a new input/output alias.
+
+    Args:
+      output_index: Iterable of int64 specifying the output index.
+      param_number: Parameter number.
+      param_index: Iterable of int64 specifying parameter index.
+    """
+    return self._builder.SetUpAlias(output_index, param_number, param_index)
+
   def GetShape(self, operand):
     return self._builder.GetShape(operand)
 

tensorflow/compiler/xla/python/xla_client_test.py

@@ -2022,6 +2022,21 @@ def testSetSharding(self):
     np.testing.assert_allclose(ans, 4.14)
 
 
+class AliasTest(ComputationTest):
+
+  def testSetUpAlias(self):
+    c = self._NewComputation()
+    p1 = c.ParameterFromNumpy(NumpyArrayF32(1.0))
+    p2 = c.ParameterFromNumpy(NumpyArrayF32(1.0))
+    out = c.Add(p1, p2)
+    c.SetUpAlias([], 0, [])
+    c = c.Build(out)
+    with self.assertRaisesRegex(RuntimeError,
+                                "Buffer aliasing is not supported "
+                                "by XLA for non-TPU backends"):
+      c.Compile()
+
+
 int_dtypes = [
     np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32,
     np.uint64