From 52c5d8c087f6a2ee91b807467860eb3e96bb6267 Mon Sep 17 00:00:00 2001
From: Laurent Mazare <laurent.mazare@gmail.com>
Date: Sat, 22 Jul 2023 08:21:28 +0200
Subject: Add the gather op. (#219)

* Start adding gather.

* Gather cpu implementation + use in simple training.

* Add scatter_add for the gradient of gather.

* Simple cpu implementation of scatter_add.

* Use gather in the simple-training backprop.
---
 candle-core/src/op.rs | 2 ++
 1 file changed, 2 insertions(+)

(limited to 'candle-core/src/op.rs')

diff --git a/candle-core/src/op.rs b/candle-core/src/op.rs
index d36aa301..de5094bd 100644
--- a/candle-core/src/op.rs
+++ b/candle-core/src/op.rs
@@ -66,6 +66,8 @@ pub(crate) enum Op {
     Reduce(Tensor, ReduceOp, Vec<usize>),
     Matmul(Tensor, Tensor),
     Embedding(Tensor, Tensor),
+    Gather(Tensor, Tensor, usize),
+    ScatterAdd(Tensor, Tensor, Tensor, usize),
     IndexSelect(Tensor, Tensor, usize),
     IndexAdd(Tensor, Tensor, Tensor, usize),
     WhereCond(Tensor, Tensor, Tensor),
-- 
cgit v1.2.3