wehub-resource-sync
109 lines
4.0 KiB
C++
109 lines
4.0 KiB
C++
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include "paddle/phi/api/ext/native_meta_tensor.h"
|
|
#include "paddle/fluid/pybind/data_type_caster.h"
|
|
#include "paddle/fluid/pybind/native_meta_tensor.h"
|
|
#include "paddle/utils/pybind.h"
|
|
#include "pybind11/functional.h"
|
|
#include "pybind11/pybind11.h"
|
|
#include "pybind11/stl.h"
|
|
|
|
namespace paddle::pybind {
|
|
|
|
void BindNativeMetaTensor(py::module* m) {
|
|
py::class_<phi::NativeMetaTensor>(*m, "NativeMetaTensor")
|
|
.def(py::init<>())
|
|
.def(py::init<const phi::NativeMetaTensor&>())
|
|
.def(py::init([](const py::object& dtype, const py::object& shape) {
|
|
DataType dt = DataType::FLOAT32;
|
|
if (!dtype.is_none()) {
|
|
dt = dtype.cast<DataType>();
|
|
}
|
|
std::vector<int64_t> dims;
|
|
if (py::isinstance<py::list>(shape) ||
|
|
py::isinstance<py::tuple>(shape)) {
|
|
dims = shape.cast<std::vector<int64_t>>();
|
|
} else {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"The shape argument must be a list or tuple of integers "
|
|
"or None, but got %s.",
|
|
py::str(shape)));
|
|
}
|
|
return phi::NativeMetaTensor(dt, phi::make_ddim(dims));
|
|
}),
|
|
py::arg("dtype") = py::none(),
|
|
py::arg("shape") = py::list())
|
|
.def(
|
|
"copy",
|
|
[](const phi::NativeMetaTensor& self) {
|
|
return phi::NativeMetaTensor(self);
|
|
},
|
|
"Create a deep copy of this tensor")
|
|
.def(
|
|
"set_shape",
|
|
[](phi::NativeMetaTensor& self, const std::vector<int64_t>& dims) {
|
|
phi::DDim ddim = phi::make_ddim(dims);
|
|
self.set_dims(ddim);
|
|
},
|
|
"Set tensor dimensions")
|
|
.def(
|
|
"set_dtype",
|
|
[](phi::NativeMetaTensor& self, const std::string& dtype_str) {
|
|
self.set_dtype(phi::StringToDataType(dtype_str));
|
|
},
|
|
"Set tensor data type from string")
|
|
.def(
|
|
"set_dtype",
|
|
[](phi::NativeMetaTensor& self, const DataType& dtype) {
|
|
self.set_dtype(dtype);
|
|
},
|
|
"Set tensor data type from DataType object")
|
|
.def_property_readonly(
|
|
"dtype",
|
|
[](const phi::NativeMetaTensor& self) -> DataType {
|
|
return self.dtype();
|
|
},
|
|
"Get tensor data type")
|
|
.def_property_readonly(
|
|
"shape",
|
|
[](const phi::NativeMetaTensor& self) -> std::vector<int64_t> {
|
|
const phi::DDim& dims = self.dims();
|
|
return common::vectorize<int64_t>(dims);
|
|
},
|
|
"Get tensor shape")
|
|
.def("__eq__",
|
|
[](const phi::NativeMetaTensor& self,
|
|
const phi::NativeMetaTensor& other) {
|
|
return self.dtype() == other.dtype() &&
|
|
self.dims() == other.dims();
|
|
})
|
|
.def("__repr__", [](const phi::NativeMetaTensor& self) {
|
|
const phi::DDim& dims = self.dims();
|
|
std::ostringstream shape_ss;
|
|
shape_ss << "[";
|
|
for (int i = 0; i < dims.size(); ++i) {
|
|
if (i > 0) {
|
|
shape_ss << ", ";
|
|
}
|
|
shape_ss << dims[i];
|
|
}
|
|
shape_ss << "]";
|
|
std::string dtype_str = phi::DataTypeToString(self.dtype());
|
|
return "NativeMetaTensor(shape=" + shape_ss.str() +
|
|
", dtype=" + dtype_str + ")";
|
|
});
|
|
}
|
|
} // namespace paddle::pybind
|