Works around the fact that we can't compile conditional initialization with XLA. It still won't work if variable initialization depends on function arguments.
Also fixes an issue where get_concrete_function would sometimes clobber variables which had already been initialized.
PiperOrigin-RevId: 231846163