fix: broadcast vectors for grad calculation

nx/lib/nx.ex

@@ -5420,9 +5420,13 @@ defmodule Nx do
       {_, [], 0} ->
         fun.(left, right)
 
-      {[devec_left, devec_right], canonical_vectorized_axes, _offset} ->
-        devec_left
-        |> fun.(devec_right)
+      {[devec_left, devec_right], canonical_vectorized_axes, offset} ->
+        leading_names = Keyword.keys(canonical_vectorized_axes)
+        l = %{devec_left | names: leading_names ++ Enum.drop(devec_left.names, offset)}
+        r = %{devec_right | names: leading_names ++ Enum.drop(devec_right.names, offset)}
+
+        l
+        |> fun.(r)
         |> vectorize(canonical_vectorized_axes)
     end
   end

nx/lib/nx/defn/expr.ex

@@ -94,13 +94,13 @@ defmodule Nx.Defn.Expr do
   def metadata(expr, metadata) when is_map(metadata) do
     case to_container_expr(expr) do
       %{data: %{context: context}} = res ->
-        expr(res, context, :metadata, [Nx.devectorize(expr), metadata])
+        expr(res, context, :metadata, [expr, metadata])
 
       t when is_tuple(t) ->
         context = elem(t, 0).data.context
 
         tuple(
-          expr(tuple_out(tuple_size(t)), context, :metadata, [Nx.devectorize(expr), metadata]),
+          expr(tuple_out(tuple_size(t)), context, :metadata, [expr, metadata]),
           Tuple.to_list(t)
         )
     end
@@ -1661,11 +1661,11 @@ defmodule Nx.Defn.Expr do
 
   defp counter_to_name(counter), do: [?a + counter]
 
-  defp to_type_shape(%{type: type, shape: shape}) do
-    brackets =
-      shape
-      |> Tuple.to_list()
-      |> Enum.map(&[?[, Integer.to_string(&1), ?]])
+  defp to_type_shape(%{vectorized_axes: vectorized_axes, type: type, shape: shape}) do
+    axes =
+      Keyword.values(vectorized_axes) ++ Tuple.to_list(shape)
+
+    brackets = Enum.map(axes, &[?[, Integer.to_string(&1), ?]])
 
     IO.iodata_to_binary([Nx.Type.to_string(type) | brackets])
   end

nx/lib/nx/defn/grad.ex

@@ -6,11 +6,10 @@ defmodule Nx.Defn.Grad do
 
   def transform(to_grad, fun, transform) do
     {to_grad, ids} =
-      Composite.traverse(to_grad, %{}, fn to_grad, ids ->
-        to_grad =
-          Expr.metadata(to_grad, %{__MODULE__ => :to_grad})
-
-        {to_grad, Map.put(ids, to_grad.data.id, :stop)}
+      Composite.traverse(to_grad, %{}, fn node, ids ->
+        node = Expr.metadata(node, %{__MODULE__ => :to_grad})
+        ids = Map.put(ids, node.data.id, :stop)
+        {node, ids}
       end)
 
     # Collect all IDs in the function environment and mark
@@ -19,36 +18,32 @@ defmodule Nx.Defn.Grad do
     {:env, env} = Function.info(fun, :env)
     ids = stop_grads(env, ids)
 
-    # save vectorized axes before devectorizing
-    expr = to_grad |> fun.()
+    expr = fun.(to_grad)
+
+    transformed_expr =
+      expr |> transform.() |> validate_expr!()
 
-    transformed_expr = transform.(expr) |> validate_expr!() |> Nx.devectorize(keep_names: false)
     {parents, nodes} = parents_tree(transformed_expr, ids)
 
     to_grad_ids = {to_grad, ids}
     grads = %{transformed_expr.data.id => [constant(1.0, transformed_expr)]}
 
     {graded, _} =
-      Composite.traverse(
-        to_grad,
-        {nodes, grads},
-        fn %{vectorized_axes: vectorized_axes} = node, acc ->
-          node
-          |> Nx.devectorize(keep_names: false)
-          |> to_grad(to_grad_ids, parents, acc)
-          |> then(fn {node, acc} ->
-            {Nx.vectorize(node, vectorized_axes), acc}
-          end)
-        end
-      )
+      Composite.traverse(to_grad, {nodes, grads}, &to_grad(&1, to_grad_ids, parents, &2))
 
     {expr, graded}
   end
 
   defp constant(float, shape) do
-    shape = Nx.shape(shape)
-    names = List.duplicate(nil, tuple_size(shape))
-    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))
+
+      t ->
+        shape = Nx.shape(t)
+        names = List.duplicate(nil, tuple_size(shape))
+        Expr.constant(%T{shape: shape, type: {:f, 32}, names: names}, float, [])
+    end
   end
 
   defp validate_expr!(%T{data: %Expr{}} = expr) do
@@ -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)
+
     pairs = grad(op, args, ans, g)
 
     if @verify_grad do
@@ -1343,9 +1340,66 @@ defmodule Nx.Defn.Grad do
 
   ## General helpers
 
