Reproduce
x = Nx.iota({2, 3, 3}, type: :f32) |> Nx.vectorize(:a)
y = Nx.iota({2, 3, 3}, type: :f32) |> Nx.vectorize(:row)
Nx.Defn.grad(x, fn x -> Nx.sum(Nx.dot(x, [1], y, [0])) end)
# ** (ArgumentError) expected length of axes (8) to match rank of shape (6)
Root cause: unbroadcast in the dot gradient doesn't account for different vectorized axes on operands.
Simple Nx.dot(x, y) without contract axes works with mixed vectorized axes. Only the contracted form
fails.
Found during #1533 / PR #1697 work. Skipped test in PR documents the expected behavior.
Reproduce
Root cause: unbroadcast in the dot gradient doesn't account for different vectorized axes on operands.
Simple Nx.dot(x, y) without contract axes works with mixed vectorized axes. Only the contracted form
fails.
Found during #1533 / PR #1697 work. Skipped test in PR documents the expected behavior.