This also concretizes some generic methods used in previous model serialization
that appeared to interact poorly with tf.function tracing.

PiperOrigin-RevId: 394295569