| 247 | } |
| 248 | |
| 249 | xla::XlaComputation BroadcastComputation( |
| 250 | const xla::Shape& shape, absl::Span<const xla::int64> dimensions) { |
| 251 | xla::XlaBuilder builder("BroadcastComputation"); |
| 252 | auto p0 = xla::Parameter(&builder, 0, shape, "P0"); |
| 253 | xla::Broadcast(p0, dimensions); |
| 254 | return builder.Build().ValueOrDie(); |
| 255 | } |
| 256 | |
| 257 | xla::XlaComputation IsEqualComputation(const xla::Shape& shape) { |
| 258 | xla::XlaBuilder builder("IsEqualComputation"); |