Enhance symbolic arithmetics

[shino16]

Closes #2736. The main purpose is to implement math.ceil, which is needed for supporting clang.arange (ref).

To test it against both int and float, I added a logic to test the type of NumberProxy's underlying values in prologue.

I took a quick side step to implement all ops tested in test_elementwise.py. This involved the following patches.

Rounding ops: ceil, floor, trunc, round

• Registered their NumberProxy versions to pythonex, backed with math.ceil/floor/trunc, builtins.round.
• Patched them to compute in the promoted type (not in int).

Bitwise ops: bitwise_not, bitwise_left_shift, bitwise_right_shift

• Registered their NumberProxy version to pythonex, backed with operator.inv/lshift/rshift
• Routed e.g. NumberProxy.__lshift__ to the corresponding clangop.

floor_divide

• Patched where, which floor_divide depends on, to support all-NumberProxy inputs

pos (as opposed to neg)

• Thunder has no implementation of pos. I changed NumberProxy.__pos__ and TensorProxy.__pos__ to just return self, which aligns with Python and PyTorch (try for x in (1000, torch.randn(3)): assert +x is x).

[kiya00]

LGTM, thank you @shino16

[Copilot]

Pull request overview

This PR enhances symbolic arithmetic support by implementing rounding operations (ceil, floor, trunc, round), bitwise shift operations, and improving where to support all-NumberProxy inputs. The main motivation is to enable math.ceil support for clang.arange. The implementation includes proper type promotion handling to ensure rounding operations return integers for number-only inputs while preserving dtypes for tensor inputs.

Key changes:

• Added rounding and bitwise operations with proper type promotion semantics
• Extended prologue to test both int and float types of NumberProxy values via symbolic value caching
• Changed __pos__ methods to return self for both NumberProxy and TensorProxy, aligning with Python/PyTorch behavior

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
thunder/tests/test_elementwise.py	Refactored tests to use parametrization for binary/unary ops and added symbolic value caching tests to verify type-based cache behavior
thunder/executors/pythonex.py	Registered rounding and bitwise operations with appropriate backing functions, added where implementation for all-NumberProxy inputs
thunder/executors/nvfuserex_impl.py	Added decorator to ensure rounding operations return int type for number-only inputs in nvFuser execution
thunder/core/utils.py	Changed type promotion for NUMBER_TO_INT kind to use promoted type instead of always returning int
thunder/core/proxies.py	Updated NumberProxy and TensorProxy __pos__ to return self, implemented bitwise shift operations, fixed type annotations
thunder/core/prims.py	Updated rounding operations to use INT_FOR_NUMBER output dtype kind, enabled where for all-NumberProxy inputs
thunder/core/jit_ext.py	Added type constraint check for symbolic value caching
thunder/clang/init.py	Updated rounding operations to use NUMBER_TO_INT type promotion, added method_name to bitwise shift operations

Comments suppressed due to low confidence (1)

thunder/core/prims.py:1

• Lines 1950-1952 duplicate the description of 'operations, like abs, map complex numbers to floats (COMPLEX_TO_FLOAT)'. Remove the duplicate sentence on lines 1951-1952 and update line 1953 to reference 'four behaviors' without the redundant explanation.

from enum import auto, Enum

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[kshitij12345]

LGTM, just few minor comments, thanks @shino16

[KaelanDt]

looks great, thanks @shino16