Several fixes/improvements to wrap_function:

1. Makes gradient tape work on captures post-pruning
2. Allows users to prune with wrapped_func.inputs even when the func includes captures
3. Makes sure to write variables to collections for better legacy compatibility in compat.v1 (currently if the outer init_scope is eager, variables won't be written to the global_variables / trainable_variables collections)

PiperOrigin-RevId: 227936352