fix(nx): read float literals at the precision of the expression around them by cash-mckeeman · Pull Request #1835 · elixir-nx/nx

cash-mckeeman · 2026-09-07T18:43:33Z

A bare float literal is annotated {:f, 32}, and that annotation is the precision it gets read back at. Binary operations promote it; nothing else does.

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

pick(Nx.tensor(1, type: :u8), Nx.tensor([1.0], type: :f64))
#=> #Nx.Tensor<f64[1] [0.10000000149011612]>

Same for clip, pad, stack, reduce, dot, and comparisons, which is how tolerances get written:

defn converged?(t), do: t <= 1.0e-8

converged?(Nx.tensor(9.999999969612645e-9, type: :f64))
#=> #Nx.Tensor<u8 0>   # the probe is below 1.0e-8 and above its f32 round

The promotion lived in binary_expr/5. Every node is built through expr/4, so it moves there and the two calls in binary_expr/5 become redundant. Comparisons answer in {:u, 8}, which can't say what precision to read the constant at, so the operands do.

Second commit: reshape/2 didn't fold a constant the way broadcast/4 does, so vectorized operands lost the promotion on the way through. Folding it also drops the wrapper nodes, so two tests asserting inspected expressions now assert smaller graphs.

One question — eager calls are untouched, so Nx.multiply(Nx.f64(2), 0.7) still gives 1.399999976158142. Fixing that reaches into Nx.Shared.binary_type/2. Want it?

🤖 Generated with Claude Code

polvalente · 2026-09-07T18:53:30Z

