Joel Schlosser a8382847f4 Support rms_norm() for NJT (#135872) ... `rms_norm()` is a nice-to-have for ViT :) This PR: * SymInt-ifies `rms_norm()`, allowing NJT to use the same decomp. * Adds torch_function-based input validation logic for nested-specific stuff (no normalization supported over the ragged dim for now) on the python NJT side. * Adds multi-dim support (on non-ragged, non-batch dims) to `mean()` for NJT. Pull Request resolved: https://github.com/pytorch/pytorch/pull/135872 Approved by: https://github.com/mikaylagawarecki ghstack dependencies: #125947		2024-09-17 18:09:20 +00:00
..
_internal	Support rms_norm() for NJT (#135872)	2024-09-17 18:09:20 +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	[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	[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