Use index in set_input

Author: rahulchaphalkar

crates/wasi-nn/src/backend/pytorch.rs

@@ -79,6 +79,7 @@ impl BackendGraph for PytorchGraph {
             module: self.module.clone(),
             inputs: Vec::new(),
             output: TchTensor::new(),
+            id_type: None,
         });
         Ok(box_.into())
     }
@@ -87,41 +88,80 @@ impl BackendGraph for PytorchGraph {
 unsafe impl Sync for PytorchExecutionContext {}
 struct PytorchExecutionContext {
     module: Arc<Mutex<tch::CModule>>,
-    inputs: Vec<tch::Tensor>,
+    inputs: Vec<Option<tch::Tensor>>,
     output: tch::Tensor,
+    id_type: Option<Id>,
 }
 
 impl BackendExecutionContext for PytorchExecutionContext {
-    fn set_input(&mut self, _index: Id, input_tensor: &Tensor) -> Result<(), BackendError> {
-        // Input index is not used in pytorch models. The forward method to a model passes the tensor/data to
-        // the appropriate layer of the model.
+    fn set_input(&mut self, id: Id, input_tensor: &Tensor) -> Result<(), BackendError> {
         let kind = input_tensor.ty.try_into()?;
         let dimensions = input_tensor
             .dimensions
             .iter()
             .map(|&dim| dim as i64)
             .collect::<Vec<_>>();
-        self.inputs.push(TchTensor::from_data_size(
-            &input_tensor.data,
-            &dimensions,
-            kind,
-        ));
-        Ok(())
+        let tensor = TchTensor::from_data_size(&input_tensor.data, &dimensions, kind);
+        match id {
+            Id::Index(i) => {
+                // Check if id_type is already set and if it matches the current id type
+                if let Some(Id::Name(_)) = self.id_type {
+                    return Err(BackendError::BackendAccess(anyhow::anyhow!(
+                        "Cannot mix u32 and str indexes"
+                    )));
+                }
+                // Set id_type if not already set
+                if self.id_type.is_none() {
+                    self.id_type = Some(Id::Index(0)); // Provide a u32 value for Index
+                }
+                let i = i as usize;
+                if i >= self.inputs.len() {
+                    self.inputs.resize_with(i + 1, || None);
+                }
+                self.inputs[i] = Some(tensor);
+                Ok(())
+            }
+            Id::Name(_) => {
+                // Check if id_type is already set and if it matches the current id type
+                if let Some(Id::Index(_)) = self.id_type {
+                    return Err(BackendError::BackendAccess(anyhow::anyhow!(
+                        "Cannot mix u32 and str indexes"
+                    )));
+                }
+                // Set id_type if not already set
+                if self.id_type.is_none() {
+                    self.id_type = Some(Id::Name(String::new())); // Provide a str value for Name
+                }
+                if self.inputs.get(0).is_some() {
+                    return Err(BackendError::BackendAccess(anyhow::anyhow!(
+                        "The pytorch backend does not support multiple named inputs"
+                    )));
+                } else {
+                    self.inputs.push(Some(tensor));
+                }
+                Ok(())
+            }
+        }
     }
 
     fn compute(&mut self) -> Result<(), BackendError> {
-        // Use forward method on the compiled module/model after locking the mutex, and pass the input tensor to it
-        self.output = self
-            .module
-            .lock()
-            .unwrap()
-            .forward_ts(&self.inputs)
-            .unwrap();
+        let inputs: Vec<tch::Tensor> = self
+            .inputs
+            .iter()
+            .enumerate()
+            .map(|(index, opt)| {
+                opt.as_ref()
+                    .expect(&format!("Input tensor at index {} not set up", index))
+                    .shallow_clone()
+            })
+            .collect();
+        // Use forward_ts method on the compiled module/model after locking the mutex, and pass the input tensor to it
+        self.output = self.module.lock().unwrap().forward_ts(&inputs).unwrap();
         Ok(())
     }
 
     fn get_output(&mut self, _index: Id) -> Result<Tensor, BackendError> {
-        // Output index is not used in pytorch models. The forward method to a model returns the output tensor.
+        // Output index is not used. The forward_ts method to a model returns a single output tensor.
         let numel = self.output.numel();
         let dimensions = self.output.size();
         let ty = self.output.kind().try_into()?;