Repository: onnx/onnx
File: onnx/defs/nn/defs.cc (C++)

Code before:
            }
            for (const auto& output : outputs) {
              graph.add_output()->set_name(output);
            }
            for (const auto& node : branch_nodes) {
              *graph.add_node() = node;
            }
            return graph;
          };

          FunctionBuilder builder(functionProto);
          // Set input tensor to the correct shape if input shape is 3D
          // NewShape = [batch_size, sequence_length, num_heads, head_size]

          builder.Add("XIn = Identity (X)")
              .Add("XShape = Shape (X)") // shape of input tensor
              .Add("Rank = Size (XShape)"); // rank of input tensor

          // Reshape tensor to 4D if input is 3D
          builder.Const1D("Zero1D", (int64_t)0)
              .Const1D("NumHeads", num_heads) // num_heads
              .Const1D("NegOne", (int64_t)(-1)) // head_size, inferred from other dimensions
              .Add("NewShape = Concat <axis = 0> (Zero1D, Zero1D, NumHeads, NegOne)");

          // Pick the correct X based on number of dimensions
          builder.Const1D("Three1D", (int64_t)3).Add("XCond = Equal(Rank, Three1D)");

          // Build nodes for then and else branches
          std::vector<NodeProto> then_nodes =
              FunctionBodyHelper::BuildNodes({{{"XReshaped"}, "Reshape", {"XIn", "NewShape"}}});

          std::vector<NodeProto> else_nodes = FunctionBodyHelper::BuildNodes(
              {{{"XTransposed"}, "Transpose", {"XIn"}, {{"perm", std::vector<int64_t>{0, 2, 1, 3}}}}});

          // Build nodes for then and else branches
          auto then_branch_graph = AddSubgraph("x_then_branch_graph", then_nodes, {"XIn", "NewShape"}, {"XReshaped"});
          auto else_branch_graph = AddSubgraph("x_else_branch_graph", else_nodes, {"XIn"}, {"XTransposed"});

          // Build nodes for then and else branches
          builder.Add("XInput = If (XCond)", "then_branch", then_branch_graph, "else_branch", else_branch_graph);

          // Rotary embedding dimension is the value along which the input is to be split
          // There are two cases for the rotary embedding dimension:
          // 1. Complete rotation: rotary embedding dimension defaults to head_size, rotary_embedding_dim = cos.shape[3]
          // * 2 or head_size
          // 2. Partial rotation: rotary embedding dimension is provided, rotary_embedding_dim = rotary_embedding_dim

          builder.Add("HeadSize = Shape <start = 3, end = 4> (XInput)");
          if (rotary_embedding_dim > 0) {
            builder.Const1D("RotaryEmbedDim", rotary_embedding_dim);
          } else {

Review comment (suggestion):
Please check if Rank is 0D or 1D here. I think it is likely 0D, as it should be (based on what Size returns above). In which case, we should be using `Three0D` here.

Code after:
          }
          if (x_tp->tensor_type().shape().dim_size() == 4) {
            is_input_4d = 1;
          } else if (x_tp->tensor_type().shape().dim_size() == 3) {
            is_input_4d = 0;
          } else {
            return false;
          }

          FunctionBuilder builder(functionProto);
          // Set input tensor to the correct shape if input shape is 3D
          // NewShape = [batch_size, sequence_length, num_heads, head_size]

          // Reshape tensor to 4D if input is 3D
          builder.Const1D("Zero1D", (int64_t)0)
              .Const1D("NumHeads", num_heads) // num_heads
              .Const1D("NegOne", (int64_t)(-1)); // head_size, inferred from other dimensions

          if (is_input_4d == 0) {
            builder.Add("NewShape = Concat <axis = 0> (Zero1D, Zero1D, NumHeads, NegOne)")
                .Add("XIn = Reshape (X, NewShape)"); // new shape of input tensor: 4D tensor
          } else {
            builder.Add("XIn = Transpose <perm = [0, 2, 1, 3]> (X)");
          }

          // Rotary embedding dimension is the value along which the input is to be split
          // There are two cases for the rotary embedding dimension:
          // 1. Complete rotation: rotary embedding dimension defaults to head_size, rotary_embedding_dim = cos.shape[3]
          // * 2 or head_size
          // 2. Partial rotation: rotary embedding dimension is provided, rotary_embedding_dim = rotary_embedding_dim

          builder.Add("HeadSize = Shape <start = 3, end = 4> (XIn)");
          if (rotary_embedding_dim > 0) {
            builder.Const1D("RotaryEmbedDim", rotary_embedding_dim);
          } else {
            builder.Add("RotaryEmbedDim = Identity(HeadSize)");
          }
          builder.Const1D("Two1D", (int64_t)2)
              .Add("NoRotateLength = Sub(HeadSize, RotaryEmbedDim)")
              .Add("RotateSplitLengths = Concat <axis = 0> (RotaryEmbedDim, NoRotateLength)");
          // shape of input to rotate = input[:,:,:,:rotary_embedding_dim]
          // shape of input not to rotate = input[:,:,:,rotary_embedding_dim:]
          builder.Add("XToRotate, XNoRotate = Split <axis = -1> (XIn, RotateSplitLengths)");

          // Gather the cos and sine matrices from the respective caches using position ids if provided.
          // Otherwise Gather op functions as an Identity op.
          // Unsqueeze applied to make cos and sin matrices have dimensions that are
          // valid for multiplication with input when is split. For cases where rotary_embedding_dim is provided,
          // slice the matrix values until that index only
          if (ctx.hasInput(3)) {
            builder