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

contiguous op / prim (#1612)

Browse Source
This commit is contained in:
Awni Hannun
2024-11-21 19:51:49 -08:00
committed by GitHub
parent 0d5e7716ad
commit dcca0d7477
11 changed files with 104 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
mlx/primitives.cpp
View File

@@ -889,6 +889,32 @@ std::pair<std::vector<array>, std::vector<int>> Conjugate::vmap(
return {{conjugate(inputs[0], stream())}, axes};
}
std::vector<array> Contiguous::vjp(
const std::vector<array>&,
const std::vector<array>& cotangents,
const std::vector<int>&,
const std::vector<array>&) {
return {cotangents};
}
std::vector<array> Contiguous::jvp(
const std::vector<array>&,
const std::vector<array>& tangents,
const std::vector<int>&) {
return {tangents};
}
std::pair<std::vector<array>, std::vector<int>> Contiguous::vmap(
const std::vector<array>& inputs,
const std::vector<int>& axes) {
return {{contiguous(inputs[0], allow_col_major_, stream())}, axes};
}
bool Contiguous::is_equivalent(const Primitive& other) const {
const Contiguous& c_other = static_cast<const Contiguous&>(other);
return allow_col_major_ == c_other.allow_col_major_;
}
array conv_weight_backward_patches(
const array& in,
const array& wt,