Skip to content

Commit 4be24b0

Browse files
BlackPeter13BlackPeter13
authored
Merge branch 'tensorflow:master' into master
2 parents 14c267f + 0c88383 commit 4be24b0

File tree

14 files changed

+328
-5
lines changed

14 files changed

+328
-5
lines changed

‎tensorflow/compiler/jit/mark_for_compilation_pass.cc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2145,6 +2145,7 @@ absl::flat_hash_set<string> GetKnownXLAAllowlistOp() {
21452145
"Roll",
21462146
"ScatterNd",
21472147
"SegmentSumV2",
2148+
"SegmentProdV2",
21482149
"SelfAdjointEigV2",
21492150
"SoftmaxCrossEntropyWithLogits",
21502151
"SpaceToBatch",

‎tensorflow/compiler/mlir/g3doc/_includes/tf_passes.md‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1116,6 +1116,18 @@ tf_device.replicate([%0, %1] as %ri: tensor<*x!tf_type.resource>) {n = 2 : i32}
11161116
tf_device.return
11171117
}
11181118
```
1119+
### `-tf-replicate-tensor-list-init-ops`: Replicate TensorList init ops for correct shape assignments in shape inference
1120+
If we pass same TensorList to a while op as multiple arguments or just use
1121+
the same TensorList at multiple places and assign different
1122+
TensorListSetItem to elements of TensorList, the shape inference is then
1123+
unable to identify the Shape of these args and thus the input TensorList
1124+
shape is unidentifiable.
1125+
All of these args are supposed to be independent and not related to original
1126+
creation of TensorList.
1127+
1128+
This pass will create multiple instances of TensorList for each arg of the
1129+
while op and each use and thus there will be not a conflict in resolving the
1130+
shape of these different inputs.
11191131
### `-tf-replicate-to-island`: Lowers device replicate to executor islands
11201132

11211133
#### Options

‎tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15241,6 +15241,64 @@ has size `k`, the number of segments.}]>:$output
1524115241
TF_DerivedOperandTypeAttr Tindices = TF_DerivedOperandTypeAttr<1>;
1524215242
}
1524315243

15244+
def TF_SegmentProdV2Op : TF_Op<"SegmentProdV2", [Pure]> {
15245+
let summary = "Computes the product along segments of a tensor.";
15246+
15247+
let description = [{
15248+
Read
15249+
[the section on segmentation](https://tensorflow.org/api_docs/python/tf/math#Segmentation)
15250+
for an explanation of segments.
15251+
15252+
Computes a tensor such that
15253+
\\(output_i = \prod_j data_j\\) where the product is over `j` such
15254+
that `segment_ids[j] == i`.
15255+
15256+
If the product is empty for a given segment ID `i`, `output[i] = 1`.
15257+
15258+
Note: That this op is currently only supported with jit_compile=True.
15259+
15260+
The only difference with SegmentProd is the additional input `num_segments`.
15261+
This helps in evaluating the output shape in compile time.
15262+
`num_segments` should be consistent with segment_ids.
15263+
e.g. Max(segment_ids) - 1 should be equal to `num_segments` for a 1-d segment_ids
15264+
With inconsistent num_segments, the op still runs. only difference is,
15265+
the output takes the size of num_segments irrespective of size of segment_ids and data.
15266+
for num_segments less than expected output size, the last elements are ignored
15267+
for num_segments more than the expected output size, last elements are assigned 1.
15268+
15269+
For example:
15270+
15271+
>>> @tf.function(jit_compile=True)
15272+
... def test(c):
15273+
... return tf.raw_ops.SegmentProdV2(data=c, segment_ids=tf.constant([0, 0, 1]), num_segments=2)
15274+
>>> c = tf.constant([[1,2,3,4], [4, 3, 2, 1], [5,6,7,8]])
15275+
>>> test(c).numpy()
15276+
array([[4, 6, 6, 4],
15277+
[5, 6, 7, 8]], dtype=int32)
15278+
}];
15279+
15280+
let arguments = (ins
15281+
TF_NumberTensor:$data,
15282+
Arg<TF_I32OrI64Tensor, [{A 1-D tensor whose size is equal to the size of `data`'s
15283+
first dimension. Values should be sorted and can be repeated.
15284+
The values must be less than `num_segments`.
15285+
15286+
Caution: The values are always validated to be sorted on CPU, never validated
15287+
on GPU.}]>:$segment_ids,
15288+
TF_I32OrI64Tensor:$num_segments
15289+
);
15290+
15291+
let results = (outs
15292+
Res<TF_NumberTensor, [{Has same shape as data, except for the first `segment_ids.rank`
15293+
dimensions, which are replaced with a single dimensionw which has size
15294+
`num_segments`.}]>:$output
15295+
);
15296+
15297+
TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
15298+
TF_DerivedOperandTypeAttr Tindices = TF_DerivedOperandTypeAttr<1>;
15299+
TF_DerivedOperandTypeAttr Tnumsegments = TF_DerivedOperandTypeAttr<2>;
15300+
}
15301+
1524415302
def TF_SegmentSumOp : TF_Op<"SegmentSum", [Pure]> {
1524515303
let summary = "Computes the sum along segments of a tensor.";
1524615304

‎tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,7 @@ bool IsOpAllowedTf2XlaFallback(Operation* op) {
232232
TypeID::get<TF::RollOp>(),
233233
TypeID::get<TF::RoundOp>(),
234234
TypeID::get<TF::SegmentSumV2Op>(),
235+
TypeID::get<TF::SegmentProdV2Op>(),
235236
TypeID::get<TF::SelectV2Op>(),
236237
TypeID::get<TF::SelfAdjointEigV2Op>(),
237238
TypeID::get<TF::SeluGradOp>(),

‎tensorflow/compiler/tests/segment_reduction_ops_test.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,10 @@ def _segmentSumV2(self, data, indices, num_segments):
4545
return self._segmentReduction(math_ops.segment_sum_v2, data, indices,
4646
num_segments)
4747

48+
def _segmentProdV2(self, data, indices, num_segments):
49+
return self._segmentReduction(math_ops.segment_prod_v2, data, indices,
50+
num_segments)
51+
4852
def _unsortedSegmentProd(self, data, indices, num_segments):
4953
return self._segmentReduction(math_ops.unsorted_segment_prod, data, indices,
5054
num_segments)
@@ -65,6 +69,30 @@ def testSegmentSum(self):
6569
np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
6670
np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 4))
6771

72+
def testSegmentProd(self):
73+
for dtype in self.numeric_types:
74+
self.assertAllClose(
75+
np.array([0, 1, 2, 60], dtype=dtype),
76+
self._segmentProdV2(
77+
np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
78+
np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 4))
79+
80+
def testSegmentProdNumSegmentsLess(self):
81+
for dtype in self.numeric_types:
82+
self.assertAllClose(
83+
np.array([0, 1, 2], dtype=dtype),
84+
self._segmentProdV2(
85+
np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
86+
np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 3))
87+
88+
def testSegmentProdNumSegmentsMore(self):
89+
for dtype in self.numeric_types:
90+
self.assertAllClose(
91+
np.array([0, 1, 2, 60, 1], dtype=dtype),
92+
self._segmentProdV2(
93+
np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
94+
np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 5))
95+
6896
def testUnsortedSegmentSum0DIndices1DData(self):
6997
for dtype in self.numeric_types:
7098
self.assertAllClose(

‎tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -150,10 +150,11 @@ REGISTER_XLA_OP(
150150
Name("UnsortedSegmentSum").CompileTimeConstantInput("num_segments"),
151151
SegmentSum</*indices_are_sorted=*/false>);
152152

153-
class UnsortedSegmentProd : public SegmentReduce {
153+
template <bool indices_are_sorted>
154+
class SegmentProd : public SegmentReduce {
154155
public:
155-
explicit UnsortedSegmentProd(OpKernelConstruction* ctx)
156-
: SegmentReduce(ctx, /*indices_are_sorted=*/false) {}
156+
explicit SegmentProd(OpKernelConstruction* ctx)
157+
: SegmentReduce(ctx, indices_are_sorted) {}
157158

158159
xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
159160
return xla::One(builder, type_);
@@ -163,7 +164,9 @@ class UnsortedSegmentProd : public SegmentReduce {
163164

164165
REGISTER_XLA_OP(
165166
Name("UnsortedSegmentProd").CompileTimeConstantInput("num_segments"),
166-
UnsortedSegmentProd);
167+
SegmentProd</*indices_are_sorted=*/false>);
168+
REGISTER_XLA_OP(Name("SegmentProdV2").CompileTimeConstantInput("num_segments"),
169+
SegmentProd</*indices_are_sorted=*/true>);
167170

168171
class UnsortedSegmentMin : public SegmentReduce {
169172
public:

‎tensorflow/core/api_def/base_api/api_def_SegmentProdV2.pbtxt‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
op {
2+
graph_op_name: "SegmentProdV2"
3+
in_arg {
4+
name: "segment_ids"
5+
description: <<END
6+
A 1-D tensor whose size is equal to the size of `data`'s
7+
first dimension. Values should be sorted and can be repeated.
8+
The values must be less than `num_segments`.
9+
10+
Caution: The values are always validated to be sorted on CPU, never validated
11+
on GPU.
12+
END
13+
}
14+
out_arg {
15+
name: "output"
16+
description: <<END
17+
Has same shape as data, except for the first `segment_ids.rank`
18+
dimensions, which are replaced with a single dimensionw which has size
19+
`num_segments`.
20+
END
21+
}
22+
summary: "Computes the product along segments of a tensor."
23+
description: <<END
24+
Read
25+
[the section on segmentation](https://tensorflow.org/api_docs/python/tf/math#Segmentation)
26+
for an explanation of segments.
27+
28+
Computes a tensor such that
29+
\\(output_i = \prod_j data_j\\) where the product is over `j` such
30+
that `segment_ids[j] == i`.
31+
32+
If the product is empty for a given segment ID `i`, `output[i] = 1`.
33+
34+
Note: That this op is currently only supported with jit_compile=True.
35+
36+
The only difference with SegmentProd is the additional input `num_segments`.
37+
This helps in evaluating the output shape in compile time.
38+
`num_segments` should be consistent with segment_ids.
39+
e.g. Max(segment_ids) - 1 should be equal to `num_segments` for a 1-d segment_ids
40+
With inconsistent num_segments, the op still runs. only difference is,
41+
the output takes the size of num_segments irrespective of size of segment_ids and data.
42+
for num_segments less than expected output size, the last elements are ignored
43+
for num_segments more than the expected output size, last elements are assigned 1.
44+
45+
For example:
46+
47+
>>> @tf.function(jit_compile=True)
48+
... def test(c):
49+
... return tf.raw_ops.SegmentProdV2(data=c, segment_ids=tf.constant([0, 0, 1]), num_segments=2)
50+
>>> c = tf.constant([[1,2,3,4], [4, 3, 2, 1], [5,6,7,8]])
51+
>>> test(c).numpy()
52+
array([[4, 6, 6, 4],
53+
[5, 6, 7, 8]], dtype=int32)
54+
55+
END
56+
}

‎tensorflow/core/api_def/python_api/api_def_SegmentProdV2.pbtxt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
op {
2+
graph_op_name: "SegmentProdV2"
3+
visibility: HIDDEN
4+
}

‎tensorflow/core/ops/compat/ops_history_v2/SegmentProdV2.pbtxt‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
op {
2+
name: "SegmentProdV2"
3+
input_arg {
4+
name: "data"
5+
type_attr: "T"
6+
}
7+
input_arg {
8+
name: "segment_ids"
9+
type_attr: "Tindices"
10+
}
11+
input_arg {
12+
name: "num_segments"
13+
type_attr: "Tnumsegments"
14+
}
15+
output_arg {
16+
name: "output"
17+
type_attr: "T"
18+
}
19+
attr {
20+
name: "T"
21+
type: "type"
22+
allowed_values {
23+
list {
24+
type: DT_FLOAT
25+
type: DT_DOUBLE
26+
type: DT_INT32
27+
type: DT_UINT8
28+
type: DT_INT16
29+
type: DT_INT8
30+
type: DT_COMPLEX64
31+
type: DT_INT64
32+
type: DT_QINT8
33+
type: DT_QUINT8
34+
type: DT_QINT32
35+
type: DT_BFLOAT16
36+
type: DT_QINT16
37+
type: DT_QUINT16
38+
type: DT_UINT16
39+
type: DT_COMPLEX128
40+
type: DT_HALF
41+
type: DT_UINT32
42+
type: DT_UINT64
43+
}
44+
}
45+
}
46+
attr {
47+
name: "Tindices"
48+
type: "type"
49+
allowed_values {
50+
list {
51+
type: DT_INT32
52+
type: DT_INT64
53+
}
54+
}
55+
}
56+
attr {
57+
name: "Tnumsegments"
58+
type: "type"
59+
default_value {
60+
type: DT_INT32
61+
}
62+
allowed_values {
63+
list {
64+
type: DT_INT32
65+
type: DT_INT64
66+
}
67+
}
68+
}
69+
}

‎tensorflow/core/ops/math_ops.cc‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
15+
#include <limits>
16+
#include <vector>
1517

1618
#include "tensorflow/core/framework/common_shape_fns.h"
1719
#include "tensorflow/core/framework/numeric_op.h"
@@ -1278,7 +1280,7 @@ REGISTER_OP("SegmentSum")
12781280
.Attr("Tindices: {int32,int64}")
12791281
.SetShapeFn(SegmentReductionShapeFn);
12801282

1281-
// TODO(hinsu): Introduce Segment{Prod,Min,Max}V2 ops, similarly.
1283+
// TODO(hinsu): Introduce Segment{Min,Max}V2 ops, similarly.
12821284
REGISTER_OP("SegmentSumV2")
12831285
.Input("data: T")
12841286
.Input("segment_ids: Tindices")
@@ -1305,6 +1307,16 @@ REGISTER_OP("SegmentProd")
13051307
.Attr("Tindices: {int32,int64}")
13061308
.SetShapeFn(SegmentReductionShapeFn);
13071309

1310+
REGISTER_OP("SegmentProdV2")
1311+
.Input("data: T")
1312+
.Input("segment_ids: Tindices")
1313+
.Input("num_segments: Tnumsegments")
1314+
.Output("output: T")
1315+
.Attr("T: numbertype")
1316+
.Attr("Tindices: {int32,int64}")
1317+
.Attr("Tnumsegments: {int32,int64} = DT_INT32")
1318+
.SetShapeFn(shape_inference::SegmentReductionWithNumSegmentsShapeFn);
1319+
13081320
REGISTER_OP("SegmentMin")
13091321
.Input("data: T")
13101322
.Input("segment_ids: Tindices")

0 commit comments

Comments
 (0)