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

Kernel generation (#614)

Browse Source 
Generate reusable element-wise kernels given a computation graph.
This commit is contained in:
Angelos Katharopoulos
2024-02-07 13:15:59 -08:00
committed by GitHub
parent 5fd11c347d
commit 28eac18571
19 changed files with 1302 additions and 459 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
mlx/array.h
View File

@@ -121,6 +121,9 @@ class array {
template <typename T>
T item();
template <typename T>
T item() const;
struct ArrayIterator {
using iterator_category = std::random_access_iterator_tag;
using difference_type = size_t;
@@ -454,6 +457,18 @@ T array::item() {
return *data<T>();
}
template <typename T>
T array::item() const {
if (size() != 1) {
throw std::invalid_argument("item can only be called on arrays of size 1.");
}
if (!is_evaled()) {
throw std::invalid_argument(
"item() const can only be called on evaled arrays");
}
return *data<T>();
}
template <typename It>
void array::init(It src) {
set_data(allocator::malloc(size() * size_of(dtype())));