PyPI - onnx - Versions diffs - 1.18.0__cp313-cp313-macosx_12_0_universal2.whl → 1.19.1rc1__cp313-cp313-macosx_12_0_universal2.whl - Mend

onnx 1.18.0__cp313-cp313-macosx_12_0_universal2.whl → 1.19.1rc1__cp313-cp313-macosx_12_0_universal2.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.

This version of onnx might be problematic. Click here for more details.

Files changed (1858) hide show

onnx/defs/tensor/defs.cc CHANGED Viewed

@@ -14,7 +14,7 @@
 namespace ONNX_NAMESPACE {
-static const char* Cast_ver19_doc = R"DOC(
+static const char* Cast_ver24_doc = R"DOC(
 The operator casts the elements of a given input tensor to a data type
 specified by the 'to' argument and returns an output tensor of the same size in
 the converted type. The 'to' argument must be one of the data types specified
@@ -52,40 +52,58 @@ if the destination type is not a float 8 type.
   * fixed point: `{1, 0}`.
   * bool: no change.
-Float 8 type were introduced to speed up the training of
+Float 8 types (E4M3FN, E4M3FNUZ, E5M2, E5M2FNUZ) were introduced to speed up the training of
 deep models. By default the conversion of a float *x* obeys
 to the following rules. `[x]` means the value rounded to
 the target mantissa width.
-| x | E4M3FN | E4M3FNUZ | E5M2 | E5M2FNUZ |
-|------|----|----|----|----|
-| 0 | 0 | 0 | 0 | 0 |
-|-0 | -0 | 0 | -0 | 0 |
-| NaN | NaN | NaN | NaN | NaN |
-| +/- Inf | +/- FLT_MAX | NaN | FLT_MAX | NaN |
-| [x] > FLT_MAX | FLT_MAX | FLT_MAX | FLT_MAX | FLT_MAX |
-| [x] < -FLT_MAX | -FLT_MAX | -FLT_MAX | -FLT_MAX | -FLT_MAX |
-| else | RNE | RNE | RNE | RNE |
+| x                 | E4M3FN   | E4M3FNUZ | E5M2     | E5M2FNUZ |
+| ----------------- | -------- | -------- | -------- | -------- |
+| 0                 | 0        | 0        | 0        | 0        |
+| -0                | -0       | 0        | -0       | 0        |
+| NaN               | NaN      | NaN      | NaN      | NaN      |
+| Inf               | FLT_MAX  | FLT_MAX  | FLT_MAX  | FLT_MAX  |
+| -Inf              | -FLT_MAX | -FLT_MAX | -FLT_MAX | -FLT_MAX |
+| \[x\] > FLT_MAX   | FLT_MAX  | FLT_MAX  | FLT_MAX  | FLT_MAX  |
+| \[x\] \< -FLT_MAX | -FLT_MAX | -FLT_MAX | -FLT_MAX | -FLT_MAX |
+| else              | RNE      | RNE      | RNE      | RNE      |
 The behavior changes if the parameter 'saturate' is set to False.
 The rules then become:
-| x | E4M3FN | E4M3FNUZ | E5M2 | E5M2FNUZ |
-|------|----|----|----|----|
-| 0 | 0 | 0 | 0 | 0 |
-|-0 | -0 | 0 | -0 | 0 |
-| NaN | NaN | NaN | NaN | NaN |
-| +/- Inf | NaN | NaN | +/- Inf | NaN |
-| [x] > FLT_MAX | NaN | NaN | Inf | NaN |
-| [x] < -FLT_MAX | NaN | NaN | -Inf | NaN |
-| else | RNE | RNE | RNE | RNE |
+| x                 | E4M3FN | E4M3FNUZ | E5M2 | E5M2FNUZ |
+| ----------------- | ------ | -------- | ---- | -------- |
+| 0                 | 0      | 0        | 0    | 0        |
+| -0                | -0     | 0        | -0   | 0        |
+| NaN               | NaN    | NaN      | NaN  | NaN      |
+| -NaN              | -NaN   | NaN      | -NaN | NaN      |
+| Inf               | NaN    | NaN      | Inf  | NaN      |
+| -Inf              | -NaN   | NaN      | -Inf | NaN      |
+| \[x\] > FLT_MAX   | NaN    | NaN      | Inf  | NaN      |
+| \[x\] \< -FLT_MAX | NaN    | NaN      | -Inf | NaN      |
+| else              | RNE    | RNE      | RNE  | RNE      |
+FLOAT8E8M0 type was introduced to enable [Microscaling (MX) formats](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf).
+When casting to FLOAT8E8M0, the rounding behavior can be specified using the `round_mode` and `saturate` attributes.
+The current CUDA behavior is to round up and saturate. Casting negative values to FLOAT8E8M0 gives undefined behavior.
+The following table describes the casting behavior of special values to FLOAT8E8M0 in the two most common cases.
+| x                 | saturate + up | non-saturate + nearest |
+| ----------------- | ------------- | ---------------------  |
+| 0                 | 0             | NaN                    |
+| -0                | Unspecified   | Unspecified            |
+| NaN               | NaN           | NaN                    |
+| Inf               | E8M0_MAX      | NaN                    |
+| x > E8M0_MAX      | E8M0_MAX      | NaN                    |
+| x \< E8M0_MIN     | E8M0_MIN      | NaN                    |
+| x \< 0            | Unspecified   | Unspecified            |
 )DOC";
 ONNX_OPERATOR_SET_SCHEMA(
     Cast,
-    23,
+    24,
     OpSchema()
-        .SetDoc(Cast_ver19_doc)
+        .SetDoc(Cast_ver24_doc)
         .Attr(
             "to",
             "The data type to which the elements of the input tensor are cast. "
@@ -95,10 +113,18 @@ ONNX_OPERATOR_SET_SCHEMA(
             "saturate",
             "The parameter defines how the conversion behaves if an input value is out of "
             "range of the destination type. It only applies for float 8 conversion "
-            "(float8e4m3fn, float8e4m3fnuz, float8e5m2, float8e5m2fnuz). It is true by default. "
-            "All cases are fully described in two tables inserted in the operator description.",
+            "(float8e4m3fn, float8e4m3fnuz, float8e5m2, float8e5m2fnuz, float8e8m0). It is true by default. "
+            "All cases are fully described in the tables inserted in the operator description.",
             AttributeProto::INT,
             static_cast<int64_t>(1))
+        .Attr(
+            "round_mode",
+            "Rounding mode for conversion to float8e8m0. It only applies to casting to float8e8m0 and is `up` by default. "
+            "`up`: round to nearest value away from zero, "
+            "`down`: round to nearest value towards zero, "
+            "`nearest`: round to nearest value and ties round up.",
+            AttributeProto::STRING,
+            std::string("up"))
         .Input(0, "input", "Input tensor to be cast.", "T1", OpSchema::Single, true, 1, OpSchema::Differentiable)
         .Output(
             0,
@@ -112,11 +138,11 @@ ONNX_OPERATOR_SET_SCHEMA(
             OpSchema::Differentiable)
         .TypeConstraint(
             "T1",
-            OpSchema::all_non_complex_tensor_types_ir11(),
+            OpSchema::all_non_complex_tensor_types_ir12(),
             "Constrain input types. Casting from complex is not supported.")
         .TypeConstraint(
             "T2",
-            OpSchema::all_non_complex_tensor_types_ir11(),
+            OpSchema::all_non_complex_tensor_types_ir12(),
             "Constrain output types. Casting to complex is not supported.")
         .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
           propagateElemTypeFromAttributeToOutput(ctx, "to", 0);
@@ -128,7 +154,7 @@ ONNX_OPERATOR_SET_SCHEMA(
           PropagateShapeDataFromInputToOutput(ctx, 0);
         }));
-static const char* CastLike_ver21_doc = R"DOC(
+static const char* CastLike_ver24_doc = R"DOC(
 The operator casts the elements of a given input tensor (the first input) to
 the same data type as the elements of the second input tensor.
 See documentation of the Cast operator for further details.
@@ -136,17 +162,26 @@ See documentation of the Cast operator for further details.
 ONNX_OPERATOR_SET_SCHEMA(
     CastLike,
-    23,
+    24,
     OpSchema()
-        .SetDoc(CastLike_ver21_doc)
+        .SetDoc(CastLike_ver24_doc)
         .Attr(
             "saturate",
             "The parameter defines how the conversion behaves if an input value is out of "
             "range of the destination type. It only applies for float 8 conversion "
-            "(float8e4m3fn, float8e4m3fnuz, float8e5m2, float8e5m2fnuz). It is true by default. "
+            "(float8e4m3fn, float8e4m3fnuz, float8e5m2, float8e5m2fnuz, float8e8m0). It is true by default. "
             "Please refer to operator Cast description for further details.",
             AttributeProto::INT,
             static_cast<int64_t>(1))
+        .Attr(
+            "round_mode",
+            "Rounding mode for conversion to float8e8m0. It only applies to casting to float8e8m0 and is `up` by default. "
+            "`up`: round to nearest value away from zero, "
+            "`down`: round to nearest value towards zero, "
+            "`nearest`: round to nearest value and ties round up. "
+            "Please refer to operator Cast description for further details.",
+            AttributeProto::STRING,
+            std::string("up"))
         .Input(0, "input", "Input tensor to be cast.", "T1", OpSchema::Single, true, 1, OpSchema::Differentiable)
         .Input(
             1,
@@ -168,11 +203,11 @@ ONNX_OPERATOR_SET_SCHEMA(
             OpSchema::Differentiable)
         .TypeConstraint(
             "T1",
-            OpSchema::all_non_complex_tensor_types_ir11(),
+            OpSchema::all_non_complex_tensor_types_ir12(),
             "Constrain input types. Casting from complex is not supported.")
         .TypeConstraint(
             "T2",
-            OpSchema::all_non_complex_tensor_types_ir11(),
+            OpSchema::all_non_complex_tensor_types_ir12(),
             "Constrain output types. Casting to complex is not supported.")
         .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
           propagateElemTypeFromInputToOutput(ctx, 1, 0);
@@ -196,7 +231,7 @@ ONNX_OPERATOR_SET_SCHEMA(
               return true;
             }));
-static const char* Reshape_ver19_doc = R"DOC(
+static const char* Reshape_ver24_doc = R"DOC(
 Reshape the input tensor similar to numpy.reshape.
 First input is the data tensor, second input is a shape tensor which specifies the output shape. It outputs the reshaped tensor.
 At most one dimension of the new shape can be -1. In this case, the value is
@@ -214,9 +249,9 @@ to -1 cannot be determined uniquely.
 ONNX_OPERATOR_SET_SCHEMA(
     Reshape,
-    23,
+    24,
     OpSchema()
-        .SetDoc(Reshape_ver19_doc)
+        .SetDoc(Reshape_ver24_doc)
         .Attr(
             "allowzero",
             "(Optional) By default, when any value in the 'shape' input is equal to zero "
@@ -236,7 +271,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             1,
             OpSchema::NonDifferentiable)
         .Output(0, "reshaped", "Reshaped data.", "T", OpSchema::Single, true, 1, OpSchema::Differentiable)
-        .TypeConstraint("T", OpSchema::all_tensor_types_ir11(), "Constrain input and output types to all tensor types.")
+        .TypeConstraint("T", OpSchema::all_tensor_types_ir12(), "Constrain input and output types to all tensor types.")
         .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
           // Type inference
           propagateElemTypeFromInputToOutput(ctx, 0, 0);
@@ -269,7 +304,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             auto* new_dim = outputShape->add_dim();
             if (targetShapeProto.dim(i).has_dim_param()) {
               // There is a tricky edge case here. It is possible that the value of
-              // symbolic dim can be -1 or 0 at runtime. In that case simply propgating this
+              // symbolic dim can be -1 or 0 at runtime. In that case simply propagating this
               // symbol can be erroneous. This should be a very rare scenario and in such a
               // case an option is to turn off data propagation during shape inference.
               new_dim->set_dim_param(targetShapeProto.dim(i).dim_param());
@@ -360,18 +395,18 @@ ONNX_OPERATOR_SET_SCHEMA(
           }
         }));
