I made a broadcasting-based implementation that works with batch dimensions: https://gist.github.com/yulkang/4a597bcc5e9ccf8c7291f8ecb776382d
1 Like
I made a broadcasting-based implementation that works with batch dimensions: https://gist.github.com/yulkang/4a597bcc5e9ccf8c7291f8ecb776382d