Joel Schlosser 8ba555ec8a Fix where() for NJT (#141500) ... **Background:** It's common to use `scalar_tensor()` in the input to `where()` to convert any scalars present to compatible tensors with matching options, *including layout*. This shows up in various places, notably including derivative formulas ([example](78491d6afc/tools/autograd/derivatives.yaml (L432-L434))). It causes problems for NJTs because they have `layout=torch.jagged` and it never makes sense to create a scalar tensor with this layout. Some of the breakage only seems to happen in CI for reasons I don't fully understand (see the revert of #140736 due to softshrink's derivative formula). **This PR:** * Allows non-contiguous NJT inputs to `where()` + adds tests for this * Handles scalar tensor / dense tensor inputs for `condition` / `other` + adds tests for this * Uses limited `broadcast_tensors()` / `broadcast_to()` support * Improves `expand()` to work on non-contig NJTs * Changes `scalar_tensor()` to use `torch.strided` instead of `torch.jagged` in both eager and torch.compile (i.e. meta registration) * Changes backward formulas for `sinc`, `pow`, `special.i1`, and `special.i1e` to uses `scalar_tensor()` instead of e.g. `zeros({})` **Alternative approach:** Update all problematic usages of `scalar_tensor()` to avoid ever passing `layout=torch.jagged`. This is an extensive change and includes `torch.where()` logic, a bunch of derivative formulas, and likely other places not yet discovered. Pull Request resolved: https://github.com/pytorch/pytorch/pull/141500 Approved by: https://github.com/malfet, https://github.com/cpuhrsch, https://github.com/soulitzer		2024-11-26 20:13:27 +00:00
..
definitions	Fix where() for NJT (#141500)	2024-11-26 20:13:27 +00:00
__init__.py
core.py	General per-SampleInput xfail / skip system (#140443)	2024-11-19 23:09:38 +00:00
refs.py	[BE][Easy][19/19] enforce style for empty lines in import segments in torch/[o-z]*/ (#129771)	2024-08-01 17:07:14 +00:00
utils.py	Remove unused Python variables in torch/[b-z]* (#136963)	2024-10-19 16:45:22 +00:00