-static const char* Shape_ver19_doc = R"DOC(
+static const char* Shape_ver24_doc = R"DOC(
 Takes a tensor as input and outputs an 1D int64 tensor containing the shape of the input tensor.
 Optional attributes start and end can be used to compute a slice of the input tensor's shape.
 If start axis is omitted, the slice starts from axis 0.
 The end axis, if specified, is exclusive (and the returned value will not include the size of that axis).
 If the end axis is omitted, the axes upto the last one will be included.
 Negative axes indicate counting back from the last axis.
-Note that axes will be clamped to the range [0, r-1], where r is the
+Note that axes will be clamped to the range [0, r], where r is the
 rank of the input tensor if they are out-of-range (after adding r in the case of
 negative axis). Thus, specifying any end value > r is equivalent to specifying an end
 value of r, and specifying any start value < -r is equivalent to specifying a start
-value of 0.
+value of 0. If start > end, the result will be an empty shape.
 Examples:
@@ -403,9 +438,9 @@ Output: [3]
 ONNX_OPERATOR_SET_SCHEMA(
     Shape,
-    23,
+    24,
     OpSchema()
-        .SetDoc(Shape_ver19_doc)
+        .SetDoc(Shape_ver24_doc)
         .Input(0, "data", "An input tensor.", "T", OpSchema::Single, true, 1, OpSchema::NonDifferentiable)
         .Output(0, "shape", "Shape of the input tensor", "T1", OpSchema::Single, true, 1, OpSchema::NonDifferentiable)
         .Attr(
@@ -421,7 +456,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             "If omitted, sizes of all axes upto (including) the last one will be included.",
             AttributeProto::INT,
             OPTIONAL_VALUE)
-        .TypeConstraint("T", OpSchema::all_tensor_types_ir11(), "Input tensor can be of arbitrary type.")
+        .TypeConstraint("T", OpSchema::all_tensor_types_ir12(), "Input tensor can be of arbitrary type.")
         .TypeConstraint("T1", {"tensor(int64)"}, "Constrain output to int64 tensor.")
         .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
           ctx.getOutputType(0)->mutable_tensor_type()->set_elem_type(TensorProto::INT64);
@@ -463,15 +498,15 @@ ONNX_OPERATOR_SET_SCHEMA(
           }
         }));
