elixir-nx / nx

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

Add stack as a callback #1482

Closed josevalim closed 6 months ago

josevalim commented 6 months 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 6 months ago

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