fix: broadcast vectors for grad calculation #1535
Conversation
nx/lib/nx/defn/grad.ex
Outdated
| Expr.constant(%T{shape: shape, type: {:f, 32}, names: names}, float, []) | ||
| case shape do | ||
| %T{vectorized_axes: [_ | _]} = t -> | ||
| Expr.tensor(Nx.fill(t, float, type: :f32)) |
There was a problem hiding this comment.
We should probably get rid of the names here too.
I also wonder if should move the check for vectorized_axes to constant. Today if someone passes vectorized_axes, Expr.constant is broken. So maybe we should create a tensor if a vectorized axes is given to tensor?
| @@ -338,6 +333,8 @@ defmodule Nx.Defn.Grad do | |||
| @verify_grad Application.compile_env(:nx, :verify_grad, false) | |||
|
|
|||
| defp update_grads(op, args, ans, g, _to_grad_ids, grads) do | |||
| args = revectorize_args(args, ans) | |||
There was a problem hiding this comment.
I would prefer to not revectorized everything on every operation. Is there any chance we could do in broadcast only?
There was a problem hiding this comment.
[unbroadcast(x, Nx.multiply(g, y), ans), unbroadcast(y, Nx.multiply(g, x), ans)]
Lines like this one make it so that g is vectorized and y is unvectorized but has axes with the same name, so things break there.
nx/lib/nx/defn/expr.ex
Outdated
| @@ -1394,6 +1394,11 @@ defmodule Nx.Defn.Expr do | |||
|
|
|||
| ## Constant helpers and related optimizations | |||
|
|
|||
| defp constant(%{vectorized_axes: [_ | _]} = out, number) do | |||
| out = %{out | names: Enum.map(out.names, fn _ -> nil end)} | |||
There was a problem hiding this comment.
I don't think this part should be done here, we should preserve the names. Sorry for the confusion.
| @@ -1343,9 +1334,77 @@ defmodule Nx.Defn.Grad do | |||
|
|
|||
| ## General helpers | |||
|
|
|||
| defp unbroadcast(%{shape: shape} = x, res, %{shape: shape}), do: {x, res} | |||
| defp revectorize_args(args, ans) do | |||
There was a problem hiding this comment.
Let's only apply this if args has more than one element and there are vectorized axes.
Also please test x * sin(y) where y is vectorized.
closes #1533