-static const char* Size_ver19_doc = R"DOC(
+static const char* Size_ver24_doc = R"DOC(
 Takes a tensor as input and outputs a int64 scalar that equals to the total number of elements of the input tensor.
 )DOC";
 ONNX_OPERATOR_SET_SCHEMA(
     Size,
-    23,
+    24,
     OpSchema()
-        .SetDoc(Size_ver19_doc)
+        .SetDoc(Size_ver24_doc)
         .Input(0, "data", "An input tensor.", "T", OpSchema::Single, true, 1, OpSchema::NonDifferentiable)
         .Output(
             0,
@@ -482,7 +517,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             true,
             1,
             OpSchema::NonDifferentiable)
-        .TypeConstraint("T", OpSchema::all_tensor_types_ir11(), "Input tensor can be of arbitrary type.")
+        .TypeConstraint("T", OpSchema::all_tensor_types_ir12(), "Input tensor can be of arbitrary type.")
         .TypeConstraint("T1", {"tensor(int64)"}, "Constrain output to int64 tensor, which should be a scalar though.")
         .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
           ctx.getOutputType(0)->mutable_tensor_type()->set_elem_type(TensorProto::INT64);
@@ -1058,7 +1093,7 @@ ONNX_OPERATOR_SET_SCHEMA(
           }
         }));
