Add tile op (#438)

This commit is contained in:
Diogo
2024-01-13 02:03:16 -05:00
committed by GitHub
parent 1b71487e1f
commit 2e29d0815b
7 changed files with 105 additions and 3 deletions

View File

@@ -3394,4 +3394,30 @@ void init_ops(py::module_& m) {
Returns:
result (array): The outer product.
)pbdoc");
m.def(
"tile",
[](const array& a, const IntOrVec& reps, StreamOrDevice s) {
if (auto pv = std::get_if<int>(&reps); pv) {
return tile(a, {*pv}, s);
} else {
return tile(a, std::get<std::vector<int>>(reps), s);
}
},
"a"_a,
"reps"_a,
py::pos_only(),
py::kw_only(),
"stream"_a = none,
R"pbdoc(
tile(a: array, reps: Union[int, List[int]], /, *, stream: Union[None, Stream, Device] = None) -> array
Construct an array by repeating ``a`` the number of times given by ``reps``.
Args:
a (array): Input array
reps (int or list(int)): The number of times to repeat ``a`` along each axis.
Returns:
result (array): The tiled array.
)pbdoc");
}