zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-20 01:18:12 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Scatter vjp (#394)

Browse Source 
* Add a first scatter vjp
* Implement the scatter_add vjp
* Add array.at to implement user friendly scatters
This commit is contained in:
Angelos Katharopoulos
2024-01-09 13:36:51 -08:00
committed by GitHub
parent e9ca65c939
commit 961435a243
7 changed files with 360 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
python/src/indexing.h
View File

@@ -12,3 +12,27 @@ using namespace mlx::core;
array mlx_get_item(const array& src, const py::object& obj);
void mlx_set_item(array& src, const py::object& obj, const ScalarOrArray& v);
array mlx_add_item(
const array& src,
const py::object& obj,
const ScalarOrArray& v);
array mlx_subtract_item(
const array& src,
const py::object& obj,
const ScalarOrArray& v);
array mlx_multiply_item(
const array& src,
const py::object& obj,
const ScalarOrArray& v);
array mlx_divide_item(
const array& src,
const py::object& obj,
const ScalarOrArray& v);
array mlx_maximum_item(
const array& src,
const py::object& obj,
const ScalarOrArray& v);
array mlx_minimum_item(
const array& src,
const py::object& obj,
const ScalarOrArray& v);