Use rank-2 linear for UMA spectral gates

[mlazos]

Summary

SpectralAtomwise retains a singleton coefficient dimension when projecting
its scalar slice. That sends [num_atoms, 1, sphere_channels] tensors through
Linear, which lowers the production shape to broadcast BMMs in the forward
and input-gradient graphs.

Squeeze the known singleton dimension before the scalar MLP. The gate already
accepts the resulting rank-2 tensor, so model semantics are unchanged while
Inductor can emit ordinary addmm/mm calls without layout conversions.

Performance

With torch.compile(dynamic=True), rank-2 improved the full endpoint from 16.98888 to 16.81807 ms with fresh caches, a 0.17081 ms (1.01%) gain, and from 16.99091 to 16.83604 ms with warm caches, a 0.15487 ms (0.91%) gain. It removed eight symbolic broadcast BMM calls.

The isolated scalar-MLP forward and input-VJP benchmark improved from 0.36388 to 0.35275 ms, a 3.1% gain.

Configuration: H100, TF32, UMA-S-1p2, 1,000 atoms, energy/forces/stress, umas_fast_gpu, merge_mole=True, external_graph_gen=False, internal graph v3, internal_graph_skin=0, compile=True, compile_dynamic_shapes=True, and full CUDA graph replay. Steady-state execution had zero recompiles and zero dynamic graph fallbacks.

Testing

PYTHONPATH=$PWD/src:/data/users/mlazos/pytorch python -m pytest -q tests/core/models/uma/test_escn_md.py
ruff check --ignore UP035 src/fairchem/core/models/uma/escn_md_block.py tests/core/models/uma/test_escn_md.py
git diff --check

All 13 tests passed.

Authored with assistance from Codex.