Fix error with int+SymBool

[jansel]

Stack from ghstack (oldest at bottom):

• Remove size asserts from fx_insert_profiling #114830
• -> Fix error with int+SymBool #114828

Fixes #104797

  File "/home/jansel/pytorch/torch/_dynamo/utils.py", line 1486, in <lambda>
    lambda: run_node(tx.output, node, args, kwargs, nnmodule)
  File "/home/jansel/pytorch/torch/_dynamo/utils.py", line 1591, in run_node
    raise RuntimeError(fn_str + str(e)).with_traceback(e.__traceback__) from e
  File "/home/jansel/pytorch/torch/_dynamo/utils.py", line 1570, in run_node
    return node.target(*args, **kwargs)
  File "/home/jansel/conda/envs/pytorch/lib/python3.10/site-packages/einops/packing.py", line 153, in unpack
    n_unknown_composed_axes = sum(x == -1 for x in lengths_of_composed_axes)
torch._dynamo.exc.TorchRuntimeError: Failed running call_function <function unpack at 0x7f644b962710>(*(FakeTensor(..., device='cuda:0', size=(1, s0*s1, 128)), [(s0, s1)], 'b * c'), **{}):
unsupported operand type(s) for +: 'int' and 'SymBool'

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/114828

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 2f4f346 with merge base f128616 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / linux-focal-py3_8-clang9-xla / test (xla, 1, 1, linux.12xlarge) (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[lezcano]

calling int(x) guards on the value of x. It would be better to implement a casting method similar to sym_float that casts a bool to an int, similar to the method sym_float.

[jansel]

1. SymPy doesn't seem to support this
2. All the usage of this pattern I found are immediately fed into conditionals, so we need to guard on it anyway

[lezcano]

Regarding 1, we can work around it by using Where

pytorch/torch/utils/_sympy/functions.py

Lines 151 to 163 in e6b3a8c

	class Where(sympy.Function):
	"""
	Good ol' ternary operator
	"""
	nargs = (3,)
	@classmethod
	def eval(cls, c, p, q):
	if c == sympy.true:
	return p
	elif c == sympy.false:
	return q

.

Regarding 2, fair enough, although it's not an extremely unusual pattern to use True as 1 when doing things like modular computations. May be good to leave a comment.

[eellison]

Since we guard on 0/1, wouldn't something like sym_float end up guarding anyway ?

[lezcano]

ah, good point :D

[lezcano]

Regarding 1, we can work around it by using Where

pytorch/torch/utils/_sympy/functions.py

Lines 151 to 163 in e6b3a8c

	class Where(sympy.Function):
	"""
	Good ol' ternary operator
	"""
	nargs = (3,)
	@classmethod
	def eval(cls, c, p, q):
	if c == sympy.true:
	return p
	elif c == sympy.false:
	return q

.

Regarding 2, fair enough, although it's not an extremely unusual pattern to use True as 1 when doing things like modular computations. May be good to leave a comment.

[jansel]

Inductor's codegen will likely choke on sympy.Where, so it might be extra work to support that. Since there isn't a real usecase needing that, I'll leave that unimplemented for now.

[jansel]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[lezcano]

I added Where and I already implemented the relevant codegen. This is used for wrapping negative indices when we have indirect indexing. But sure, this is good as-is for now.