Merge pull request #39 from gregszumel/bool_tensors

Boolean input support

Author: mortont

native/ortex/src/model.rs

@@ -9,8 +9,9 @@
 //! ```
 
 use crate::tensor::OrtexTensor;
-use crate::utils::map_opt_level;
-use std::convert::{TryFrom, TryInto};
+use crate::utils::{is_bool_input, map_opt_level};
+use std::convert::TryInto;
+use std::iter::zip;
 
 use ort::{Error, ExecutionProviderDispatch, Session};
 use rustler::resource::ResourceArc;
@@ -83,27 +84,46 @@ pub fn run(
     model: ResourceArc<OrtexModel>,
     inputs: Vec<ResourceArc<OrtexTensor>>,
 ) -> Result<Vec<(ResourceArc<OrtexTensor>, Vec<usize>, Atom, usize)>, Error> {
-    // TODO: can we handle an error more elegantly than just .unwrap()?
+    // Grab the session and run a forward pass with it
+    let session: &ort::Session = &model.session;
 
     let mut ortified_inputs: Vec<ort::SessionInputValue> = Vec::new();
-    for input in inputs {
-        let derefed_input: &OrtexTensor = &input;
-        let v: ort::SessionInputValue = derefed_input.try_into()?;
-        ortified_inputs.push(v);
-    }
 
-    // Grab the session and run a forward pass with it
-    let session: &ort::Session = &model.session;
+    for (elixir_input, onnx_input) in zip(inputs, &session.inputs) {
+        let derefed_input: &OrtexTensor = &elixir_input;
+        if is_bool_input(&onnx_input.input_type) {
+            // this assumes that the boolean input isn't huge -- we're cloning it twice;
+            // once below, once in the try_into()
+            let boolified_input: &OrtexTensor = &derefed_input.clone().to_bool();
+            let v: ort::SessionInputValue = boolified_input.try_into()?;
+            ortified_inputs.push(v);
+        } else {
+            let v: ort::SessionInputValue = derefed_input.try_into()?;
+            ortified_inputs.push(v);
+        }
+    }
 
     // Construct a Vec of ModelOutput enums based on the DynOrtTensor data type
     let outputs = session.run(&ortified_inputs[..])?;
-    outputs
-        .iter()
-        .map(|(_name, val)| {
-            let ortextensor: OrtexTensor = OrtexTensor::try_from(val)?;
-            let shape = ortextensor.shape();
-            let (dtype, bits) = ortextensor.dtype();
-            Ok((ResourceArc::new(ortextensor), shape, dtype, bits))
-        })
-        .collect()
+    let mut collected_outputs = Vec::new();
+
+    for output_descriptor in &session.outputs {
+        let output_name: &str = &output_descriptor.name;
+        let val = outputs.get(output_name).expect(
+            &format!(
+                "Expected {} to be in the outputs, but didn't find it",
+                output_name
+            )[..],
+        );
+
+        // NOTE: try_into impl here will implicitly map bool outputs to u8 outputs
+        let ortextensor: OrtexTensor = val.try_into()?;
+        let shape = ortextensor.shape();
+        let (dtype, bits) = ortextensor.dtype();
+
+        let collected_output = (ResourceArc::new(ortextensor), shape, dtype, bits);
+        collected_outputs.push(collected_output)
+    }
+
+    Ok(collected_outputs)
 }

native/ortex/src/tensor.rs

@@ -26,6 +26,9 @@ pub enum OrtexTensor {
     bf16(Array<half::bf16, IxDyn>),
     f32(Array<f32, IxDyn>),
     f64(Array<f64, IxDyn>),
+    // the bool input is for internal use only.
+    // Any Nx facing ops should panic if called on a bool input
+    bool(Array<bool, IxDyn>),
 }
 
 impl OrtexTensor {
@@ -43,6 +46,7 @@ impl OrtexTensor {
             OrtexTensor::bf16(y) => y.shape().to_owned(),
             OrtexTensor::f32(y) => y.shape().to_owned(),
             OrtexTensor::f64(y) => y.shape().to_owned(),
+            _ => panic!("Can't convert this type to Nx format"),
         }
     }
 
@@ -108,6 +112,7 @@ impl OrtexTensor {
                     .into_shape_with_order(shape)
                     .map_err(|e| rustler::Error::Term(Box::new(e.to_string())))?,
             )),
+            _ => panic!("Can't convert this type to Nx format"),
         }
     }
 
@@ -125,6 +130,7 @@ impl OrtexTensor {
             OrtexTensor::bf16(_) => (ortex_atoms::bf(), 16),
             OrtexTensor::f32(_) => (ortex_atoms::f(), 32),
             OrtexTensor::f64(_) => (ortex_atoms::f(), 64),
+            _ => panic!("Can't convert this type to Nx format"),
         }
     }
 
