Skip to content

Commit

Permalink
feat: support optional ValueTypes
Browse files Browse the repository at this point in the history
According to the C API, these only appear in model metadata, so no need to implement another `Value<T>`.
  • Loading branch information
decahedron1 committed Sep 5, 2024
1 parent dc08a6b commit 52422ae
Showing 1 changed file with 12 additions and 1 deletion.
13 changes: 12 additions & 1 deletion src/value/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,9 @@ pub enum ValueType {
key: TensorElementType,
/// The map value type.
value: TensorElementType
}
},
/// An optional value, which may or may not contain a [`Value`].
Optional(Box<ValueType>)
}

impl ValueType {
Expand All @@ -110,6 +112,15 @@ impl ValueType {
ortsys![unsafe CastTypeInfoToMapTypeInfo(typeinfo_ptr, &mut info_ptr)?; nonNull(info_ptr)];
unsafe { extract_data_type_from_map_info(info_ptr)? }
}
ort_sys::ONNXType::ONNX_TYPE_OPTIONAL => {
let mut info_ptr: *const ort_sys::OrtOptionalTypeInfo = std::ptr::null_mut();
ortsys![unsafe CastTypeInfoToOptionalTypeInfo(typeinfo_ptr, &mut info_ptr)?; nonNull(info_ptr)];

let mut contained_type: *mut ort_sys::OrtTypeInfo = std::ptr::null_mut();
ortsys![unsafe GetOptionalContainedTypeInfo(info_ptr, &mut contained_type)?; nonNull(contained_type)];

ValueType::Optional(Box::new(ValueType::from_type_info(contained_type)?))
}
_ => unreachable!()
};
ortsys![unsafe ReleaseTypeInfo(typeinfo_ptr)];
Expand Down

0 comments on commit 52422ae

Please sign in to comment.