Some more XLATensor cleanup. (#1110)

Author: pschuh

Sources/CX10/xla_tensor_ops_wrapper.cc

@@ -144,6 +144,11 @@ xla::XlaOp LowerBinaryValueOp(xla::XlaOp lhs, xla::XlaOp rhs) {
   return T(lhs, rhs);
 }
 
+std::vector<xla::XlaOp> LowerBroadcastTensors(xla::XlaOp lhs, xla::XlaOp rhs) {
+  std::tie(lhs, rhs) = XlaHelpers::PromoteValues(lhs, rhs);
+  return {lhs, rhs};
+}
+
 xla::XlaOp LowerSqueeze(xla::XlaOp input, int dim) {
   if (dim == -1) return SqueezeAllTrivialDimensions(input);
   XLA_CHECK_GE(dim, 0);
@@ -324,9 +329,7 @@ xla::XlaOp LowerTfUnsortedSegmentSum(xla::XlaOp data, xla::XlaOp indices,
                                combine);
 }
 
-xla::XlaOp LowerTfStatelessRandomUniform(xla::Shape shape, xla::XlaOp seeds,
-                                         xla::XlaOp minval, xla::XlaOp maxval,
-                                         LoweringContext* loctx = nullptr) {
+xla::BitGeneratorTy GetBestGenerator(LoweringContext* loctx = nullptr) {
   xla::BitGeneratorTy generator;
   if (!loctx || loctx->device().hw_type == swift_xla::DeviceType::TPU) {
     generator = xla::ThreeFryBitGenerator;
@@ -336,6 +339,30 @@ xla::XlaOp LowerTfStatelessRandomUniform(xla::Shape shape, xla::XlaOp seeds,
       return xla::PhiloxBitGenerator(key, state, shape);
     };
   }
+  return generator;
+}
+
+xla::XlaOp LowerTfStatelessRandomNormal(xla::Shape shape, xla::XlaOp seeds,
+                                        at::ScalarType dtype,
+                                        LoweringContext* loctx = nullptr) {
+  auto generator = GetBestGenerator(loctx);
+  xla::XlaOp seed0 = xla::Reshape(xla::Slice(seeds, {0}, {1}, {1}), {});
+  xla::XlaOp seed1 = xla::Reshape(xla::Slice(seeds, {1}, {2}, {1}), {});
+  xla::XlaOp initial_state =
+      xla::ConstantR0WithType(seeds.builder(), xla::U64, 0);
+  xla::XlaOp key = ConvertElementType(seed0, xla::U64) |
+                   ShiftLeft(ConvertElementType(seed1, xla::U64),
+                             ConstantR0WithType(seeds.builder(), xla::U64, 32));
+  xla::XlaOp normal =
+      xla::NormalFloatingPointDistribution(key, initial_state, generator, shape)
+          .value;
+  return normal;
+}
+
+xla::XlaOp LowerTfStatelessRandomUniform(xla::Shape shape, xla::XlaOp seeds,
+                                         xla::XlaOp minval, xla::XlaOp maxval,
+                                         LoweringContext* loctx = nullptr) {
+  auto generator = GetBestGenerator(loctx);
   xla::XlaOp seed0 = xla::Reshape(xla::Slice(seeds, {0}, {1}, {1}), {});
   xla::XlaOp seed1 = xla::Reshape(xla::Slice(seeds, {1}, {2}, {1}), {});
   xla::XlaOp key = ConvertElementType(seed0, xla::U64) |

Sources/CX10/xla_tensor_ops_wrapper_generated.cc.inc

@@ -296,14 +296,6 @@ OpaqueXLATensor* XLATensor_arange(XLAScalar start, XLAScalar end,
                         ToScalarType(type));
   return new XLATensor(out);
 }
-OpaqueXLATensor_pair XLATensor_broadcast_tensors(OpaqueXLATensor* a,
-                                                 OpaqueXLATensor* b) {
-  OpaqueXLATensor_pair result;
-  auto output = XLATensor::broadcast_tensors({*a, *b});
-  result.x = new XLATensor(output[0]);
-  result.y = new XLATensor(output[1]);
-  return result;
-}
 OpaqueXLATensorArrayRef XLATensor_cross_replica_sum(
     OpaqueXLATensorArrayRef inputs, double scale) {
   auto token = swift_xla::ir::MakeNode<swift_xla::ir::ops::Token>();
@@ -334,17 +326,6 @@ OpaqueXLATensor* XLATensor_linspace(XLAScalar start, XLAScalar stop,
 OpaqueXLATensor* XLATensor_replica_id(const struct CDevice device) {
   return new XLATensor(XLATensor::xla_replica_id(ConvertDevice(device)));
 }
-OpaqueXLATensor* XLATensor_tf_StatelessRandomNormal(
-    Int64ArrayRef size, OpaqueXLATensor* seeds, const struct CDevice device,
-    enum XLATensorScalarType type) {
-  return new XLATensor(XLATensor::tf_StatelessRandomNormal(
-      size.slice(), *seeds, ConvertDevice(device), ToScalarType(type)));
-}
-OpaqueXLATensor* XLATensor_threshold_backward(OpaqueXLATensor* grad_output,
-                                              OpaqueXLATensor* input,
-                                              float threshold) {
-  return XLATensor_threshold(input, grad_output, threshold, 0);
-}
 OpaqueXLATensor* XLATensor_to(OpaqueXLATensor* a, const CDevice* device,
                               Optional_XLAScalarType dtype) {
   return new XLATensor(XLATensor::to(*a, AsOptional(device), dtype.value()));

Sources/CX10/xla_tensor_wrapper.cc

@@ -401,16 +401,13 @@ XLA_API OpaqueXLATensor*
 XLATensor_tf_OneHot(OpaqueXLATensor* indices, OpaqueXLATensor* on_value,
                     OpaqueXLATensor* off_value, int64_t depth, int64_t axis);
 XLA_API OpaqueXLATensor* XLATensor_tf_StatelessRandomNormal(
-    Int64ArrayRef size, OpaqueXLATensor* seeds, const struct CDevice device,
-    enum XLATensorScalarType type);
+    Int64ArrayRef size, OpaqueXLATensor* seeds, enum XLATensorScalarType type);
 XLA_API OpaqueXLATensor* XLATensor_tf_StatelessRandomUniform(
     Int64ArrayRef size, OpaqueXLATensor* seeds, OpaqueXLATensor* minvalue,
     OpaqueXLATensor* maxvalue);
 XLA_API OpaqueXLATensor*
 XLATensor_tf_UnsortedSegmentSum(OpaqueXLATensor* data, OpaqueXLATensor* indices,
                                 int64_t num_segments);
-XLA_API OpaqueXLATensor* XLATensor_threshold_backward(
-    OpaqueXLATensor* grad_output, OpaqueXLATensor* input, float threshold);
 XLA_API OpaqueXLATensor* XLATensor_threshold(
     OpaqueXLATensor* input, OpaqueXLATensor* output, float threshold, float value);
 XLA_API OpaqueXLATensor* XLATensor_truncated_normal(OpaqueXLATensor* input);

Sources/CX10/xla_tensor_wrapper.h

@@ -125,6 +125,20 @@ extension _RawXLA {
     return Tensor(_xlaHandle: XLATensor_atanh(input.xlaHandle))
   }
 
+  public static func broadcast_tensors<
+    T: TensorFlowScalar
+  >(
+    _ lhs: Tensor<T>,
+    _ rhs: Tensor<T>
+  ) -> (Tensor<T>, Tensor<T>) {
+    defer { _fixLifetime(lhs) }
+    defer { _fixLifetime(rhs) }
+    checkSameDevice(lhs.device, rhs.device)
+    checkSamePrecision(lhs, rhs)
+    let tuple_output = XLATensor_broadcast_tensors(lhs.xlaHandle, rhs.xlaHandle)
+    return (Tensor(_xlaHandle: tuple_output.x), Tensor(_xlaHandle: tuple_output.y))
+  }
+
   public static func concat<
     T: TensorFlowScalar
   >(
@@ -474,6 +488,17 @@ extension _RawXLA {
     return Tensor(_xlaHandle: XLATensor_logicalAnd(lhs.xlaHandle, rhs.xlaHandle))
   }
 
+  static func logicalCast<
+    Srct: TensorFlowScalar,
+    Dstt: TensorFlowScalar
+  >(
+    _ input: Tensor<Srct>,
+    destType: ScalarType
+  ) -> Tensor<Dstt> {
+    defer { _fixLifetime(input) }
+    return Tensor(_xlaHandle: XLATensor_logical_cast(input.xlaHandle, destType))
+  }
+
   public static func logicalNot(
     _ input: Tensor<Bool>
   ) -> Tensor<Bool> {
@@ -639,6 +664,16 @@ extension _RawXLA {
     }
   }
 
+  static func physicalCast<
+    T: TensorFlowScalar
+  >(
+    _ input: Tensor<T>,
+    destType: ScalarType
+  ) -> Tensor<T> {
+    defer { _fixLifetime(input) }
+    return Tensor(_xlaHandle: XLATensor_physical_cast(input.xlaHandle, destType))
+  }
+
   public static func pow<
     T: TensorFlowNumeric
   >(
@@ -990,7 +1025,7 @@ extension _RawXLA {
   }
 
   static func tf_MirrorPad<
-    T: TensorFlowNumeric
+    T: TensorFlowScalar
   >(
     _ input: Tensor<T>,
     _ padding: [Int64],
@@ -1023,20 +1058,86 @@ extension _RawXLA {
     Ti: TensorFlowInteger,
     T: TensorFlowScalar
   >(
-    indices: Tensor<Ti>,
-    on_value: Tensor<T>,
-    off_value: Tensor<T>,
-    depth: Int64,
-    axis: Int64
+    _ indices: Tensor<Ti>,
+    _ onValue: Tensor<T>,
+    _ offValue: Tensor<T>,
+    _ depth: Int64,
+    _ axis: Int64
   ) -> Tensor<T> {
     defer { _fixLifetime(indices) }
-    defer { _fixLifetime(on_value) }
-    defer { _fixLifetime(off_value) }
-    checkSameDevice(indices.device, on_value.device)
-    checkSameDevice(indices.device, off_value.device)
+    defer { _fixLifetime(onValue) }
+    defer { _fixLifetime(offValue) }
+    checkSameDevice(indices.device, onValue.device)
+    checkSameDevice(indices.device, offValue.device)
     return Tensor(
       _xlaHandle: XLATensor_tf_OneHot(
-        indices.xlaHandle, on_value.xlaHandle, off_value.xlaHandle, depth, axis))
+        indices.xlaHandle, onValue.xlaHandle, offValue.xlaHandle, depth, axis))
+  }
+
+  static func tf_StatelessRandomNormal<
+    T: TensorFlowScalar,
+    Ti: TensorFlowIndex
+  >(
+    _ shape: [Int64],
+    _ seeds: Tensor<Ti>,
+    dtype: ScalarType
+  ) -> Tensor<T> {
+    defer { _fixLifetime(seeds) }
+    return shape.withArrayRef { shape in
+      return Tensor(_xlaHandle: XLATensor_tf_StatelessRandomNormal(shape, seeds.xlaHandle, dtype))
+    }
+  }
+
+  static func tf_StatelessRandomUniform<
+    T: TensorFlowScalar,
+    Ti: TensorFlowIndex
+  >(
+    _ shape: [Int64],
+    _ seeds: Tensor<Ti>,
+    _ minvalue: Tensor<T>,
+    _ maxvalue: Tensor<T>
+  ) -> Tensor<T> {
+    defer { _fixLifetime(seeds) }
+    defer { _fixLifetime(minvalue) }
+    defer { _fixLifetime(maxvalue) }
+    checkSameDevice(seeds.device, minvalue.device)
+    checkSameDevice(seeds.device, maxvalue.device)
+    return shape.withArrayRef { shape in
+      return Tensor(
+        _xlaHandle: XLATensor_tf_StatelessRandomUniform(
+          shape, seeds.xlaHandle, minvalue.xlaHandle, maxvalue.xlaHandle))
+    }
+  }
+
+  static func tf_UnsortedSegmentSum<
+    T: TensorFlowNumeric,
+    Ti: TensorFlowIndex
+  >(
+    _ data: Tensor<T>,
+    indicies: Tensor<Ti>,
+    numSegments: Int64
+  ) -> Tensor<T> {
+    defer { _fixLifetime(data) }
+    defer { _fixLifetime(indicies) }
+    checkSameDevice(data.device, indicies.device)
+    return Tensor(
+      _xlaHandle: XLATensor_tf_UnsortedSegmentSum(data.xlaHandle, indicies.xlaHandle, numSegments))
+  }
+
+  static func threshold<
+    T: TensorFlowNumeric
+  >(
+    _ input: Tensor<T>,
+    output: Tensor<T>,
+    threshold: Float,
+    value: Float
+  ) -> Tensor<T> {
+    defer { _fixLifetime(input) }
+    defer { _fixLifetime(output) }
+    checkSameDevice(input.device, output.device)
+    checkSamePrecision(input, output)
+    return Tensor(
+      _xlaHandle: XLATensor_threshold(input.xlaHandle, output.xlaHandle, threshold, value))
   }
 
   public static func topk<
@@ -1052,6 +1153,15 @@ extension _RawXLA {
     return (Tensor(_xlaHandle: tuple_output.x), Tensor(_xlaHandle: tuple_output.y))
   }
 
+  static func truncatedNormal<
+    T: FloatingPoint & TensorFlowScalar
+  >(
+    _ input: Tensor<T>
+  ) -> Tensor<T> {
+    defer { _fixLifetime(input) }
+    return Tensor(_xlaHandle: XLATensor_truncated_normal(input.xlaHandle))
+  }
+
   public static func updateSlice<
     T: TensorFlowScalar
   >(
@@ -1085,6 +1195,20 @@ extension _RawXLA {
       _xlaHandle: XLATensor_where(condition.xlaHandle, input.xlaHandle, other.xlaHandle))
   }
 
+  static func xlaPad<
+    T: TensorFlowScalar
+  >(
+    _ input: Tensor<T>,
+    paddingValue: AnyScalar,
+    paddingConfig: [PaddingConfigDimension]
+  ) -> Tensor<T> {
+    defer { _fixLifetime(input) }
+    return paddingConfig.withArrayRef { paddingConfig in
+      return Tensor(
+        _xlaHandle: XLATensor_xla_pad(input.xlaHandle, paddingValue.xlaScalar, paddingConfig))
+    }
+  }
+
   public static func xlaSlice<
     T: TensorFlowScalar
   >(

Sources/x10/swift_bindings/RawOpsXLAGenerated.swift

@@ -210,7 +210,7 @@ extension Array where Element: AnyTensor {
 }
 
 extension Array where Element == PaddingConfigDimension {
-  func withPaddingConfig<Result>(_ body: (inout PaddingConfig) -> Result) -> Result {
+  func withArrayRef<Result>(_ body: (inout PaddingConfig) -> Result) -> Result {
     defer { _fixLifetime(self) }
     return withUnsafeBufferPointer {
       (_ dimensions: UnsafeBufferPointer<PaddingConfigDimension>) -> Result in
@@ -330,11 +330,6 @@ extension XLATensor {
         start.xlaScalar, stop.xlaScalar, num, cdevice, type))
   }
 
-  static func logicalCast(_ input: XLATensor, destType: XLATensorScalarType) -> XLATensor {
-    defer { _fixLifetime(input) }
-    return XLATensor(_handle: XLATensor_logical_cast(input.handle, destType))
-  }
-
   static func maxpool(
     _ input: XLATensor,
     _ ksize: [Int64],
@@ -368,104 +363,10 @@ extension XLATensor {
     }
   }
 
-  static func mirrorPad(_ input: XLATensor, _ padding: [Int64], _ mode: TFMirrorPadMode)
-    -> XLATensor
-  {
-    defer { _fixLifetime(input) }
-    return padding.withArrayRef { padding in
-      XLATensor(_handle: XLATensor_tf_MirrorPad(input.handle, padding, mode))
-    }
-  }
-
-  static func mirrorPadGrad(
-    _ grad_output: XLATensor, _ inputSize: [Int64], _ padding: [Int64], _ mode: TFMirrorPadMode
-  )
-    -> XLATensor
-  {
-    defer { _fixLifetime(grad_output) }
-    return inputSize.withArrayRef { inputSize in
-      padding.withArrayRef { padding in
-        XLATensor(
-          _handle: XLATensor_tf_MirrorPadGrad(grad_output.handle, inputSize, padding, mode))
-      }
-    }
-  }
-
-  static func physicalCast(_ input: XLATensor, destType: XLATensorScalarType) -> XLATensor {
-    defer { _fixLifetime(input) }
-    return XLATensor(_handle: XLATensor_physical_cast(input.handle, destType))
-  }
-
   static func replica_id(_ device: Device) -> XLATensor {
     return XLATensor(_handle: XLATensor_replica_id(device.cdevice))
   }
 
-  static func tf_OneHot(
-    _ indices: XLATensor, _ on_value: XLATensor, _ off_value: XLATensor, _ depth: Int64,
-    _ axis: Int64
-  ) -> XLATensor {
-    defer { _fixLifetime(indices) }
-    defer { _fixLifetime(on_value) }
-    defer { _fixLifetime(off_value) }
-    return XLATensor(
-      _handle: XLATensor_tf_OneHot(indices.handle, on_value.handle, off_value.handle, depth, axis))
-  }
-
-  static func tf_StatelessRandomNormal(
-    _ dims: [Int64],
-    _ seeds: XLATensor,
-    _ dtype: XLAScalarType.Type,
-    _ device: Device
-  ) -> XLATensor {
-    defer { _fixLifetime(seeds) }
-    let cdevice = device.cdevice
-    return dims.withArrayRef { dims in
-      XLATensor(
-        _handle: XLATensor_tf_StatelessRandomNormal(
-          dims, seeds.handle, cdevice,
-          dtype.xlaTensorScalarType))
-    }
-  }
-
-  static func tf_StatelessRandomUniform(
-    _ dims: [Int64],
-    _ seeds: XLATensor,
-    _ minvalue: XLATensor,
-    _ maxvalue: XLATensor,
-    _ dtype: XLAScalarType.Type,
-    _ device: Device
-  ) -> XLATensor {
-    defer { _fixLifetime(seeds) }
-    return dims.withArrayRef { dims in
-      XLATensor(
-        _handle: XLATensor_tf_StatelessRandomUniform(
-          dims, seeds.handle, minvalue.handle, maxvalue.handle))
-    }
-  }
-
-  static func tf_UnsortedSegmentSum(
-    _ data: XLATensor, _ indices: XLATensor, _ numSegments: Int64
-  ) -> XLATensor {
-    defer { _fixLifetime(data) }
-    defer { _fixLifetime(indices) }
-    return XLATensor(
-      _handle: XLATensor_tf_UnsortedSegmentSum(data.handle, indices.handle, numSegments))
-  }
-
-  static func threshold_backward(_ grad_output: XLATensor, _ input: XLATensor, _ threshold: Float)
-    -> XLATensor
-  {
-    defer { _fixLifetime(grad_output) }
-    defer { _fixLifetime(input) }
-    return XLATensor(
-      _handle: XLATensor_threshold_backward(grad_output.handle, input.handle, threshold))
-  }
-
-  static func truncatedNormal(_ input: XLATensor) -> XLATensor {
-    defer { _fixLifetime(input) }
-    return XLATensor(_handle: XLATensor_truncated_normal(input.handle))
-  }
-
   static func to(
     _ a: XLATensor, _ device: Device?, _ dtype: XLAScalarType.Type?
   ) -> XLATensor {
@@ -477,15 +378,6 @@ extension XLATensor {
     }
   }
 
-  static func xlaPad(
-    _ input: XLATensor, paddingValue: XLAScalarType, paddingConfig: [PaddingConfigDimension]
-  ) -> XLATensor {
-    defer { _fixLifetime(input) }
-    return paddingConfig.withPaddingConfig { paddingConfig in
-      XLATensor(_handle: XLATensor_xla_pad(input.handle, paddingValue.xlaScalar, paddingConfig))
-    }
-  }
-
   struct StridedSliceSpec {
     let begin: [Int64]
     let end: [Int64]

Sources/x10/swift_bindings/XLATensor.swift

@@ -26,6 +26,7 @@ public enum _RawXLA {
   public typealias Padding1 = _RawTFEager.Padding1
   public typealias Mode5 = _RawTFEager.Mode1
   typealias AnyScalar = XLAScalarType
+  typealias ScalarType = XLATensorScalarType
 
   private static func canonicalDims(_ dims: [Int64], _ rank: Int64) -> [Int64] {
     dims.map { $0 < 0 ? $0 + rank : $0 }
@@ -541,7 +542,7 @@ public enum _RawXLA {
     _ x: Tensor<Srct>,
     truncate: Bool = false
   ) -> Tensor<Dstt> {
-    Tensor(_xla: XLATensor.logicalCast(x.xlaTensor, destType: Dstt.xlaTensorScalarType))
+    return logicalCast(x, destType: Dstt.xlaTensorScalarType)
   }
 
   /// Concatenates tensors along one dimension.
@@ -1916,9 +1917,7 @@ public enum _RawXLA {
     mode: Mode5
   ) -> Tensor<T> {
     let linearizedPaddings = paddings.scalars.map { Int64($0) }
-    return Tensor(
-      _xla: XLATensor.mirrorPad(
-        input.xlaTensor, reversedPaddings(linearizedPaddings), convertMirrorPadMode(mode)))
+    return tf_MirrorPad(input, reversedPaddings(linearizedPaddings), convertMirrorPadMode(mode))
   }
 
   /// Gradient op for `MirrorPad` op. This op folds a mirror-padded tensor.
@@ -1963,10 +1962,9 @@ public enum _RawXLA {
       let totalPadding = linearizedPaddings[2 * dim] + linearizedPaddings[2 * dim + 1]
       return Int64(grad.shape.dimensions[dim]) - totalPadding
     }
-    return Tensor(
-      _xla: XLATensor.mirrorPadGrad(
-        grad.xlaTensor, inputDimensions, reversedPaddings(linearizedPaddings),
-        convertMirrorPadMode(mode)))
+    return tf_MirrorPadGrad(
+      grad, inputDimensions, reversedPaddings(linearizedPaddings),
+      convertMirrorPadMode(mode))
   }
 
   /// Returns the truth value of (x != y) element-wise.
@@ -2091,12 +2089,7 @@ public enum _RawXLA {
     offValue: Tensor<T>,
     axis: Int64 = -1
   ) -> Tensor<T> {
-    checkSameDevice(onValue, offValue)
-    checkSameDevice(indices.device, onValue.device)
-    checkSamePrecision(onValue, offValue)
-    return Tensor(
-      _xla: XLATensor.tf_OneHot(
-        indices.xlaTensor, onValue.xlaTensor, offValue.xlaTensor, depth, axis))
+    return tf_OneHot(indices, onValue, offValue, depth, axis)
   }
   public static func oneHot<
     T: TensorFlowScalar,
@@ -2108,12 +2101,7 @@ public enum _RawXLA {
     offValue: Tensor<T>,
     axis: Int64 = -1
   ) -> Tensor<T> {
-    checkSameDevice(onValue, offValue)
-    checkSameDevice(indices.device, onValue.device)
-    checkSamePrecision(onValue, offValue)
-    return Tensor(
-      _xla: XLATensor.tf_OneHot(
-        indices.xlaTensor, onValue.xlaTensor, offValue.xlaTensor, Int64(depth.scalarized()), axis))
+    return tf_OneHot(indices, onValue, offValue, Int64(depth.scalarized()), axis)
   }
 
   /// Returns a tensor of ones with the same shape and type as x.
@@ -2250,7 +2238,7 @@ public enum _RawXLA {
   public static func physicalCast<T: TensorFlowScalar, R: TensorFlowScalar>(
     _ input: Tensor<T>, destType: R.Type
   ) -> Tensor<T> {
-    Tensor(_xla: XLATensor.physicalCast(input.xlaTensor, destType: destType.xlaTensorScalarType))
+    physicalCast(input, destType: destType.xlaTensorScalarType)
   }
 
   /// Computes the product of elements across dimensions of a tensor.
@@ -2379,9 +2367,7 @@ public enum _RawXLA {
     gradients: Tensor<T>,
     features: Tensor<T>
   ) -> Tensor<T> {
-    checkSameDevice(gradients, features)
-    checkSamePrecision(gradients, features)
-    return Tensor(_xla: XLATensor.threshold_backward(gradients.xlaTensor, features.xlaTensor, 0))
+    return threshold(features, output: gradients, threshold: 0, value: 0)
   }
 
   public static func replicaId(_ device: Device) -> Tensor<Int32> {
@@ -3138,10 +3124,8 @@ public enum _RawXLA {
     seed: Tensor<Tseed>,
     device: Device
   ) -> Tensor<Dtype> {
-    Tensor(
-      _xla: XLATensor.tf_StatelessRandomNormal(
-        shape.scalars.map { Int64($0) },
-        seed.xlaTensor, Dtype.self, device))
+    tf_StatelessRandomNormal(
+      shape.scalars.map { Int64($0) }, seed, dtype: Dtype.xlaTensorScalarType)
   }
 
   public static func statelessRandomNormal<
@@ -3179,11 +3163,10 @@ public enum _RawXLA {
     seed: Tensor<Tseed>,
     device: Device
   ) -> Tensor<Dtype> {
-    Tensor(
-      _xla: XLATensor.tf_StatelessRandomUniform(
-        shape.scalars.map { Int64($0) },
-        seed.xlaTensor, Tensor<Dtype>(0, on: device).xlaTensor,
-        Tensor<Dtype>(1, on: device).xlaTensor, Dtype.self, device))
+    tf_StatelessRandomUniform(
+      shape.scalars.map { Int64($0) },
+      seed, Tensor<Dtype>(0, on: device),
+      Tensor<Dtype>(1, on: device))
   }
 
   public static func statelessRandomUniform<
@@ -3224,11 +3207,7 @@ public enum _RawXLA {
     maxval: Tensor<Dtype>,
     device: Device
   ) -> Tensor<Dtype> {
-    Tensor(
-      _xla: XLATensor.tf_StatelessRandomUniform(
-        shape.scalars.map { Int64($0) },
-        seed.xlaTensor, minval.xlaTensor,
-        maxval.xlaTensor, Dtype.self, device))
+    tf_StatelessRandomUniform(shape.scalars.map { Int64($0) }, seed, minval, maxval)
   }
 
   public static func statelessRandomUniformInt<
@@ -3272,11 +3251,10 @@ public enum _RawXLA {
   ) -> Tensor<Dtype> {
     let minval = Tensor<Dtype>(Dtype.leastNormalMagnitude, on: device)
     let maxval = Tensor<Dtype>(1, on: device)
-    let uniform = XLATensor.tf_StatelessRandomUniform(
+    let uniform = tf_StatelessRandomUniform(
       shape.scalars.map { Int64($0) },
-      seed.xlaTensor, minval.xlaTensor,
-      maxval.xlaTensor, Dtype.self, device)
-    return Tensor(_xla: XLATensor.truncatedNormal(uniform))
+      seed, minval, maxval)
+    return truncatedNormal(uniform)
   }
 
   public static func statelessTruncatedNormal<
@@ -3548,8 +3526,7 @@ public enum _RawXLA {
     if !dimensionsToReverse.isEmpty {
       grad = flip(grad, dims: dimensionsToReverse)
     }
-    return Tensor(
-      _xla: XLATensor.xlaPad(grad.xlaTensor, paddingValue: 0, paddingConfig: paddingConfig))
+    return xlaPad(grad, paddingValue: 0, paddingConfig: paddingConfig)
   }
 
   /// Computes the sum of elements across dimensions of a tensor.
@@ -3819,9 +3796,7 @@ public enum _RawXLA {
             + String(segmentIds.shape.dimensions[dim]))
       }
     }
-    return Tensor(
-      _xla: XLATensor.tf_UnsortedSegmentSum(
-        data.xlaTensor, segmentIds.xlaTensor, Int64(numSegments)))
+    return tf_UnsortedSegmentSum(data, indicies: segmentIds, numSegments: Int64(numSegments))
   }
 
   /// Returns 0 if x == 0, and x / y otherwise, elementwise.

Sources/x10/swift_bindings/apis/RawOpsManual.swift

@@ -23,7 +23,7 @@
     "TFMirrorPadMode":
         ("enum TFMirrorPadMode", lambda name: f"ToTFMirrorPadMode({name})",
          "tensorflow::MirrorPadMode"),
-    "PaddingConfig":
+    "[PaddingConfigDimension]":
         ("PaddingConfig", lambda name: f"ToXLAPaddingConfig({name})",
          "xla::PaddingConfig"),
 }
@@ -229,6 +229,14 @@ def c_function_define(op):
     else:
       first_tensor = tensor_args[0][0]
 
+  def listify(l):
+    if type(l) is list:
+      return l
+    return [l]
+
+  dtypes = (([None] * op["n_results"])
+            if "result_dtype" not in op else listify(op["result_dtype"]))
+
   def format_arg_def(arg):
     name, stype, _ = arg
     if stype == "Tensor": return "OpaqueXLATensor* " + name
@@ -253,8 +261,18 @@ def format_arg_ref(arg):
     if stype == "[Tensor]":
       return name + "_ir_value"
     if name == "shape":
+      relement_type = f"{first_tensor}->shape().get().element_type()"
+      result_dtype_arg = None
+      if dtypes[0] and first_tensor != dtypes[0]:
+        for arg in args:
+          if arg[0] == dtypes[0]:
+            result_dtype_arg = arg
+      if result_dtype_arg:
+        relement_type = (
+            f"swift_xla::MakeXlaPrimitiveType({format_arg_ref(result_dtype_arg)},"
+            f" /*device=*/nullptr)")
       return ("swift_xla::MakeArrayShapeFromDimensions(shape.slice(), {}, " +
-              f"{first_tensor}->shape().get().element_type(), "
+              f"{relement_type}, "
               f"{first_tensor}->GetDevice().hw_type)")
     if stype in builtin_types:
       return builtin_types[stype][1](name)
@@ -298,14 +316,6 @@ def unpack_arg(arg):
         f"""{op["c_name"]} has unsupported number of return values {op["n_results"]}"""
     )
 
-  def listify(l):
-    if type(l) is list:
-      return l
-    return [l]
-
-  dtypes = (([None] * op["n_results"])
-            if "result_dtype" not in op else listify(op["result_dtype"]))
-
   def format_result(result_i=0, dtype=None):
     if not dtype:
       dtype = dtypes[result_i]

Sources/x10/swift_bindings/generate_ops.py

@@ -60,6 +60,10 @@
   lower_fn: xla::Atanh
   generics: {T: FloatingPoint & TensorFlowScalar}
 
+- def: "broadcast_tensors(_ lhs: Tensor<T>, _ rhs: Tensor<T>) -> (Tensor<T>, Tensor<T>)"
+  lower_fn: LowerBroadcastTensors
+  generics: {T: TensorFlowScalar}
+
 - def: "cat(_ input: [Tensor<T>], dim: Int64) -> Tensor<T>"
   extras: ["canonicalize dim input CanonicalizeCat"]
   swift_name: concat
@@ -235,11 +239,14 @@
   x10_enum: at::aten::logical_and
   lower_fn: LowerBinaryOp<xla::And>
 
-- def: "logical_cast(_ input: Tensor, dtype: ScalarType) -> Tensor"
+- def: "logical_cast(_ input: Tensor<Srct>, destType: ScalarType) -> Tensor<Dstt>"
   x10_enum: xla_symbols::cast
+  generics: {Srct: TensorFlowScalar, Dstt: TensorFlowScalar}
+  swift_name: logicalCast
+  protection: internal
   shape_fn: ShapeLogicalCast
   lower_fn: LowerLogicalCast
-  result_dtype: dtype
+  result_dtype: destType
 
 - def: "logicalNot(_ input: Tensor<Bool>) -> Tensor<Bool>"
   x10_enum: at::aten::bitwise_not
@@ -314,7 +321,10 @@
   generics: {T: TensorFlowScalar}
   lower_fn: xla::Transpose
 
-- def: "physical_cast(input: Tensor, dtype: ScalarType) -> Tensor"
+- def: "physical_cast(_ input: Tensor<T>, destType: ScalarType) -> Tensor<T>"
+  generics: {T: TensorFlowScalar}
+  swift_name: physicalCast
+  protection: internal
   x10_enum: xla_symbols::cast
   shape_fn: ShapeLogicalCast
   lower_fn: LowerLogicalCast
@@ -458,7 +468,7 @@
 
 - def: "tf_MirrorPad(_ input: Tensor<T>, _ padding: [Int64], _ mode: TFMirrorPadMode) -> Tensor<T>"
   x10_enum: at::aten::tf_mirror_pad
-  generics: {T: TensorFlowNumeric}
+  generics: {T: TensorFlowScalar}
   protection: internal
   lower_fn: BuildMirrorPad
 
@@ -468,25 +478,39 @@
   protection: internal
   lower_fn: BuildMirrorPadBackward
 
-- def: "tf_OneHot(indices: Tensor<Ti>, on_value: Tensor<T>, off_value: Tensor<T>, depth: Int64, axis: Int64) -> Tensor<T>"
-  result_dtype: on_value
+- def: "tf_OneHot(_ indices: Tensor<Ti>, _ onValue: Tensor<T>, _ offValue: Tensor<T>, _ depth: Int64, _ axis: Int64) -> Tensor<T>"
+  result_dtype: onValue
   generics: {Ti: TensorFlowInteger, T: TensorFlowScalar}
   protection: internal
   x10_enum: at::aten::tf_one_hot
   lower_fn: BuildOneHot
 
-- def: "tf_StatelessRandomUniform(shape: [Int64], seeds: Tensor, minvalue: Tensor, maxvalue: Tensor) -> Tensor"
+- def: "tf_StatelessRandomNormal(_ shape: [Int64], _ seeds: Tensor<Ti>, dtype: ScalarType) -> Tensor<T>"
+  x10_enum: at::aten::tf_stateless_random_normal
+  generics: {T: TensorFlowScalar, Ti: TensorFlowIndex}
+  extras: ["shape_fn shape", "needs_lowering_context"]
+  lower_fn: LowerTfStatelessRandomNormal
+  protection: internal
+  result_dtype: dtype
+
+- def: "tf_StatelessRandomUniform(_ shape: [Int64], _ seeds: Tensor<Ti>, _ minvalue: Tensor<T>, _ maxvalue: Tensor<T>) -> Tensor<T>"
   x10_enum: at::aten::tf_stateless_random_uniform
+  generics: {T: TensorFlowScalar, Ti: TensorFlowIndex}
   extras: ["shape_fn shape", "needs_lowering_context"]
   lower_fn: LowerTfStatelessRandomUniform
+  protection: internal
   result_dtype: minvalue
 
-- def: "tf_UnsortedSegmentSum(data: Tensor, indicies: Tensor, num_segments: Int64) -> Tensor"
+- def: "tf_UnsortedSegmentSum(_ data: Tensor<T>, indicies: Tensor<Ti>, numSegments: Int64) -> Tensor<T>"
+  generics: {T: TensorFlowNumeric, Ti: TensorFlowIndex}
   x10_enum: at::aten::tf_unsorted_segment_sum
+  protection: internal
   lower_fn: LowerTfUnsortedSegmentSum
 
-- def: "threshold(input: Tensor, output: Tensor, threshold: Float, value: Float) -> Tensor"
+- def: "threshold(_ input: Tensor<T>, output: Tensor<T>, threshold: Float, value: Float) -> Tensor<T>"
   x10_enum: at::aten::threshold_backward
+  generics: {T: TensorFlowNumeric}
+  protection: internal
   shape_fn: input
   lower_fn: BuildThreshold
 
@@ -496,8 +520,11 @@
   generics: {T: FloatingPoint & TensorFlowScalar}
   result_dtype: [input, Long]
 
-- def: "truncated_normal(input: Tensor) -> Tensor"
+- def: "truncated_normal(_ input: Tensor<T>) -> Tensor<T>"
   x10_enum: at::aten::xla_truncated_normal
+  generics: {T: FloatingPoint & TensorFlowScalar}
+  protection: internal
+  swift_name: truncatedNormal
   shape_fn: input
   lower_fn: tensorflow::TruncatedNormal
 
@@ -513,7 +540,10 @@
   swift_name: where_
   lower_fn: LowerWhere
 
-- def: "xla_pad(input: Tensor, padding_value: AnyScalar, padding_config: PaddingConfig) -> Tensor"
+- def: "xla_pad(_ input: Tensor<T>, paddingValue: AnyScalar, paddingConfig: [PaddingConfigDimension]) -> Tensor<T>"
+  generics: {T: TensorFlowScalar}
+  protection: internal
+  swift_name: xlaPad
   lower_fn: LowerPad
 
 - def: "xla_slice(_ input: Tensor<T>, start_indices: [Int64], limit_indices: [Int64], strides: [Int64]) -> Tensor<T>"