frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[ez] add docblock to cast_symbool_to_symint_guardless (#154400)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/154400
Approved by: https://github.com/laithsakka
This commit is contained in:
bobrenjc93
2025-05-28 12:46:06 -07:00
committed by PyTorch MergeBot
parent f6dcc45c44
commit e7318b863d
1 changed files with 14 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
torch/fx/experimental/symbolic_shapes.py
View File

@ -2385,6 +2385,20 @@ def _sympy_cast_symbool_to_symint_guardless(x: SympyBoolean) -> sympy.Expr:
def cast_symbool_to_symint_guardless(
symbool: Union[bool, torch.SymBool]
) -> Union[int, torch.SymInt]:
"""
Converts a SymBool or bool to a SymInt or int without introducing guards.
This function maps True to 1 and False to 0, preserving the symbolic nature
of the input when it's a SymBool. Unlike regular casting which might introduce
guards, this function performs the conversion without adding any guards.
Args:
symbool: A boolean value, either a concrete bool or symbolic SymBool
Returns:
The corresponding integer value (1 for True, 0 for False) as either
a concrete int or symbolic SymInt
"""
if isinstance(symbool, bool):
return 1 if symbool else 0
int_sym = _sympy_cast_symbool_to_symint_guardless(symbool.node.expr)