saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-12 17:57:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Improve shape inference for OneHot (#4452)

Browse source 
* Improve shape inference for OneHot

Attempt to get the depth parameter before adding a new symbolic dimension.

* Update symbolic shape infer

* Nit
This commit is contained in:
Oliver Rausch 2020-08-07 23:05:20 +02:00 committed by GitHub
parent 9c729d1719
commit 9c3153acd6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
onnxruntime/core/providers/nuphar/scripts/symbolic_shape_infer.py
View file

@ -776,10 +776,13 @@ class SymbolicShapeInference:
vi.CopyFrom(helper.make_tensor_value_info(node.output[0], vi.type.tensor_type.elem_type, [input_rank, nz_len]))
def _infer_OneHot(self, node):
shape = self._get_shape(node, 0)
sympy_shape = self._get_sympy_shape(node, 0)
depth = self._try_get_value(node, 1)
axis = get_attribute(node, 'axis', -1)
axis = handle_negative_axis(axis, len(shape)+1)
new_shape = shape[:axis] + [self._new_symbolic_dim_from_output(node)] + shape[axis:]
axis = handle_negative_axis(axis, len(sympy_shape) + 1)
new_shape = get_shape_from_sympy_shape(
sympy_shape[:axis] + [self._new_symbolic_dim_from_output(node) if not is_literal(depth) else depth] +
sympy_shape[axis:])
vi = self.known_vi_[node.output[0]]
vi.CopyFrom(helper.make_tensor_value_info(node.output[0], self.known_vi_[node.input[2]].type.tensor_type.elem_type, new_shape))