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
..
_internal	Fix where() for NJT (#141500)	2024-11-26 20:13:27 +00:00
__init__.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
_comparison.py	Strict shape checking for NJTs with TestCase.assertEqual() (#131898)	2024-07-30 20:05:48 +00:00
_creation.py	Ensure noncontiguous tensor creation tests offsetting (#136396)	2024-10-02 00:40:43 +00:00
_utils.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