saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Don't check cache shape when using dynamic axis (#10090)

Browse source 
Co-authored-by: Ubuntu <wy@linux-v100.aidmrjtolptuzevavgwhrapqcd.jx.internal.cloudapp.net>
This commit is contained in:
Ye Wang 2021-12-20 21:19:29 -08:00 committed by GitHub
parent 4e9e01cb3c
commit 7a1bdc2052
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 16 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
onnxruntime/core/graph/contrib_ops/contrib_defs.cc
View file

@ -550,20 +550,21 @@ void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx
if (cache_dims.size() != 4) {
fail_shape_inference("key and value cache shall be 4 dimensions");
}
if (!cache_dims[0].has_dim_value() ||
!cache_dims[1].has_dim_value() ||
!cache_dims[2].has_dim_value() ||
!cache_dims[3].has_dim_value()) {
fail_shape_inference("key and value cache dimensions value shall not be null");
}
ONNX_NAMESPACE::TensorShapeProto new_cache_shape;
*new_cache_shape.add_dim() = cache_shape.dim(0);
*new_cache_shape.add_dim() = cache_shape.dim(1);
new_cache_shape.add_dim();
*new_cache_shape.add_dim() = cache_shape.dim(3);
// has_dim_value() will return false if value is dynamic
if (cache_dims[0].has_dim_value() &&
cache_dims[1].has_dim_value() &&
cache_dims[2].has_dim_value() &&
cache_dims[3].has_dim_value()) {
updateOutputShape(ctx, 1, new_cache_shape);
updateOutputShape(ctx, 2, new_cache_shape);
ONNX_NAMESPACE::TensorShapeProto new_cache_shape;
*new_cache_shape.add_dim() = cache_shape.dim(0);
*new_cache_shape.add_dim() = cache_shape.dim(1);
new_cache_shape.add_dim();
*new_cache_shape.add_dim() = cache_shape.dim(3);
updateOutputShape(ctx, 1, new_cache_shape);
updateOutputShape(ctx, 2, new_cache_shape);
}
}
}
}
@ -611,7 +612,7 @@ void BeamSearchShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) {
if (!(input_ids_dims[0].has_dim_value() && input_ids_dims[1].has_dim_value())) {
return;
}
int64_t batch_size = input_ids_dims[0].dim_value();
int64_t sequence_length = input_ids_dims[1].dim_value();
@ -883,7 +884,7 @@ Some boolean parameters are passed by runtime input for generic purpose
.TypeConstraint("B", {"tensor(bool)"}, "Constrain key_padding_mask to bool tensors.")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
DecoderAttentionTypeAndShapeInference(ctx);
});
});
static const char* EmbedLayerNormalization_ver1_doc = R"DOC(
EmbedLayerNormalization is the fusion of embedding layer in BERT model, with optional mask processing.