| 235 | } |
| 236 | |
| 237 | xla::XlaComputation AddAndSubTuple() { |
| 238 | xla::XlaBuilder builder("AddAndSubTuple"); |
| 239 | auto p0 = xla::Parameter(&builder, 0, xla::ShapeUtil::MakeShape(xla::F32, {}), |
| 240 | "P0"); |
| 241 | auto p1 = xla::Parameter(&builder, 1, xla::ShapeUtil::MakeShape(xla::F32, {}), |
| 242 | "P1"); |
| 243 | auto sum = xla::Add(p0, p1); |
| 244 | auto sub = xla::Sub(p0, p1); |
| 245 | xla::Tuple(&builder, {sum, sub}); |
| 246 | return builder.Build().ValueOrDie(); |
| 247 | } |
| 248 | |
| 249 | xla::XlaComputation BroadcastComputation( |
| 250 | const xla::Shape& shape, absl::Span<const xla::int64> dimensions) { |