@@ -142,6 +148,7 @@ impl OrtexTensor {
             OrtexTensor::bf16(y) => get_bytes(y),
             OrtexTensor::f32(y) => get_bytes(y),
             OrtexTensor::f64(y) => get_bytes(y),
+            _ => panic!("Can't convert this type to Nx format"),
         };
         contents
     }
@@ -173,6 +180,25 @@ impl OrtexTensor {
             OrtexTensor::bf16(y) => OrtexTensor::bf16(slice_array(y, &slice_specs).to_owned()),
             OrtexTensor::f32(y) => OrtexTensor::f32(slice_array(y, &slice_specs).to_owned()),
             OrtexTensor::f64(y) => OrtexTensor::f64(slice_array(y, &slice_specs).to_owned()),
+            _ => panic!("Can't convert this type to Nx format"),
+        }
+    }
+
+    pub fn to_bool(self) -> OrtexTensor {
+        match self {
+            OrtexTensor::u8(y) => {
+                let bool_tensor = y.to_owned().mapv(|x| match x {
+                    0 => false,
+                    1 => true,
+                    _ => {
+                        panic!(
+                            "Tried to convert a u8 tensor to bool, but not every element is 0 or 1"
+                        )
+                    }
+                });
+                OrtexTensor::bool(bool_tensor)
+            }
+            t => panic!("Can't convert this type {:?} to bool", t.dtype()),
         }
     }
 }
@@ -253,8 +279,10 @@ impl TryFrom<&Value> for OrtexTensor {
             ort::TensorElementType::String => {
                 todo!("Can't return string tensors")
             }
+            // map the output into u8 space
             ort::TensorElementType::Bool => {
-                todo!("Can't return bool tensors")
+                let nd_array = e.try_extract_tensor::<bool>()?.into_owned();
+                OrtexTensor::u8(nd_array.mapv(|x| x as u8))
             }
         };
 
@@ -278,11 +306,32 @@ impl TryFrom<&OrtexTensor> for ort::SessionInputValue<'_> {
             OrtexTensor::u16(arr) => arr.clone().try_into()?,
             OrtexTensor::u32(arr) => arr.clone().try_into()?,
             OrtexTensor::u64(arr) => arr.clone().try_into()?,
+            OrtexTensor::bool(arr) => arr.clone().try_into()?,
         };
         Ok(r.into())
     }
 }
 
+impl Clone for OrtexTensor {
+    fn clone(&self) -> Self {
+        match self {
+            OrtexTensor::s8(t) => OrtexTensor::s8(t.clone()),
+            OrtexTensor::s16(t) => OrtexTensor::s16(t.clone()),
+            OrtexTensor::s32(t) => OrtexTensor::s32(t.clone()),
+            OrtexTensor::s64(t) => OrtexTensor::s64(t.clone()),
+            OrtexTensor::bf16(t) => OrtexTensor::bf16(t.clone()),
+            OrtexTensor::f16(t) => OrtexTensor::f16(t.clone()),
+            OrtexTensor::f32(t) => OrtexTensor::f32(t.clone()),
+            OrtexTensor::f64(t) => OrtexTensor::f64(t.clone()),
+            OrtexTensor::u8(t) => OrtexTensor::u8(t.clone()),
+            OrtexTensor::u16(t) => OrtexTensor::u16(t.clone()),
+            OrtexTensor::u32(t) => OrtexTensor::u32(t.clone()),
+            OrtexTensor::u64(t) => OrtexTensor::u64(t.clone()),
+            OrtexTensor::bool(t) => OrtexTensor::bool(t.clone()),
+        }
+    }
+}
+
 // Currently only supports concatenating tenors of the same type.
 //
 // This is a similar structure to the above match clauses, except each function

native/ortex/src/utils.rs

@@ -116,3 +116,12 @@ pub fn map_opt_level(opt: i32) -> GraphOptimizationLevel {
         _ => GraphOptimizationLevel::Disable,
     }
 }
+
+pub fn is_bool_input(inp: &ort::ValueType) -> bool {
+    match inp {
+        ort::ValueType::Tensor { ty, .. } => ty == &ort::TensorElementType::Bool,
+        ort::ValueType::Map { value, .. } => value == &ort::TensorElementType::Bool,
+        ort::ValueType::Sequence(boxed_input) => is_bool_input(boxed_input),
+        ort::ValueType::Optional(boxed_input) => is_bool_input(boxed_input),
+    }
+}