-
Notifications
You must be signed in to change notification settings - Fork 39
Description
As has been discussed elsewhere, it can be quite convenient to be able to describe a graph's input and output tensor dimensions and type. For example, a user of wasi-nn currently has to know the exact specifications of the ML graph being used at compile-time. This is inconvenient if one wants to dynamically specify the model. Also, users of wasi-nn in languages like Rust must allocate the exact-size buffers to contain the output tensors. These issues can be resolved by allowing users to inspect some details of the graph's inputs and outputs.
I believe wasi-nn could be improved by adding the following WIT:
record tensor-description {
dimensions: tensor-dimensions,
tensor-type: tensor-type,
}
get-num-inputs: func(graph: graph) -> u32
get-input-description: func(graph: graph, index: u32) -> expected<tensor-description, error>
get-num-outputs: func(graph: graph) -> u32
get-output-description: func(graph: graph, index: u32) -> expected<tensor-description, error>
One additional refactoring this suggests is to use tensor-description
in the tensor
record (see here). With these additional functions in the specification, we could properly create high-level bindings as is suggested here: bytecodealliance/wasi-nn#68.