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

Implemented Cholesky on CPU (#1119)

Browse Source
This commit is contained in:
Luca Arnaboldi
2024-05-17 21:31:59 +02:00
committed by GitHub
parent 6a9b584f3d
commit b3ec792380
12 changed files with 239 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
mlx/primitives.h
View File

@@ -2093,4 +2093,20 @@ class Inverse : public UnaryPrimitive {
void eval(const std::vector<array>& inputs, array& output);
};
class Cholesky : public UnaryPrimitive {
public:
explicit Cholesky(Stream stream, bool upper)
: UnaryPrimitive(stream), upper_(upper) {};
void eval_cpu(const std::vector<array>& inputs, array& out) override;
void eval_gpu(const std::vector<array>& inputs, array& out) override;
DEFINE_VMAP()
DEFINE_PRINT(Cholesky)
private:
void eval(const std::vector<array>& inputs, array& output);
bool upper_;
};
} // namespace mlx::core