-static const char* Transpose_ver13_doc = R"DOC(
+static const char* Transpose_ver24_doc = R"DOC(
 Transpose the input tensor similar to numpy.transpose. For example, when
 perm=(1, 0, 2), given an input tensor of shape (1, 2, 3), the output shape
 will be (2, 1, 3).
@@ -1066,9 +1101,9 @@ will be (2, 1, 3).
 ONNX_OPERATOR_SET_SCHEMA(
     Transpose,
-    23,
+    24,
     OpSchema()
-        .SetDoc(Transpose_ver13_doc)
+        .SetDoc(Transpose_ver24_doc)
         .Attr(
             "perm",
             "A list of integers. By default, reverse the dimensions, "
@@ -1078,7 +1113,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             OPTIONAL_VALUE)
         .Input(0, "data", "An input tensor.", "T", OpSchema::Single, true, 1, OpSchema::Differentiable)
         .Output(0, "transposed", "Transposed output.", "T", OpSchema::Single, true, 1, OpSchema::Differentiable)
-        .TypeConstraint("T", OpSchema::all_tensor_types_ir11(), "Constrain input and output types to all tensor types.")
+        .TypeConstraint("T", OpSchema::all_tensor_types_ir12(), "Constrain input and output types to all tensor types.")
         .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
           propagateElemTypeFromInputToOutput(ctx, 0, 0);
           if (!hasNInputShapes(ctx, 1)) {
@@ -1700,7 +1735,7 @@ ONNX_OPERATOR_SET_SCHEMA(
           }
         }));
