experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Don't allow reductions to have repeated reduction dimensions

Browse Source 
And fix the TF->XLA bridge to have the same check.

PiperOrigin-RevId: 319319777
Change-Id: Ia2c515582fbee4c7b1c190fad416c23a5c7d5678
This commit is contained in:
Sanjoy Das 2020-07-01 16:40:38 -07:00 committed by TensorFlower Gardener
parent 7c224307e5
commit f962228824
4 changed files with 35 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
tensorflow/compiler/tests/reduce_ops_test.py
View File

@ -168,6 +168,17 @@ class ReduceOpsTest(xla_test.XLATestCase, parameterized.TestCase):
self._testReduction(math_ops.reduce_any, np.any, np.bool, self.BOOL_DATA,
index_dtype)
def testReduceSumWithDuplicateAxes(self, index_dtype):
with self.session() as sess:
with self.test_scope():
a = array_ops.placeholder(np.float32)
index = array_ops.placeholder(np.int32)
out = math_ops.reduce_sum(a, index)
with self.assertRaisesWithPredicateMatch(
errors_impl.InvalidArgumentError,
'Axes contains duplicate dimension'):
sess.run(out, {a: [10, 20, 30], index: [0, 0]})
class ReduceOpPrecisionTest(xla_test.XLATestCase):

5
tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc
View File

@ -80,6 +80,11 @@ void XlaReductionOp::Compile(XlaOpKernelContext* ctx) {
" for input with ", data_shape.dims(),
" dimension(s)"));
index = (index + data_shape.dims()) % data_shape.dims();
OP_REQUIRES(
ctx, !bitmap[index],
errors::InvalidArgument(
"Invalid reduction arguments: Axes contains duplicate dimension: ",
index));
bitmap[index] = true;
xla_axes.push_back(index);
}

10
tensorflow/compiler/xla/service/shape_inference.cc
View File

@ -2127,8 +2127,14 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation,
TF_RETURN_IF_ERROR(VerifyReducerShape(to_apply, init_values, element_types,
num_reduced_args));
std::set<int64> dimensions_to_reduce_set(dimensions_to_reduce.begin(),
dimensions_to_reduce.end());
absl::flat_hash_set<int64> dimensions_to_reduce_set;
for (int64 dim_to_reduce : dimensions_to_reduce) {
if (!dimensions_to_reduce_set.insert(dim_to_reduce).second) {
return InvalidArgument("Duplicate reduction dimension: %d",
dim_to_reduce);
}
}
std::vector<int64> new_dimensions;
std::vector<bool> new_is_dynamic;
for (int i = 0; i < arg.rank(); ++i) {

11
tensorflow/compiler/xla/service/shape_inference_test.cc
View File

@ -1020,6 +1020,17 @@ TEST_F(ReduceShapeInferenceTest, ErrorElementTypeVsApplyType) {
HasSubstr("0-th parameter shape differs"));
}
TEST_F(ReduceShapeInferenceTest, ReduceWithRepeatedReduceDimension) {
ProgramShape to_apply = ShapeUtil::MakeProgramShape({f32_, f32_}, f32_);
Shape arg_shape = ShapeUtil::MakeShape(F32, {5, 3});
auto inferred_status = ShapeInference::InferReduceShape(
{&arg_shape, &f32_},
/*dimensions_to_reduce=*/{0, 0}, to_apply);
EXPECT_FALSE(inferred_status.ok());
EXPECT_THAT(inferred_status.status().error_message(),
HasSubstr("Duplicate reduction dimension: 0"));
}
TEST_F(ShapeInferenceTest, InferSliceShapeRank2) {
Shape matrix_shape = ShapeUtil::MakeShape(F32, {128, 64});
auto inferred_status =