-  defp unbroadcast(%{shape: shape} = x, res, %{shape: shape}), do: {x, res}
+  defp revectorize_args(args, ans) do
+    names_ans =
+      Enum.with_index(Keyword.keys(ans.vectorized_axes) ++ ans.names, fn name, idx ->
+        if(name, do: name, else: {:ans, idx})
+      end)
+
+    for arg <- args do
+      case arg do
+        %T{names: names} ->
+          names = Enum.with_index(names, fn name, idx -> if(name, do: {name, idx}) end)
+
+          vectorized_axes =
+            names
+            |> Enum.reduce([], fn
+              nil, acc ->
+                acc
+
+              {name, _idx}, acc ->
+                if name in names_ans do
+                  [name | acc]
+                else
+                  acc
+                end
+            end)
+            |> Enum.reverse()
+
+          Nx.vectorize(arg, vectorized_axes)
+
+        arg ->
+          arg
+      end
+    end
+  end
+
+  defp unbroadcast(x, res, ans) do
+    vectorized_axes_x = Keyword.keys(x.vectorized_axes)
+    vectorized_axes_ans = Keyword.keys(ans.vectorized_axes)
+
+    permutation =
+      vectorized_axes_ans
+      |> Enum.with_index()
+      |> Enum.sort_by(fn {axis, _idx} -> axis in vectorized_axes_x end)
+      |> Enum.map(fn {_, idx} -> idx end)
+
+    num_vectorized_axes = length(permutation)
+
+    inner_axes = Enum.to_list(num_vectorized_axes..(num_vectorized_axes + Nx.rank(res) - 1)//1)
+
+    res =
+      res
+      |> Nx.devectorize()
+      |> Nx.transpose(axes: permutation ++ inner_axes)
+      |> Nx.vectorize(vectorized_axes_x)
+
+    grad_broadcast(x, res, ans)
+  end
+
+  defp grad_broadcast(%{shape: shape} = x, res, %{shape: shape}), do: {x, res}
 
-  defp unbroadcast(%{shape: shape} = x, res, %{shape: new_shape}) do
+  defp grad_broadcast(%{shape: shape} = x, res, %{shape: new_shape}) do
     axes = Nx.Shape.broadcast_axes(shape, new_shape)
     {x, grad_broadcast(x, new_shape, axes, res)}
   end

nx/test/nx/defn/grad_test.exs

@@ -4238,21 +4238,70 @@ defmodule Nx.Defn.GradTest do
   end
 
   describe "vectorization" do
-    test "supports vectorization" do
+    test "supports combination of vectorized and non-vectorized tensors" do
+      x = Nx.tensor([[1, 2, 3], [4, 5, 6]]) |> Nx.vectorize(:x)
+      y = 1
+
+      grad = Nx.Defn.grad(y, fn y -> Nx.add(x, y) end)
+
+      assert grad == Nx.tensor([3.0, 3.0]) |> Nx.vectorize([:x])
+    end
+
+    test "supports heterogenous vectorization combinations" do
       x = Nx.tensor([[1, 2, 3], [4, 5, 6]])
       y = Nx.tensor([10, 20])
 
       # first case: y is vectorized scalar, x is vectorized vectors, different vectorized axis names
       # expected result: equivalent to fully broadcasting one tensor onto the other
       x_vec = Nx.vectorize(x, :x)
       y_vec = Nx.vectorize(y, :y)
-      {grad_x_vec, grad_y_vec} = Nx.Defn.grad({x_vec, y_vec}, fn {a, b} -> Nx.multiply(a, b) end)
 
+      grad_fun = fn x, y ->
+        Nx.Defn.grad({x, y}, fn {a, b} -> Nx.multiply(a, b) end)
+      end
+
+      {grad_x_vec, grad_y_vec} = grad_fun.(x_vec, y_vec)
+
+      # Explicit assertion on the results
       assert grad_x_vec ==
-               Nx.tensor([[30.0, 30.0, 30.0], [30.0, 30.0, 30.0]])
-               |> Nx.vectorize(x_vec.vectorized_axes)
+               Nx.tensor([
+                 [
+                   [10.0, 10.0, 10.0],
+                   [20.0, 20.0, 20.0]
+                 ],
+                 [
+                   [10.0, 10.0, 10.0],
+                   [20.0, 20.0, 20.0]
+                 ]
+               ])
+               |> Nx.vectorize([:x, :y])
+
+      assert grad_y_vec ==
+               Nx.tensor([
+                 [6.0, 6.0],
+                 [15.0, 15.0]
+               ])
+               |> Nx.vectorize([:x, :y])
+
+      # Conceptual assertion: the result should be equivalent to calling Nx.Defn.grad with
+      # each cross-entry of the combined vectors [(x0, y0), (x0, y1), (x1, y0), (x1, y1)]
 
-      assert grad_y_vec == Nx.tensor([21.0, 21.0]) |> Nx.vectorize(y_vec.vectorized_axes)
+      {x0y0_wrt_x, x0y0_wrt_y} = grad_fun.(x[0], y[0])
+      {x0y1_wrt_x, x0y1_wrt_y} = grad_fun.(x[0], y[1])
+      {x1y0_wrt_x, x1y0_wrt_y} = grad_fun.(x[1], y[0])
+      {x1y1_wrt_x, x1y1_wrt_y} = grad_fun.(x[1], y[1])
+
+      assert grad_x_vec ==
+               [x0y0_wrt_x, x0y1_wrt_x, x1y0_wrt_x, x1y1_wrt_x]
+               |> Nx.stack()
+               |> Nx.reshape({2, 2, 3})
+               |> Nx.vectorize([:x, :y])
+
+      assert grad_y_vec ==
+               [x0y0_wrt_y, x0y1_wrt_y, x1y0_wrt_y, x1y1_wrt_y]
+               |> Nx.stack()
+               |> Nx.reshape({2, 2})
+               |> Nx.vectorize([:x, :y])
 
       # second case: y is vectorized scalar, x is vectorized vectors, same vectorized axis name
       # expected result: equivalent to "row-wise" broadcasting