You’re right! I just brainfarted for a bit because Nx.broadcast(w, x) doesn’t work directly.
You can do something like:
iex(4)> t1.shape
{2, 3}
iex(5)> t2.shape
{5, 3, 4}
iex(6)> Nx.broadcast(t1, Tuple.insert_at(t1.shape, 0, Nx.axis_size(t2, 0)))
#Nx.Tensor<
s64[5][2][3]
...
>
Note that inside defn you’ll probably need to define a separate deftransformp for applying this broadcast.
Vectorization will come with its own separate set of difficulties if you want to support a function that receives a maybe vectorized tensor, though, so each approach has its own trade-offs.
