Skip to content

Commit

Permalink
chore: document the builder functions
Browse files Browse the repository at this point in the history
  • Loading branch information
pixelspark committed May 22, 2023
1 parent e95dfed commit d3536f7
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions wonnx/src/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,13 @@ impl<'model> From<&TensorRef<'model>> for Input<'model> {
}

impl<'model> TensorRef<'model> {
/// Element-wise addition of two tensors (ONNX Add operator)
pub fn add(&self, rhs: &Self) -> Self {
assert_eq!(self.output_shape, rhs.output_shape);
self.binary_op(rhs, "Add", self.output_shape.clone())
}

/// Element-wise negation of this tensor (ONNX Neg operator)
pub fn neg(&self) -> Self {
self.unary_mapping_op("Neg")
}
Expand All @@ -54,6 +56,7 @@ impl<'model> TensorRef<'model> {
}
}

/// Convolution (ONNX Conv operator)
pub fn conv(&self, weights: &Self, kernel_shape: &[usize], output_dims: &[usize]) -> Self {
let output_shape = Shape::from(self.output_shape.data_type, output_dims);
let mut op = OperatorDefinition::new(
Expand Down Expand Up @@ -95,6 +98,7 @@ impl<'model> TensorRef<'model> {
}
}

/// Create a tensor reference representing input supplied at inference time.
pub fn input<'model, S: ToString>(
name: S,
scalar_type: ScalarType,
Expand All @@ -114,6 +118,7 @@ pub fn input<'model, S: ToString>(
}
}

/// Create a tensor reference containing static data (included in the model).
pub fn tensor<'model, S: ToString>(
name: S,
dims: &[usize],
Expand Down

0 comments on commit d3536f7

Please sign in to comment.