nest.flatten was mangling some attribute lists

Fixes #40895.

PiperOrigin-RevId: 339965951
Change-Id: Ibdbbb5a12b2d06a077b62902764ca4849a2b8d2a