| 607 | } // namespace |
| 608 | |
| 609 | PYBIND11_MODULE(xla_extension, m) { |
| 610 | // Caution: import_array1 works by initializing a static variable |
| 611 | // (PyArray_API) which is *defined* in a NumPy header. import_array1() must |
| 612 | // therefore be called from the *same translation unit* as any users of |
| 613 | // NumPy C APIs. |
| 614 | auto init_numpy = []() -> bool { |
| 615 | // import_array1 might look like a function. It's not. It's a macro that |
| 616 | // calls `return`, which is why we wrap it in this strange-looking lambda. |
| 617 | import_array1(false); |
| 618 | return true; |
| 619 | }; |
| 620 | if (!init_numpy() || !InitializeNumpyAPIForTypes()) { |
| 621 | throw std::runtime_error("Unable to initialize Numpy API"); |
| 622 | } |
| 623 | |
| 624 | // Types |
| 625 | py::enum_<PrimitiveType>(m, "PrimitiveType") |
| 626 | .value("PRIMITIVE_TYPE_INVALID", PRIMITIVE_TYPE_INVALID) |
| 627 | .value("PRED", PRED) |
| 628 | .value("S8", S8) |
| 629 | .value("S16", S16) |
| 630 | .value("S32", S32) |
| 631 | .value("S64", S64) |
| 632 | .value("U8", U8) |
| 633 | .value("U16", U16) |
| 634 | .value("U32", U32) |
| 635 | .value("U64", U64) |
| 636 | .value("F16", F16) |
| 637 | .value("BF16", BF16) |
| 638 | .value("F32", F32) |
| 639 | .value("F64", F64) |
| 640 | .value("C64", C64) |
| 641 | .value("C128", C128) |
| 642 | .value("TUPLE", TUPLE) |
| 643 | .value("OPAQUE_TYPE", OPAQUE_TYPE) |
| 644 | .value("TOKEN", TOKEN); |
| 645 | |
| 646 | m.def("bfloat16_dtype", Bfloat16Dtype); |
| 647 | |
| 648 | // Shapes |
| 649 | py::class_<Shape> shape_class(m, "Shape"); |
| 650 | shape_class |
| 651 | .def(py::init([](const string& s) { |
| 652 | return absl::make_unique<Shape>(ValueOrThrow(ParseShape(s))); |
| 653 | })) |
| 654 | .def_static( |
| 655 | "tuple_shape", |
| 656 | [](std::vector<Shape> shapes) -> Shape { |
| 657 | return ShapeUtil::MakeTupleShape(shapes); |
| 658 | }, |
| 659 | "Constructs a tuple shape.") |
| 660 | .def_static( |
| 661 | "array_shape", |
| 662 | [](PrimitiveType type, py::object dims_seq, |
| 663 | absl::optional<py::object> layout_seq) -> Shape { |
| 664 | std::vector<int64> dims = IntSequenceToVector(dims_seq); |
| 665 | if (layout_seq) { |
| 666 | std::vector<int64> layout = IntSequenceToVector(*layout_seq); |
nothing calls this directly
no test coverage detected