Skip to content

Pass through non-Expr tensors in stop_grad/custom_grad/metadata #1801

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Chapaman wants to merge 2 commits into from
Closed

Pass through non-Expr tensors in stop_grad/custom_grad/metadata #1801

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3f97047
Pass through non-Expr tensors in stop_grad/custom_grad/metadata
Chapaman Aug 7, 2026
eb27b7d
Keep stop_grad/custom_grad wrapping; gate only in Expr.metadata
Chapaman Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 22 additions & 9 deletions nx/lib/nx/defn/expr.ex
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,20 +99,33 @@ defmodule Nx.Defn.Expr do
inspection.
"""
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])
if expr_container?(expr) do
case to_container_expr(expr) do
%{data: %{context: context}} = res ->
expr(res, context, :metadata, [Nx.devectorize(expr), metadata])

t when is_tuple(t) ->
context = elem(t, 0).data.context
t when is_tuple(t) ->
context = elem(t, 0).data.context

tuple(
expr(tuple_out(tuple_size(t)), context, :metadata, [Nx.devectorize(expr), metadata]),
Tuple.to_list(t)
)
tuple(
expr(tuple_out(tuple_size(t)), context, :metadata, [Nx.devectorize(expr), metadata]),
Tuple.to_list(t)
)
end
else
# BinaryBackend.block/4 (and similar) may re-run callbacks that call
# stop_grad/custom_grad/metadata; concrete tensors must pass through.
expr
end
end

defp expr_container?(container) do
Composite.reduce(container, true, fn
%T{data: %Expr{}}, true -> true
_, _ -> false
end)
end

@doc """
Creates a tuple with elements in `list` that points to tuple
expression `expr`.
Expand Down
Loading