+  defp collect_float_types(args) do
+    Enum.flat_map(args, fn
+      %T{type: type} -> if Nx.Type.float?(type), do: [type], else: []
+      list when is_list(list) -> collect_float_types(list)


I don't think we should be traversing nested lists here. These are generally option arguments and not tensor arguments. Do we have a test that fails if we remove this?

Removed from collect_float_types — comparisons are the only caller and they're binary, so it never fired.

Kept in upcast_float_constants, where concatenate and stack pass their tensors in a list; without it Nx.stack([t_f64, 0.1]) regresses. Nothing covered that, so I added a test.

polvalente · 2026-09-07T18:55:23Z

+    if Nx.Type.float?(type) do
+      type
+    else
+      args
+      |> collect_float_types()
+      |> case do
+        [] -> type
+        [first | rest] -> Enum.reduce(rest, first, &Nx.Type.merge/2)
+      end
+    end


I think this function should just use tensor and args directly through the else block. Otherwise this might end up staying at a lower precision.

General question for the PR: what do we want to happen if we get :c64 and :c128 vs a float constant? What happens if an :f64 constant encounters a :c64 tensor?

I think in other parts of the repo, the constant adapts to the tensor (Nx.Type.merge_number/2) with the exception being where the tensor cannot hold the constant type - the fallback in that function catches cases where a tensor with integer type widens to {:f, 32}.

That feels right to me in general and in this PR. I tried to faithfully implement according to that principle.

Digging through this a bit, though, I found that when something calls to_tensor that the Nx.Type.merge/2 function is called. And that function runs against the {:f, 32} placeholder given by Nx.Type.infer/1.

All code paths get that placeholder, but it's to_tensor that is destructive in the sense that it rounds the input right after, I believe. Going through defn, the type is just stamped and the input can be recovered.

Maybe that's the real fix needed; making those two code paths agree.

polvalente · 2026-09-07T20:43:38Z

+  # op answering in a tuple has no type to contribute.
+  defp constant_read_type(%T{type: type}, args) do
+    types = operand_float_types(args)
+    types = if match?({:tuple, _}, type), do: types, else: [type | types]


If a tuple can appear there, I think this would justify recursion into the tuple arguments

polvalente · 2026-09-07T20:49:08Z

+    case types do
+      [] -> type
+      [first | rest] -> Enum.reduce(rest, first, &Nx.Type.merge/2)
+    end


I think if we include the tuple argument types, we can always use the reduce here

polvalente · 2026-09-07T20:49:59Z

+  defp operand_float_types(args) do
+    Enum.flat_map(args, fn
+      %T{data: %Expr{op: :constant}} -> []
+      %T{type: type} -> if Nx.Type.float?(type), do: [type], else: []
+      _ -> []
+    end)


I think we should return all types for simplicity

polvalente · 2026-09-07T20:50:22Z

+  defp upcast_float_constants(args, type) do
+    Enum.map(args, fn
+      %T{data: %Expr{op: :constant}} = t -> maybe_upcast_float_constant(t, type)
+      list when is_list(list) -> upcast_float_constants(list, type)


I still think traversing lists here is not the way to go

I'll take another pass

A bare float literal is annotated {:f, 32}, and that annotation is the precision it gets read back at. The promotion added for binary operations lived in binary_expr/5, so every other op still read the f32 round: select, clip, pad, stack, reduce, dot, and comparisons. Every expression node is built through expr/4, so the promotion moves there and the two calls in binary_expr/5 become redundant. Comparisons answer in {:u, 8}, which cannot say what precision the constant is read at, so there the operands do. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Eu3cbRr9M6j5qj4Ps4Hgcg

The output type alone cannot say what precision a constant should be read at. Comparisons answer in {:u, 8}, and an op answering in a tuple has no type to contribute, so the target now merges the output type with the other operands. The literal's own {:f, 32} is the annotation being corrected, so it is left out of that merge, which keeps a narrow operand narrow instead of widening it through f32. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Eu3cbRr9M6j5qj4Ps4Hgcg

Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com>

Walking every list argument in expr/4 reached option lists as well as tensor ones. Only concatenate and stack pass their tensors in a list, so they promote their own, and the generic walk goes away. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Eu3cbRr9M6j5qj4Ps4Hgcg

A tuple output type carries no precision, so it was being filtered out of the list it had just been added to. Ask whether the output type is a float instead: a float output is already the merge of everything the node was given, and every other output type defers to the operands. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Eu3cbRr9M6j5qj4Ps4Hgcg

expr/4 computed a read type for every node it built, and the answer went nowhere on the ones carrying no literal. Fold the check and the promotion into one function and ask it once; concatenate and stack, which were passing the output type by hand, go through the same door. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Eu3cbRr9M6j5qj4Ps4Hgcg

cash-mckeeman · 2026-09-09T15:59:41Z

Simplified to avoid {:tuple, size} - it's what elem/2 takes as an argument, and what while and cond return. Nx.Type.float?/1 is already false for it.

Also folded the check and the promotion into one function, so expr/4 only asks for a literal's precision when there's a literal.

polvalente reviewed Sep 7, 2026

View reviewed changes

Comment thread nx/lib/nx/defn/expr.ex

cash-mckeeman changed the title ~~fix: read float literals at the precision of the expression around them~~ fix(nx): read float literals at the precision of the expression around them Sep 7, 2026

cash-mckeeman force-pushed the fix-float-literal-precision branch from 1f38814 to f8977ce Compare September 7, 2026 19:10

cash-mckeeman mentioned this pull request Sep 7, 2026

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

Merged

polvalente reviewed Sep 7, 2026

View reviewed changes

Comment thread nx/lib/nx/defn/expr.ex Outdated

polvalente reviewed Sep 7, 2026

View reviewed changes

cash-mckeeman and others added 4 commits September 7, 2026 19:37

Update nx/lib/nx/defn/expr.ex

f224b75

Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com>

cash-mckeeman force-pushed the fix-float-literal-precision branch from c80d00f to 777d459 Compare September 8, 2026 00:38

cash-mckeeman and others added 2 commits September 9, 2026 09:56

polvalente approved these changes Sep 10, 2026

View reviewed changes

polvalente merged commit 4fba859 into elixir-nx:main Sep 10, 2026
9 checks passed

Conversation

cash-mckeeman commented Sep 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cash-mckeeman commented Sep 9, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

cash-mckeeman commented Sep 7, 2026 •

edited

Loading