Skip to content

fix: Compile non-finite constants in JIT programs - #153

Merged
polvalente merged 1 commit into
elixir-nx:mainfrom
georgeguimaraes:fix/infinite-constants
Oct 23, 2026
Merged

polvalente merged 1 commit into
elixir-nx:mainfrom
georgeguimaraes:fix/infinite-constants

Conversation

@georgeguimaraes

@georgeguimaraes georgeguimaraes commented Oct 19, 2026 •

Copy link
Copy Markdown
Member

A compiled function that contains an infinite or NaN constant fails when it compiles:

x = Nx.tensor([1.0, -2.0], backend: EMLX.Backend)
Nx.Defn.jit(fn x -> Nx.select(Nx.greater(x, 0), Nx.Constants.neg_infinity(:f32), x) end, compiler: EMLX).(x)
# ** (ArithmeticError) bad argument in arithmetic expression
#     lib/emlx/native/expr.ex:3332: anonymous fn/1 in EMLX.Native.Expr.to_native/1

Nx keeps the special values as atoms: :infinity, :neg_infinity and :nan. to_native/1 multiplied each constant value by 1.0 to make the NIF wire format, and this raised for an atom.

Erlang floats cannot hold NaN or infinity, so the atoms now go to the NIF as they are. A fine::Decoder for std::tuple<double, mlx::core::Dtype> in compiler.hpp maps them to the C double values infinity, -infinity and quiet NaN. Other constants do not change.

The new tests in test/emlx/compiler_test.exs cover the three atoms in f32, f16 and bf16. They check the values of the results, including NaN propagation. The full suite passes: 2697 tests, 7 excluded.

@georgeguimaraes georgeguimaraes changed the title fix: compile non-finite constants in JIT programs fix: Compile non-finite constants in JIT programs Oct 19, 2026
Nx keeps the IEEE 754 special values as atoms: :infinity, :neg_infinity and :nan. In compiled programs, to_native/1 multiplied each constant value by 1.0 to make the NIF wire format, and this raised ArithmeticError for an atom.

Erlang floats cannot hold NaN or infinity, so the atoms now go to the NIF as they are. A decoder in compiler.hpp maps them to the C double values infinity, -infinity and quiet NaN.

The new tests cover the three atoms in f32, f16 and bf16, and check the values of the results.
@polvalente
polvalente merged commit 92a3143 into elixir-nx:main Oct 23, 2026
10 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