mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-09 15:08:02 -05:00
Add Tensor.scatter_reduce (#8947)
* pytorch scatter -> scatter_reduce * WIP scatter_reduce implementation * _pre_scatter return type hint * split out src, mask to satisfy linter * Add src cast back in * dict of lambdas instead of ifs * sum and prod reduction ops with include_self * add reduce arg error message * add amax and amin reduction ops * Fix include_self for higher dims * Simplify * Simplify amax and amin too * Pull include_self logic out into _inv_mask function * reduce arg cannot be None for scatter_reduce * Fix self-mask issue * Add mean reduce op * Add tests * any() not needed here * remove comment * End support for Tensor src with reduce arg in tinygrad scatter * Process index, dim inside actual functions * Add scatter_reduce to onnx * Add excluded onnx ScatterElements reduction tests back in * Save 2 lines on the mask helpers * Update docs * Add include_self=False tests * cleanup * Remove unneeded helper function --------- Co-authored-by: chenyu <chenyu@fastmail.com>
This commit is contained in:
@@ -32,6 +32,7 @@
|
||||
::: tinygrad.Tensor.tril
|
||||
::: tinygrad.Tensor.interpolate
|
||||
::: tinygrad.Tensor.scatter
|
||||
::: tinygrad.Tensor.scatter_reduce
|
||||
|
||||
## Neural Network (functional)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user