| 954 | } |
| 955 | |
| 956 | py::object _broadcast_cpp(py::handle input, py::handle args) { |
| 957 | py::object shape = _expand_args(args); |
| 958 | py::list dims; |
| 959 | bool all_imm; |
| 960 | if (PyList_Check(shape.ptr()) || PyTuple_Check(shape.ptr())) { |
| 961 | dims = py::reinterpret_steal<py::list>(PySequence_List(shape.ptr())); |
| 962 | mgb_assert(!dims.is_none()); |
| 963 | all_imm = true; |
| 964 | py::object inp_shape = py::none(); |
| 965 | size_t inp_ndim; |
| 966 | for (size_t i = 0; i < dims.size(); ++i) { |
| 967 | py::object dim = dims[i]; |
| 968 | if (dim.is_none()) { |
| 969 | ptrdiff_t right = (ptrdiff_t)i - dims.size(); |
| 970 | if (inp_shape.is_none()) { |
| 971 | inp_shape = input.attr("shape"); |
| 972 | mgb_assert(!inp_shape.is_none()); |
| 973 | inp_ndim = py::len(inp_shape); |
| 974 | } |
| 975 | if ((ptrdiff_t)inp_ndim + right < 0) { |
| 976 | throw py::value_error("size connot be `None` for new axis"); |
| 977 | } |
| 978 | dim = inp_shape.attr("__getitem__")(right); |
| 979 | dims[i] = dim; |
| 980 | } |
| 981 | if (py::int_::check_(dim)) { |
| 982 | if (dim.cast<long>() < 0) { |
| 983 | throw py::value_error(ssprintf( |
| 984 | "expect shape[%zu] >= 0 or use `None` to auto infer, got " |
| 985 | "%s", |
| 986 | i, py::repr(dims[i]).cast<std::string>().c_str())); |
| 987 | } |
| 988 | } else { |
| 989 | all_imm = false; |
| 990 | } |
| 991 | } |
| 992 | shape = dims; |
| 993 | } else { |
| 994 | all_imm = false; |
| 995 | } |
| 996 | bool fastpath = all_imm && enable_fastpath(input); |
| 997 | if ((!fastpath) && (!is_tensor(shape))) { |
| 998 | shape = _astensor1d_cpp( |
| 999 | shape, py::cast((mgb::DType)dtype::Int32()), input.attr("device"), |
| 1000 | input); |
| 1001 | } |
| 1002 | std::shared_ptr<OpDef> op; |
| 1003 | SmallVector<PyObject*> p(2); |
| 1004 | if (fastpath) { |
| 1005 | std::vector<int32_t> shape_vec; |
| 1006 | for (auto&& dim : dims) { |
| 1007 | shape_vec.push_back(dim.cast<long>()); |
| 1008 | } |
| 1009 | op = Broadcast::make(shape_vec); |
| 1010 | } else { |
| 1011 | op = Broadcast::make(); |
| 1012 | p.push_back(shape.ptr()); |
| 1013 | } |
no test coverage detected