pub fn unit_output(op: &Op) -> Option<(usize, usize)>
The units an op writes its result in, (units, width), when it can write them as a part of every unit of a wider tensor: a norm, a rotation, a per-head product.
(units, width)