zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Make ops aware of rvalues: astype/as_strided/copy/full (#895)

Browse Source 
When compositing transforms lots of temporary of arrays will be created
and passed to next primitive, and by making ops accepting args by value
we can avoid lots of copies of temporary arrays.
This commit is contained in:
Cheng
2024-03-28 14:35:55 +09:00
committed by GitHub
parent aca7584635
commit bab5386306
3 changed files with 50 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
mlx/ops.h
View File

@@ -41,40 +41,33 @@ array linspace(
StreamOrDevice s = {});
/** Convert an array to the given data type. */
array astype(const array& a, Dtype dtype, StreamOrDevice s = {});
array astype(array a, Dtype dtype, StreamOrDevice s = {});
/** Create a view of an array with the given shape and strides. */
array as_strided(
const array& a,
array a,
std::vector<int> shape,
std::vector<size_t> strides,
size_t offset,
StreamOrDevice s = {});
/** Copy another array. */
array copy(const array& a, StreamOrDevice s = {});
array copy(array a, StreamOrDevice s = {});
/** Fill an array of the given shape with the given value(s). */
array full(
const std::vector<int>& shape,
const array& vals,
std::vector<int> shape,
array vals,
Dtype dtype,
StreamOrDevice s = {});
array full(
const std::vector<int>& shape,
const array& vals,
StreamOrDevice s = {});
array full(std::vector<int> shape, array vals, StreamOrDevice s = {});
template <typename T>
array full(
const std::vector<int>& shape,
T val,
Dtype dtype,
StreamOrDevice s = {}) {
return full(shape, array(val, dtype), to_stream(s));
array full(std::vector<int> shape, T val, Dtype dtype, StreamOrDevice s = {}) {
return full(std::move(shape), array(val, dtype), to_stream(s));
}
template <typename T>
array full(const std::vector<int>& shape, T val, StreamOrDevice s = {}) {
return full(shape, array(val), to_stream(s));
array full(std::vector<int> shape, T val, StreamOrDevice s = {}) {
return full(std::move(shape), array(val), to_stream(s));
}
/** Fill an array of the given shape with zeros. */