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

shapeless slice update and broadcast when possible (#1727)

Browse Source
This commit is contained in:
Awni Hannun
2024-12-23 11:25:15 -08:00
committed by GitHub
parent 0308e9af71
commit ebfe64b92d
6 changed files with 43 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
mlx/primitives.h
View File

@@ -533,6 +533,8 @@ class Broadcast : public UnaryPrimitive {
DEFINE_PRINT(Broadcast)
bool is_equivalent(const Primitive& other) const override;
std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
private:
Shape shape_;
@@ -1943,6 +1945,7 @@ class SliceUpdate : public UnaryPrimitive {
DEFINE_GRADS()
DEFINE_PRINT(SliceUpdate)
bool is_equivalent(const Primitive& other) const override;
DEFINE_INPUT_OUTPUT_SHAPE()
private:
Shape start_indices_;