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

Add view op (#1179)

Browse Source 
* add view primitive

* nit

* fix view
This commit is contained in:
Awni Hannun
2024-06-04 08:05:27 -07:00
committed by GitHub
parent 81def6ac76
commit ea9090bbc4
14 changed files with 202 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
mlx/primitives.h
View File

@@ -2021,6 +2021,22 @@ class Uniform : public UnaryPrimitive {
void eval(const std::vector<array>& inputs, array& out);
};
class View : public UnaryPrimitive {
public:
explicit View(Stream stream, Dtype dtype)
: UnaryPrimitive(stream), dtype_(dtype) {};
void eval_cpu(const std::vector<array>& inputs, array& out) override;
void eval_gpu(const std::vector<array>& inputs, array& out) override;
DEFINE_VMAP()
void print(std::ostream& os) override;
bool is_equivalent(const Primitive& other) const override;
private:
Dtype dtype_;
};
class Transpose : public UnaryPrimitive {
public:
explicit Transpose(Stream stream, const std::vector<int>& axes)