When exporting StableHLO to SavedModel, also include the original var (…

…#5711)

* Pass bundle object to make_tf_function, so that it will include tf.variables as args to the XlaCallModule op.

* fix format

* When exporting StableHLO to SavedModel, also include the original variable tensor name from pytorch.

torch_xla/tf_saved_model_integration.py

@@ -81,7 +81,8 @@ def save_stablehlo_graph_as_tf(
   bundle = copy.deepcopy(stablehlo_program._bundle)
   tfm = tf.Module()
   bundle.state_dict = {
-      k: tf.Variable(v, trainable=False) for k, v in bundle.state_dict.items()
+      k: tf.Variable(v, trainable=False, name=k)
+      for k, v in bundle.state_dict.items()
   }
   bundle.additional_constants = [
       tf.Variable(v, trainable=False) for v in bundle.additional_constants