-static const char* Squeeze_ver13_doc = R"DOC(
+static const char* Squeeze_ver24_doc = R"DOC(
 Remove single-dimensional entries from the shape of a tensor.
 Takes an input `axes` with a list of axes to squeeze.
 If `axes` is not provided, all the single dimensions will be removed from
@@ -1709,9 +1744,9 @@ the shape. If an axis is selected with shape entry not equal to one, an error is
 ONNX_OPERATOR_SET_SCHEMA(
     Squeeze,
-    23,
+    24,
     OpSchema()
-        .SetDoc(Squeeze_ver13_doc)
+        .SetDoc(Squeeze_ver24_doc)
         .Input(
             0,
             "data",
@@ -1724,7 +1759,7 @@ ONNX_OPERATOR_SET_SCHEMA(
         .Input(
             1,
             "axes",
-            "List of integers indicating the dimensions to squeeze. Negative value means counting dimensions "
+            "1D tensor of integers indicating the dimensions to squeeze. Negative value means counting dimensions "
             "from the back. Accepted range is [-r, r-1] where r = rank(data).",
             "tensor(int64)",
             OpSchema::Optional,
@@ -1742,8 +1777,8 @@ ONNX_OPERATOR_SET_SCHEMA(
             OpSchema::Differentiable)
         .TypeConstraint(
             "T",
-            OpSchema::all_tensor_types_ir11(),
-            "Constrain input and output types to all tensor types up to IRv11.")
+            OpSchema::all_tensor_types_ir12(),
+            "Constrain input and output types to all tensor types up to IRv12.")
         .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
           propagateElemTypeFromInputToOutput(ctx, 0, 0);
           if (!hasNInputShapes(ctx, 1)) {
@@ -1803,7 +1838,7 @@ ONNX_OPERATOR_SET_SCHEMA(
           PropagateShapeDataFromInputToOutput(ctx, 0);
         }));
-static const char* Unsqueeze_ver13_doc = R"DOC(
+static const char* Unsqueeze_ver24_doc = R"DOC(
 Insert single-dimensional entries to the shape of an input tensor (`data`).
 Takes one required input `axes` - which contains a list of dimension indices and this operator will insert a dimension of value `1` into the corresponding index of the output tensor (`expanded`).
@@ -1818,14 +1853,14 @@ The order of values in `axes` does not matter and can come in any order.
 ONNX_OPERATOR_SET_SCHEMA(
     Unsqueeze,
-    23,
+    24,
     OpSchema()
-        .SetDoc(Unsqueeze_ver13_doc)
+        .SetDoc(Unsqueeze_ver24_doc)
         .Input(0, "data", "Original tensor", "T", OpSchema::Single, true, 1, OpSchema::Differentiable)
         .Input(
             1,
             "axes",
-            "List of integers indicating the dimensions to be inserted. Negative value means counting dimensions "
+            "1D tensor of integers indicating the dimensions to be inserted. Negative value means counting dimensions "
             "from the back. Accepted range is [-r, r-1] where r = rank(expanded).",
             "tensor(int64)",
             OpSchema::Single,
@@ -1843,8 +1878,8 @@ ONNX_OPERATOR_SET_SCHEMA(
             OpSchema::Differentiable)
         .TypeConstraint(
             "T",
-            OpSchema::all_tensor_types_ir11(),
-            "Constrain input and output types to all tensor types up to IRv11.")
+            OpSchema::all_tensor_types_ir12(),
+            "Constrain input and output types to all tensor types up to IRv12.")
         .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
           propagateElemTypeFromInputToOutput(ctx, 0, 0);
           if (!hasNInputShapes(ctx, 1)) {
@@ -2205,13 +2240,13 @@ If `keep_aspect_ratio_policy` is `"stretch"`, the original aspect ratio is disre
 If `keep_aspect_ratio_policy` is `"not_larger"`, the sizes are adjusted so that no extent of the output is larger than the specified size, while keeping the original aspect ratio:
 ```
 scale = Min(sizes[i] / in_size[d])
-out_size[d] = round_int(scale * in_size[i])
+out_size[d] = round_int(scale * in_size[d])
 ```
 If `keep_aspect_ratio_policy` is `"not_smaller"`, the sizes are adjusted so that no extent of the output is smaller than the specified size, while keeping the original aspect ratio:
 ```
 scale = Max(sizes[i] / in_size[d])
-out_size[d] = round_int(scale * in_size[i])
+out_size[d] = round_int(scale * in_size[d])
 ```
 For non-resizable axes (those not specified in `axes`), the output size will be equal to the input size.
@@ -2702,7 +2737,7 @@ ONNX_OPERATOR_SET_SCHEMA(
 ONNX_OPERATOR_SET_SCHEMA(
     Identity,
-    23,
+    24,
     OpSchema()
         .SetDoc("Identity operator")
         .Input(0, "input", "Input tensor", "V", OpSchema::Single, true, 1, OpSchema::Differentiable)
@@ -2710,7 +2745,7 @@ ONNX_OPERATOR_SET_SCHEMA(
         .TypeConstraint(
             "V",
             []() {
-              auto t = OpSchema::all_tensor_types_ir11();
+              auto t = OpSchema::all_tensor_types_ir12();
               auto s = OpSchema::all_tensor_sequence_types();
               auto o = OpSchema::all_optional_types();
               t.insert(t.end(), s.begin(), s.end());
@@ -3570,7 +3605,7 @@ ONNX_OPERATOR_SET_SCHEMA(
           }
         }));
-static const char* Pad_ver19_doc = R"DOC(
+static const char* Pad_ver24_doc = R"DOC(
 Given a tensor containing the data to be padded (`data`), a tensor containing the number of start and end pad values for axis (`pads`), (optionally) a `mode`, and (optionally) `constant_value`,
 a padded tensor (`output`) is generated.
@@ -3675,12 +3710,12 @@ output = [
 ONNX_OPERATOR_SET_SCHEMA(
     Pad,
-    23,
+    24,
     OpSchema().FillUsing(PadDocGenerator(
-        Pad_ver19_doc,
+        Pad_ver24_doc,
         "Supported modes: `constant`(default), `reflect`, `edge`, `wrap`",
-        OpSchema::all_tensor_types_ir11(),
-        "Constrain input and output types to all tensor types up to IRv11.")));
+        OpSchema::all_tensor_types_ir12(),
+        "Constrain input and output types to all tensor types up to IRv12.")));
 static const char* Trilu_ver14_doc = R"DOC(
 Given a 2-D matrix or batches of 2-D matrices, returns the upper or lower triangular part of the tensor(s).
@@ -3917,4 +3952,142 @@ ONNX_OPERATOR_SET_SCHEMA(
           return true;
         }));
+static const char* TensorScatter_ver24_doc = R"DOC(
+TensorScatter is a generic tensor update operation, motivated by the requirements for KV cache updates for Attention
+ops commonly found in LLMs. It is a functional operation that models an in-place update to a KV cache buffer.
+The past and present cache tensors have the same shape (batch_size, D1, D2, ..., max_sequence_length, ..., Dn), with
+the sequence dimension (indicated by the `axis` attribute) being max_sequence_length, so the sizes of these tensors do
+not need to grow between iterations. The `update` tensor's shape only differs from the cache tensors in the sequence
+dimension: (batch_size, D1, D2, ..., sequence_length, ..., Dn), where sequence_length <= max_sequence_length.
+The optional `write_indices` input indicates the write index for each sample in the batch, assumed to be zero
+if not provided. When the `mode` attribute is set to "circular", the write index is modulo max_sequence_length.
+The operation can be described using the following pseudocode:
+```
+for prefix_idx in np.ndindex(past_cache.shape[:axis]):
+    batch_idx = prefix_idx[0]
+    for sequence_idx in range(sequence_length):
+        cache_idx = (*prefix_idx, write_indices[batch_idx] + sequence_idx)
+        if mode == "circular":
+            cache_idx = tuple(np.mod(np.asarray(cache_idx), max_sequence_length))
+        update_idx = (*prefix_idx, sequence_idx)
+        present_cache[cache_idx] = update[update_idx]
+```
+During the prefill phase of attention, only the first two inputs are needed. During the decode phase, `write_indices`
+is also needed so that the incoming key or value update can be appended after the last valid token for each sample
+in the batch.
+)DOC";
+ONNX_OPERATOR_SET_SCHEMA(
+    TensorScatter,
+    24,
+    OpSchema()
+        .SetDoc(TensorScatter_ver24_doc)
+        .Attr(
+            "axis",
+            "Sequence dimension of the `past_cache` and `update` tensors. It cannot be 0 (the batch dimension). Default is -2.",
+            AttributeProto::INT,
+            static_cast<int64_t>(-2))
+        .Attr(
+            "mode",
+            "Write mode of cache update. Supported modes include `linear` and `circular`. `linear` mode requires "
+            "write_indices+sequence_length<=max_sequence_length. For `circular` mode, the updates happen in "
+            "wrap-around fashion, ie, the update index is modulo `max_sequence_length`",
+            AttributeProto::STRING,
+            std::string("linear"))
+        .Input(
+            0,
+            "past_cache",
+            "Past state cache for key or value with shape `(batch_size, D1, D2, ..., max_sequence_length, ..., Dn)`.",
+            "T",
+            OpSchema::Single,
+            true,
+            1,
+            OpSchema::Differentiable)
+        .Input(
+            1,
+            "update",
+            "New update tensor with shape `(batch_size, D1, D2, ..., sequence_length, ..., Dn)`.",
+            "T",
+            OpSchema::Single,
+            true,
+            1,
+            OpSchema::Differentiable)
+        .Input(
+            2,
+            "write_indices",
+            "Write indices for the incoming update tensor in the cache. Shape is `(batch_size,)`. Assumed to be all zeros if not provided.",
+            "tensor(int64)",
+            OpSchema::Optional,
+            true,
+            1,
+            OpSchema::NonDifferentiable)
+        .Output(
+            0,
+            "present_cache",
+            "Updated cache. Same shape as `past_cache`.",
+            "T",
+            OpSchema::Single,
+            true,
+            1,
+            OpSchema::Differentiable)
+        .TypeConstraint("T", OpSchema::all_tensor_types_ir12(), "Constrain input and output types to any tensor type.")
+        .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+          propagateElemTypeFromInputToOutput(ctx, 0, 0);
+          // Check that updates and past_cache have the same rank
+          if (hasInputShape(ctx, 0) && hasInputShape(ctx, 1)) {
+            const auto& past_cache_shape = getInputShape(ctx, 0);
+            const auto& update_shape = getInputShape(ctx, 1);
+            const int past_cache_rank = static_cast<int>(past_cache_shape.dim_size());
+            const int update_rank = static_cast<int>(update_shape.dim_size());
+            if (past_cache_rank != update_rank) {
+              fail_shape_inference(
+                  "past_cache and update must have the same rank. past_cache_rank=",
+                  past_cache_rank,
+                  ", update_rank=",
+                  update_rank);
+            }
+            // Check that all dimensions match except the axis dimension
+            // Read axis from attribute, default is -2
+            auto axis_attr = ctx.getAttribute("axis");
+            int64_t axis = axis_attr ? axis_attr->i() : -2;
+            // Handle negative axis
+            if (axis < 0) {
+              axis += past_cache_rank;
+            }
+            // Validate axis is within bounds
+            if (axis <= 0 || axis >= past_cache_rank) {
+              fail_shape_inference("axis ", axis, " is out of range for rank ", past_cache_rank);
+            }
+            for (int i = 0; i < past_cache_rank; ++i) {
+              if (i != axis) {
+                // All dimensions except axis should match
+                Dim dim;
+                unifyInputDim(ctx, 0, i, dim);
+                unifyInputDim(ctx, 1, i, dim);
+              }
+            }
+          }
+          // Check that write_indices is of shape (batch_size,)
+          if (hasInputShape(ctx, 2)) {
+            checkInputRank(ctx, 2, 1);
+            Dim batch_size;
+            unifyInputDim(ctx, 0, 0, batch_size);
+            unifyInputDim(ctx, 2, 0, batch_size);
+          }
+          if (hasNInputShapes(ctx, 1)) {
+            propagateShapeFromInputToOutput(ctx, 0, 0);
+          }
+        }));
 } // namespace ONNX_NAMESPACE