Fix Torch ONNX export not respecting InputSpec.name (#21646)

* fix(torch-export): ONNX export not respecting InputSpec.name

* Update onnx.py

* Pythonifying the code

* api-gen

* idek

* Inputspec Attribute Error

* Inputspec Attribute Error

* InputSpec missing dtype

* layer not built yet

Author: Flakes342

keras/src/export/onnx.py

@@ -80,6 +80,10 @@ def export_onnx(
                 "The model provided has never called. "
                 "It must be called at least once before export."
             )
+    input_names = [
+        getattr(spec, "name", None) or f"input_{i}"
+        for i, spec in enumerate(input_signature)
+    ]
 
     if backend.backend() in ("tensorflow", "jax"):
         from keras.src.utils.module_utils import tf2onnx
@@ -143,6 +147,7 @@ def export_onnx(
                     sample_inputs,
                     verbose=actual_verbose,
                     opset_version=opset_version,
+                    input_names=input_names,
                     dynamo=True,
                 )
                 if hasattr(onnx_program, "optimize"):
@@ -161,6 +166,7 @@ def export_onnx(
                     filepath,
                     verbose=actual_verbose,
                     opset_version=opset_version,
+                    input_names=input_names,
                 )
     else:
         raise NotImplementedError(

keras/src/export/onnx_test.py

@@ -14,6 +14,7 @@
 from keras.src import testing
 from keras.src import tree
 from keras.src.export import onnx
+from keras.src.layers.input_spec import InputSpec as InputSpec
 from keras.src.saving import saving_lib
 from keras.src.testing.test_utils import named_product
 
@@ -269,3 +270,31 @@ def test_export_with_opset_version(self, opset_version):
         if opset_version is not None:
             onnx_model = onnx_lib.load(temp_filepath)
             self.assertEqual(onnx_model.opset_import[0].version, opset_version)
+
+    def test_export_with_input_names(self):
+        """Test ONNX export uses InputSpec.name for input names."""
+        import onnx as onnx_lib
+
+        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
+        model = get_model("sequential")
+        batch_size = 3 if backend.backend() != "torch" else 1
+        ref_input = np.random.normal(size=(batch_size, 10)).astype("float32")
+        ref_output = model(ref_input)
+
+        # Test with custom input name
+        input_spec = [
+            InputSpec(
+                name="custom_input", shape=(batch_size, 10), dtype="float32"
+            )
+        ]
+        onnx.export_onnx(model, temp_filepath, input_signature=input_spec)
+
+        onnx_model = onnx_lib.load(temp_filepath)
+        input_names = [input.name for input in onnx_model.graph.input]
+        self.assertIn("custom_input", input_names)
+
+        ort_session = onnxruntime.InferenceSession(temp_filepath)
+        ort_inputs = {
+            k.name: v for k, v in zip(ort_session.get_inputs(), [ref_input])
+        }
+        self.assertAllClose(ref_output, ort_session.run(None, ort_inputs)[0])