Fix rotary embedding to follow spec on attribute constraints (#26044)

### Description
<!-- Describe your changes. -->

`num_heads` is not necessarily required from users when input shape is
4D.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

To follow ONNX spec,
https://github.com/onnx/onnx/blob/main/docs/Operators.md#RotaryEmbedding,
the original constraints on attributes were wrong.

NOTE: 3 rotary embedding tests are expected to be wrong until next
release.

Author: titaiwangms

onnxruntime/core/providers/cpu/llm/rotary_embedding.cc

@@ -30,10 +30,6 @@ RotaryEmbedding<T>::RotaryEmbedding(const OpKernelInfo& info) : OpKernel(info) {
   num_heads = static_cast<int>(info.GetAttrOrDefault<int64_t>("num_heads", 0));
   rotary_embedding_dim = static_cast<int>(info.GetAttrOrDefault<int64_t>("rotary_embedding_dim", 0));
   interleaved = (info.GetAttrOrDefault<int64_t>("interleaved", 0) == 1);  // Turn 0/1 into bool
-
-  if (rotary_embedding_dim > 0) {
-    ORT_ENFORCE(num_heads > 0, "num_heads must be provided if rotary_embedding_dim is specified");
-  }
 }
 
 // TODO: rotary embedding in place
@@ -111,6 +107,15 @@ Status RotaryEmbedding<T>::Compute(OpKernelContext* context) const {
   // Optional position_ids input, can be nullptr
   const Tensor* position_ids = context->Input<Tensor>(3);
 
+  // If rotary_embedding_dim is set (>0) and num_heads attribute not provided (==0),
+  // we can only proceed if input is 4D (B, num_heads, S, head_size) so num_heads can be inferred.
+  if (rotary_embedding_dim > 0 && num_heads <= 0) {
+    const auto& dims = X->Shape().GetDims();
+    ORT_ENFORCE(dims.size() == 4,
+                "Attribute 'num_heads' must be provided when 'rotary_embedding_dim' is specified "
+                "and input is not rank-4 (batch, num_heads, sequence, head).");
+  }
+
   RotaryParameters parameters = {};
   ORT_RETURN_IF_ERROR(rotary_embedding_helper::CheckInputs<Tensor>(X,
                                                                    position_ids,

onnxruntime/core/providers/cuda/llm/rotary_embedding.cc

@@ -44,6 +44,15 @@ Status RotaryEmbedding<T>::ComputeInternal(OpKernelContext* context) const {
   const Tensor* sin_cache = context->Input<Tensor>(2);
   const Tensor* position_ids = context->Input<Tensor>(3);  // Optional, can be nullptr
 
+  // If rotary_embedding_dim is set (>0) and num_heads attribute not provided (==0),
+  // we can only proceed if input is 4D (B, num_heads, S, head_size) so num_heads can be inferred.
+  if (rotary_embedding_dim > 0 && num_heads <= 0) {
+    const auto& dims = input->Shape().GetDims();
+    ORT_ENFORCE(dims.size() == 4,
+                "Attribute 'num_heads' must be provided when 'rotary_embedding_dim' is specified "
+                "and input is not rank-4 (batch, num_heads, sequence, head).");
+  }
+
   RotaryParameters parameters = {};
   ORT_RETURN_IF_ERROR(rotary_embedding_helper::CheckInputs<Tensor>(input,
                                                                    position_ids,

onnxruntime/test/onnx/TestCase.cc

@@ -1408,6 +1408,9 @@ std::unique_ptr<std::set<BrokenTest>> GetBrokenTests(const std::string& provider
     broken_tests->insert({"gridsample_volumetric_nearest_align_corners_1", "unknown version"});
     broken_tests->insert({"rotary_embedding_no_position_ids_expanded", "unknown version"});
     broken_tests->insert({"rotary_embedding_no_position_ids_interleaved_expanded", "unknown version"});
+    broken_tests->insert({"rotary_embedding_no_position_ids_rotary_dim", "unknown version"});
+    broken_tests->insert({"rotary_embedding_with_interleaved_rotary_dim", "unknown version"});
+    broken_tests->insert({"rotary_embedding_with_rotary_dim", "unknown version"});
     // Fails since QNN SDK 2.17.0:
     // expected 7.70947 (40f6b3f3), got 7.84096 (40fae920), diff: 0.131491, tol=0.00870947 idx=419. 100 of 1715 differ
     broken_tests->insert({"facedetection_op8_qdq", "result differs"});

onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc

@@ -123,12 +123,10 @@
         "^test_if_opt",
         "^test_loop16_seq_none",
         "^test_identity_opt",
+        // rotary dim should be fixed in onnx==1.19.1
         "^test_rotary_embedding_no_position_ids_rotary_dim",
         "^test_rotary_embedding_with_interleaved_rotary_dim",
         "^test_rotary_embedding_with_rotary_dim",
-        "^test_rotary_embedding_3d_input_expanded",
-        "^test_rotary_embedding_interleaved_expanded",
-        "^test_rotary_embedding_no_position_ids_interleaved_expanded",
         "^test_rotary_embedding_expanded", //webgpu
         "^test_rotary_embedding_no_position_ids_expanded", //webgpu
         // Following tests are for opset 16 ops and are not yet implemented in ORT