Ensure standard shape for triadd layernorm

4ac2919f · Paul · 6d34b90f · 4ac2919f
Commit 4ac2919f authored Jun 06, 2022 by Paul
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 0 deletions

src/targets/gpu/fuse_ops.cpp src/targets/gpu/fuse_ops.cpp +5 -0

No files found.
--- a/src/targets/gpu/fuse_ops.cpp
+++ b/src/targets/gpu/fuse_ops.cpp
@@ -255,6 +255,11 @@ MIGRAPHX_REGISTER_OP(hip_layernorm)

 struct hip_triadd_layernorm : ternary_device<hip_triadd_layernorm, &device::triadd_layernorm>
 {
+    shape compute_shape(const std::vector<shape>& inputs) const
+    {
+        check_shapes{inputs, *this}.has(4).standard();
+        return inputs[0];
+    }
    // Empty finalize to skip dimension reduction
    void finalize(context&, const shape&, const std::vector<shape>&) {}
 };