message: Add tf.BoostedTreesBucketize and tf.SparseSegmentSum to TF MLIR ODS.
diff: @@ -1315,6 +1315,26 @@ for dtype in dtype_list:
  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
  }
  +def TF_BoostedTreesBucketizeOp : TF_Op<"BoostedTreesBucketize", [NoSideEffect, SameVariadicOperandSize]> {
  + let summary = "Bucketize each feature based on bucket boundaries.";
  +
  + let description = [{
  +An op that returns a list of float tensors, where each tensor represents the
  +bucketized values for a single feature.
  + }];
  +
  + let arguments = (ins
  + Variadic<TF_Float32Tensor>:$float_values,
  + Variadic<TF_Float32Tensor>:$bucket_boundaries
  + );
  +
  + let results = (outs
  + Variadic<TF_Int32Tensor>:$buckets
  + );
  +
  + TF_DerivedOperandSizeAttr num_features = TF_DerivedOperandSizeAttr<0>;
  +}
  +
  def TF_BroadcastArgsOp : TF_Op<"BroadcastArgs", [NoSideEffect]> {
  let summary = "Return the shape of s0 op s1 with broadcast.";
  @@ -11512,6 +11532,56 @@ See `tf.sparse.segment_sum` for usage examples.
  TF_DerivedOperandTypeAttr Tidx = TF_DerivedOperandTypeAttr<1>;
  }
  +def TF_SparseSegmentSumOp : TF_Op<"SparseSegmentSum", [NoSideEffect]> {
  + let summary = "Computes the sum along sparse segments of a tensor.";
  +
  + let description = [{
  +Read
  +[the section on segmentation](https://tensorflow.org/api_docs/python/tf/math#Segmentation)
  +for an explanation of segments.
  +
  +Like `SegmentSum`, but `segment_ids` can have rank less than `data`'s first
  +dimension, selecting a subset of dimension 0, specified by `indices`.
  +
  +For example:
  +
  +```python
  +c = tf.constant([[1,2,3,4], [-1,-2,-3,-4], [5,6,7,8]])
  +
  +# Select two rows, one segment.
  +tf.sparse_segment_sum(c, tf.constant([0, 1]), tf.constant([0, 0]))
  +# => [[0 0 0 0]]
  +
  +# Select two rows, two segment.
  +tf.sparse_segment_sum(c, tf.constant([0, 1]), tf.constant([0, 1]))
  +# => [[ 1 2 3 4]
  +# [-1 -2 -3 -4]]
  +
  +# Select all rows, two segments.
  +tf.sparse_segment_sum(c, tf.constant([0, 1, 2]), tf.constant([0, 0, 1]))
  +# => [[0 0 0 0]
  +# [5 6 7 8]]
  +
  +# Which is equivalent to:
  +tf.segment_sum(c, tf.constant([0, 0, 1]))
  +```
  + }];
  +
  + let arguments = (ins
  + TF_IntOrFpTensor:$data,
  + TF_I32OrI64Tensor:$indices,
  + TF_I32OrI64Tensor:$segment_ids
  + );
  +
  + let results = (outs
  + TF_IntOrFpTensor:$output
  + );
  +
  + TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
  + TF_DerivedOperandTypeAttr Tsegmentids = TF_DerivedOperandTypeAttr<2>;
  + TF_DerivedOperandTypeAttr Tidx = TF_DerivedOperandTypeAttr<1>;
  +}
  +
  def TF_SparseSoftmaxCrossEntropyWithLogitsOp : TF_Op<"SparseSoftmaxCrossEntropyWithLogits", [NoSideEffect]> {
  let summary = [{
  Computes softmax cross entropy cost and gradients to backpropagate.
  