| 308 | } |
| 309 | |
| 310 | void BuildOpsSubmodule(py::module* m) { |
| 311 | // ops submodule, containing free functions that add operators to an |
| 312 | // XlaBuilder. |
| 313 | py::module ops = m->def_submodule("ops", "XLA operations"); |
| 314 | |
| 315 | ops.def("AfterAll", &AfterAll); |
| 316 | ops.def( |
| 317 | "AllReduce", |
| 318 | static_cast<XlaOp (*)( |
| 319 | XlaOp, const XlaComputation&, absl::Span<const ReplicaGroup>, |
| 320 | const absl::optional<ChannelHandle>&, const absl::optional<Shape>&)>( |
| 321 | &AllReduce)); |
| 322 | ops.def("AllToAll", &AllToAll); |
| 323 | ops.def("CollectivePermute", &CollectivePermute); |
| 324 | ops.def("CreateToken", &CreateToken); |
| 325 | ops.def("CrossReplicaSum", |
| 326 | static_cast<XlaOp (*)(XlaOp, absl::Span<const ReplicaGroup>)>( |
| 327 | &CrossReplicaSum)); |
| 328 | ops.def("BitcastConvertType", &BitcastConvertType, py::arg("operand"), |
| 329 | py::arg("new_element_type")); |
| 330 | ops.def("Broadcast", &Broadcast, py::arg("operand"), py::arg("sizes")); |
| 331 | ops.def("BroadcastInDim", &BroadcastInDim, py::arg("operand"), |
| 332 | py::arg("shape"), py::arg("broadcast_dimensions")); |
| 333 | ops.def("Call", &Call); |
| 334 | ops.def("Cholesky", &Cholesky, py::arg("a"), py::arg("lower") = true); |
| 335 | ops.def("Clamp", &Clamp); |
| 336 | ops.def("Collapse", &Collapse, py::arg("operand"), py::arg("dimensions")); |
| 337 | ops.def("ConcatInDim", &ConcatInDim); |
| 338 | ops.def("Conditional", |
| 339 | static_cast<XlaOp (*)(XlaOp, absl::Span<const XlaComputation* const>, |
| 340 | absl::Span<const XlaOp>)>(&Conditional)); |
| 341 | ops.def("Conditional", |
| 342 | static_cast<XlaOp (*)(XlaOp, XlaOp, const XlaComputation&, XlaOp, |
| 343 | const XlaComputation&)>(&Conditional)); |
| 344 | ops.def("ConstantLiteral", &ConstantLiteral); |
| 345 | ops.def("ConvGeneralDilated", &ConvGeneralDilated, py::arg("lhs"), |
| 346 | py::arg("rhs"), py::arg("window_strides"), py::arg("padding"), |
| 347 | py::arg("lhs_dilation"), py::arg("rhs_dilation"), |
| 348 | py::arg("dimension_numbers"), py::arg("feature_group_count") = 1, |
| 349 | py::arg("batch_group_count") = 1, |
| 350 | py::arg("precision_config") = nullptr); |
| 351 | ops.def("ConvertElementType", &ConvertElementType, py::arg("operand"), |
| 352 | py::arg("new_element_type")); |
| 353 | // TODO(phawkins): remove CustomCall after callers are updated to use |
| 354 | // CustomCallWithLayout. |
| 355 | ops.def("CustomCall", &CustomCallWithLayout); |
| 356 | ops.def("CustomCallWithLayout", &CustomCallWithLayout); |
| 357 | ops.def("Dot", &Dot, py::arg("lhs"), py::arg("rhs"), |
| 358 | py::arg("precision_config") = nullptr); |
| 359 | ops.def("DotGeneral", &DotGeneral, py::arg("lhs"), py::arg("rhs"), |
| 360 | py::arg("dimension_numbers"), py::arg("precision_config") = nullptr); |
| 361 | ops.def("DynamicSlice", |
| 362 | static_cast<XlaOp (*)(XlaOp, absl::Span<const XlaOp>, |
| 363 | absl::Span<const int64>)>(&DynamicSlice)); |
| 364 | ops.def("DynamicUpdateSlice", |
| 365 | static_cast<XlaOp (*)(XlaOp, XlaOp, absl::Span<const XlaOp>)>( |
| 366 | &DynamicUpdateSlice)); |
| 367 |
no test coverage detected