Skip to content

fix(nx): keep a reshaped float constant a constant - #1837

Merged
polvalente merged 1 commit into
elixir-nx:mainfrom
cash-mckeeman:fold-reshaped-constants
Sep 7, 2026
Merged

polvalente merged 1 commit into
elixir-nx:mainfrom
cash-mckeeman:fold-reshaped-constants

Conversation

@cash-mckeeman

@cash-mckeeman cash-mckeeman commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

broadcast/4 folds a constant it's handed; reshape/2 didn't, so a literal reaching an op through a reshape arrived as a plain node.

defn pick(pred, t), do: Nx.select(pred, 0.1, t)

pick(Nx.vectorize(pred, :a), Nx.vectorize(t, :a))

Vectorized operands take that route. Folding also drops the reshape and broadcast nodes wrapping the literal, so two tests asserting inspected expressions now assert smaller graphs.

Split out of #1835.

🤖 Generated with Claude Code

Nx.Defn.Expr.broadcast/4 folds a constant it is given; reshape/2 did not,
so a literal reaching an op through a reshape arrived as a plain node and
missed the constant promotion. Vectorized operands take that route:

    defn pick(pred, t), do: Nx.select(pred, 0.1, t)
    pick(Nx.vectorize(pred, :a), Nx.vectorize(t_f64, :a))

answered 0.10000000149011612 instead of 0.1.

Folding it drops the reshape and broadcast nodes that wrapped the literal,
so two tests asserting inspected expressions were updated to the smaller
graphs.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Eu3cbRr9M6j5qj4Ps4Hgcg
@cash-mckeeman cash-mckeeman changed the title fix: keep a reshaped float constant a constant fix(nx): keep a reshaped float constant a constant Sep 7, 2026
@polvalente
polvalente merged commit 29e99e6 into elixir-nx:main Sep 7, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants