elixir-nx / nx

Multi-dimensional arrays (tensors) and numerical definitions for Elixir
2.55k stars 190 forks source link

Add stack as a callback #1482

Closed josevalim closed 1 month ago

josevalim commented 1 month ago

Generally speaking we want to reduce the number of callbacks but stack causes one additional copy of every tensor and it is used very frequently.

One option is to use optional callbacks, but concatenate/stack are the only operations that require a variable number of arguments, which is hard to include as part of the optional/evaluator system, so it its own callback.

Closes #1441.

josevalim commented 1 month ago

:green_heart: :blue_heart: :purple_heart: :yellow_heart: :heart: