Correctly set layouts for resource variables.

Previously if a resource variable has non-descending layout, the
layout will not be set correctly and a default layout will be
used. This changes fixes this issue by:

- Add a xla shape field to represent on device shape and layout.
- Set the above field when assigning to a resource variable.
- When generating program, use the layout in the above field to set the layout of the program.

PiperOrigin-RevId: 233826352