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

Dynamic broadcasting for shapeless compile/export (#1722)

Browse Source 
* working towards dynamic broadcast

* shapeless broadcast

* fix build + nits

* use broadcast arrays in quantize matmul

* some cleanup / consistency

* mend

* some comments

* add vjp, jvp for broadcast axes
This commit is contained in:
Awni Hannun
2025-01-09 11:04:24 -08:00
committed by GitHub
parent ec36bfa317
commit 1ccaf80575
20 changed files with 471 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
mlx/compile_impl.h
View File

@@ -27,7 +27,8 @@ bool compile_available_for_device(const Device& device);
std::pair<std::vector<array>, std::vector<array>> compile_trace(
const std::function<std::vector<array>(const std::vector<array>&)>& fun,
const std::vector<array>& inputs);
const std::vector<array>& inputs,
bool shapeless);
using ParentsMap =
std::unordered_map<std::uintptr_t, std::vector<std::pair<array, int>>>;