"true_computation=%TransposeDot, false_computation=%TransposeDot");
}
-TEST_F(HloInstructionTest, StringifyGather) {
+TEST_F(HloInstructionTest, StringifyGather_0) {
Shape input_tensor_shape = ShapeUtil::MakeShape(F32, {50, 49, 48, 47, 46});
Shape gather_indices_tensor_shape =
ShapeUtil::MakeShape(S64, {10, 9, 8, 7, 5});
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7, 8},
/*elided_window_dims=*/{},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 27, 26}));
HloModule module(TestName());
"s64[10,9,8,7,5]{4,3,2,1,0} %gather_indices), "
"output_window_dims={4,5,6,7,8}, elided_window_dims={}, "
"gather_dims_to_operand_dims={0,1,2,3,4}, "
- "window_bounds={30,29,28,27,26}");
+ "index_vector_dim=4, window_bounds={30,29,28,27,26}");
+}
+
+TEST_F(HloInstructionTest, StringifyGather_1) {
+ Shape input_tensor_shape = ShapeUtil::MakeShape(F32, {50, 49, 48, 47, 46});
+ Shape gather_indices_tensor_shape =
+ ShapeUtil::MakeShape(S64, {10, 9, 5, 7, 6});
+ Shape gather_result_shape =
+ ShapeUtil::MakeShape(F32, {10, 9, 7, 6, 30, 29, 28, 27, 26});
+
+ HloComputation::Builder builder("Gather");
+ HloInstruction* input = builder.AddInstruction(
+ HloInstruction::CreateParameter(0, input_tensor_shape, "input_tensor"));
+ HloInstruction* gather_indices =
+ builder.AddInstruction(HloInstruction::CreateParameter(
+ 1, gather_indices_tensor_shape, "gather_indices"));
+
+ HloInstruction* gather_instruction =
+ builder.AddInstruction(HloInstruction::CreateGather(
+ gather_result_shape, input, gather_indices,
+ HloInstruction::MakeGatherDimNumbers(
+ /*output_window_dims=*/{4, 5, 6, 7, 8},
+ /*elided_window_dims=*/{},
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/2),
+ /*window_bounds=*/{30, 29, 28, 27, 26}));
+
+ HloModule module(TestName());
+ module.AddEntryComputation(builder.Build());
+
+ EXPECT_EQ(gather_instruction->ToString(),
+ "%gather = f32[10,9,7,6,30,29,28,27,26]{8,7,6,5,4,3,2,1,0} "
+ "gather(f32[50,49,48,47,46]{4,3,2,1,0} %input_tensor, "
+ "s64[10,9,5,7,6]{4,3,2,1,0} %gather_indices), "
+ "output_window_dims={4,5,6,7,8}, elided_window_dims={}, "
+ "gather_dims_to_operand_dims={0,1,2,3,4}, "
+ "index_vector_dim=2, window_bounds={30,29,28,27,26}");
}
} // namespace
class GatherShapeInferenceTest : public ShapeInferenceTest {
protected:
+ const Shape s64_scalar_ = ShapeUtil::MakeShape(S64, {});
+ const Shape s64_vector_5_ = ShapeUtil::MakeShape(S64, {5});
const Shape s64_vector_32_ = ShapeUtil::MakeShape(S64, {32});
const Shape s64_4d_tensor_10_9_8_7_1_ =
ShapeUtil::MakeShape(S64, {10, 9, 8, 7, 1});
const Shape s64_4d_tensor_10_9_8_7_5_ =
ShapeUtil::MakeShape(S64, {10, 9, 8, 7, 5});
+ const Shape s64_4d_tensor_5_10_9_7_6_ =
+ ShapeUtil::MakeShape(S64, {5, 10, 9, 7, 6});
+ const Shape s64_4d_tensor_10_9_5_7_6_ =
+ ShapeUtil::MakeShape(S64, {10, 9, 5, 7, 6});
const Shape f32_5d_tensor_50_49_48_47_46_ =
ShapeUtil::MakeShape(F32, {50, 49, 48, 47, 46});
const Shape tuple_shape_ = ShapeUtil::MakeTupleShape(
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{0},
/*elided_window_dims=*/{1},
- /*gather_dims_to_operand_dims=*/{1}),
+ /*gather_dims_to_operand_dims=*/{1},
+ /*index_vector_dim=*/1),
/*window_bounds=*/{64, 1}));
EXPECT_TRUE(
ShapeUtil::Equal(gather_shape, ShapeUtil::MakeShape(F32, {64, 32})))
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{1},
/*elided_window_dims=*/{0},
- /*gather_dims_to_operand_dims=*/{0}),
+ /*gather_dims_to_operand_dims=*/{0},
+ /*index_vector_dim=*/1),
/*window_bounds=*/{1, 48}));
EXPECT_TRUE(
ShapeUtil::Equal(gather_shape, ShapeUtil::MakeShape(F32, {32, 48})))
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4},
/*elided_window_dims=*/{0},
- /*gather_dims_to_operand_dims=*/{0}),
+ /*gather_dims_to_operand_dims=*/{0},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{1, 48}));
EXPECT_TRUE(ShapeUtil::Equal(gather_shape,
ShapeUtil::MakeShape(F32, {10, 9, 8, 7, 48})))
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7, 8},
/*elided_window_dims=*/{},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 27, 26}));
EXPECT_TRUE(ShapeUtil::Equal(
gather_shape,
<< ShapeUtil::HumanString(gather_shape);
}
+TEST_F(GatherShapeInferenceTest, NonDefaultGatherIndicesLeafDim_A) {
+ TF_ASSERT_OK_AND_ASSIGN(
+ Shape gather_shape,
+ ShapeInference::InferGatherShape(
+ f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_5_7_6_,
+ HloInstruction::MakeGatherDimNumbers(
+ /*output_window_dims=*/{4, 5, 6, 7, 8},
+ /*elided_window_dims=*/{},
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/2),
+ /*window_bounds=*/{30, 29, 28, 27, 26}));
+
+ EXPECT_TRUE(ShapeUtil::Equal(
+ gather_shape,
+ ShapeUtil::MakeShape(F32, {10, 9, 7, 6, 30, 29, 28, 27, 26})))
+ << ShapeUtil::HumanString(gather_shape);
+}
+
+TEST_F(GatherShapeInferenceTest, NonDefaultGatherIndicesLeafDim_B) {
+ TF_ASSERT_OK_AND_ASSIGN(
+ Shape gather_shape,
+ ShapeInference::InferGatherShape(
+ f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_5_10_9_7_6_,
+ HloInstruction::MakeGatherDimNumbers(
+ /*output_window_dims=*/{4, 5, 6, 7, 8},
+ /*elided_window_dims=*/{},
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/0),
+ /*window_bounds=*/{30, 29, 28, 27, 26}));
+
+ EXPECT_TRUE(ShapeUtil::Equal(
+ gather_shape,
+ ShapeUtil::MakeShape(F32, {10, 9, 7, 6, 30, 29, 28, 27, 26})))
+ << ShapeUtil::HumanString(gather_shape);
+}
+
+TEST_F(GatherShapeInferenceTest, NoOutputGatherDims) {
+ // This is equivalent to a dynamic slice.
+ TF_ASSERT_OK_AND_ASSIGN(
+ Shape gather_shape,
+ ShapeInference::InferGatherShape(
+ f32_5d_tensor_50_49_48_47_46_, s64_vector_5_,
+ HloInstruction::MakeGatherDimNumbers(
+ /*output_window_dims=*/{0, 1, 2, 3, 4},
+ /*elided_window_dims=*/{},
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/0),
+ /*window_bounds=*/{30, 29, 28, 27, 26}));
+
+ EXPECT_TRUE(ShapeUtil::Equal(gather_shape,
+ ShapeUtil::MakeShape(F32, {30, 29, 28, 27, 26})))
+ << ShapeUtil::HumanString(gather_shape);
+}
+
+TEST_F(GatherShapeInferenceTest, ScalarGatherIndices) {
+ // The gather indices "tensor" is a scalar S here that's used to slice out
+ // [S,0,0,0,0]..[S,30,29,28,27] into a [30,29,28,27] shaped result.
+ TF_ASSERT_OK_AND_ASSIGN(Shape gather_shape,
+ ShapeInference::InferGatherShape(
+ f32_5d_tensor_50_49_48_47_46_, s64_scalar_,
+ HloInstruction::MakeGatherDimNumbers(
+ /*output_window_dims=*/{0, 1, 2, 3},
+ /*elided_window_dims=*/{0},
+ /*gather_dims_to_operand_dims=*/{0},
+ /*index_vector_dim=*/0),
+ /*window_bounds=*/{1, 30, 29, 28, 27}));
+
+ EXPECT_TRUE(ShapeUtil::Equal(gather_shape,
+ ShapeUtil::MakeShape(F32, {30, 29, 28, 27})))
+ << ShapeUtil::HumanString(gather_shape);
+}
+
TEST_F(GatherShapeInferenceTest, TupleShapedTensorInput) {
StatusOr<Shape> statusor = ShapeInference::InferGatherShape(
tuple_shape_, s64_vector_32_,
HloInstruction::MakeGatherDimNumbers(/*output_window_dims=*/{0},
/*elided_window_dims=*/{1},
- /*gather_dims_to_operand_dims=*/{1}),
+ /*gather_dims_to_operand_dims=*/{1},
+ /*index_vector_dim=*/1),
/*window_bounds=*/{64, 1});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(statusor.status().error_message(),
s64_vector_32_, tuple_shape_,
HloInstruction::MakeGatherDimNumbers(/*output_window_dims=*/{0},
/*elided_window_dims=*/{1},
- /*gather_dims_to_operand_dims=*/{1}),
+ /*gather_dims_to_operand_dims=*/{1},
+ /*index_vector_dim=*/0),
/*window_bounds=*/{64, 1});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(statusor.status().error_message(),
<< statusor.status();
}
-TEST_F(GatherShapeInferenceTest, ScalarGatherIndicesInput) {
- StatusOr<Shape> statusor = ShapeInference::InferGatherShape(
- s64_vector_32_, s32_,
- HloInstruction::MakeGatherDimNumbers(/*output_window_dims=*/{0},
- /*elided_window_dims=*/{1},
- /*gather_dims_to_operand_dims=*/{1}),
- /*window_bounds=*/{64, 1});
- ASSERT_FALSE(statusor.ok());
- EXPECT_THAT(statusor.status().error_message(),
- HasSubstr("Gather indices parameter must at least of rank 1"))
- << statusor.status();
-}
-
TEST_F(GatherShapeInferenceTest, FloatingPointGatherIndicesInput) {
StatusOr<Shape> statusor = ShapeInference::InferGatherShape(
s64_vector_32_, vector_32_,
HloInstruction::MakeGatherDimNumbers(/*output_window_dims=*/{0},
/*elided_window_dims=*/{1},
- /*gather_dims_to_operand_dims=*/{1}),
+ /*gather_dims_to_operand_dims=*/{1},
+ /*index_vector_dim=*/0),
/*window_bounds=*/{64, 1});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(statusor.status().error_message(),
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 8, 7},
/*elided_window_dims=*/{},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 27, 26});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7, 7},
/*elided_window_dims=*/{},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 27, 26});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 99, 100, 101},
/*elided_window_dims=*/{},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 27, 26});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(statusor.status().error_message(),
}
TEST_F(GatherShapeInferenceTest,
+ InvalidGatherDimNumbers_WindowIndexBarelyOutOfBounds) {
+ StatusOr<Shape> statusor = ShapeInference::InferGatherShape(
+ f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_,
+ HloInstruction::MakeGatherDimNumbers(
+ /*output_window_dims=*/{4, 5, 6, 7, 9},
+ /*elided_window_dims=*/{},
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
+ /*window_bounds=*/{30, 29, 28, 27, 26});
+ ASSERT_FALSE(statusor.ok());
+ EXPECT_THAT(statusor.status().error_message(),
+ HasSubstr("Window index 4 in gather op is out of bounds"))
+ << statusor.status();
+}
+
+TEST_F(GatherShapeInferenceTest,
InvalidGatherDimNumbers_MismatchingElidedWindowDims) {
StatusOr<Shape> statusor = ShapeInference::InferGatherShape(
f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_,
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7, 8},
/*elided_window_dims=*/{4},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 27, 26});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7, 8},
/*elided_window_dims=*/{0, 1, 2, 3, 19},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 27, 26});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(statusor.status().error_message(),
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7, 8},
/*elided_window_dims=*/{0, 1, 2, 3, 3},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 27, 26});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7, 8},
/*elided_window_dims=*/{},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 27, 26});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(
statusor.status().error_message(),
- HasSubstr(
- "There must be exactly as many elements in "
- "gather_dims_to_operand_dims "
- "as there are elements in the last dimension of %gather_indices"))
+ HasSubstr("Gather op has 4 elements in gather_dims_to_operand_dims and "
+ "the bound of dimension index_vector_dim=4 of "
+ "gather_indices is 5. These two numbers must be equal."))
<< statusor.status();
}
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7, 8},
/*elided_window_dims=*/{},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 7}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 7},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 27, 26});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7, 8},
/*elided_window_dims=*/{},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 3}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 3},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 27, 26});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7, 8},
/*elided_window_dims=*/{2, 1},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{1, 1, 28, 27, 26});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(statusor.status().error_message(),
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7},
/*elided_window_dims=*/{2},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 1, 300, 26});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(statusor.status().error_message(),
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7, 8},
/*elided_window_dims=*/{},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 26});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(
HloInstruction::MakeGatherDimNumbers(
/*output_window_dims=*/{4, 5, 6, 7},
/*elided_window_dims=*/{1},
- /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}),
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/4),
/*window_bounds=*/{30, 29, 28, 26, 20});
ASSERT_FALSE(statusor.ok());
EXPECT_THAT(statusor.status().error_message(),
<< statusor.status();
}
+TEST_F(GatherShapeInferenceTest, OutOfBoundsGatherIndicesLeafDim) {
+ StatusOr<Shape> statusor = ShapeInference::InferGatherShape(
+ f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_5_7_6_,
+ HloInstruction::MakeGatherDimNumbers(
+ /*output_window_dims=*/{4, 5, 6, 7, 8},
+ /*elided_window_dims=*/{},
+ /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
+ /*index_vector_dim=*/32),
+ /*window_bounds=*/{30, 29, 28, 27, 26});
+
+ ASSERT_FALSE(statusor.ok());
+ EXPECT_THAT(statusor.status().error_message(),
+ HasSubstr("Gather index leaf dimension must be within [0, "
+ "rank(gather_indices) + 1)"))
+ << statusor.status();
+}
+
} // namespace
} // namespace xla
: : : indices of the slices we're :
: : : we're stitching together into :
: : : the output tensor. :
+|`index_vector_dim` | `int64` | The dimension in |
+: : : `gather_indices` that contains :
+: : : the starting indices. :
|`output_window_dims` | `ArraySlice<int64>` | The set of dimensions in the |
: : : output shape that are _window :
: : : dimensions_ (defined below). :
: : : `output_window_dims`) and the window :
: : : dimensions that are elided (via :
: : : `elided_window_dims`). :
-|`gather_dims_to_operand_dims` | `ArraySlice<int64>` | A dimension map (the |
+|`gather_dims_to_operand_dims` | `ArraySlice<int64>` | A dimension map (the |
: : : array is interpreted as mapping `i` to :
: : : `gather_dims_to_operand_dims[i]`) from :
: : : the gather indices in `gather_indices` to :
: : : the operand index space. It has to be :
: : : one-to-one and total. :
-If `gather_indices` is a vector with `N` elements then we implicitly reshape it
-to a tensor of shape `[N,1]` before proceeding.
-
For every index `Out` in the output tensor, we compute two things (more
precisely described later):
- - An index into the first `gather_indices.rank` - `1` dimensions of
- `gather_indices`, which gives us a starting index of a slice, _operand
- slice_, in the operand tensor.
+ - An index into `gather_indices.rank` - `1` dimensions of `gather_indices`,
+ which gives us a starting index of a slice, _operand slice_, in the operand
+ tensor. These `gather_indices.rank` - `1` dimensions are all the dimensions
+ in `gather_indices` except `index_vector_dim`.
- A _window index_ that has the same rank as the operand. This index is
composed of the values in `Out` at dimensions `output_window_dims`, embedded
The output is a tensor of rank `output_window_dims.size` + `gather_indices.rank`
- `1`. Additionally, as a shorthand, we define `output_gather_dims` of type
`ArraySlice<int64>` as the set of dimensions in the output shape but not in
-`output_window_dims`, in ascending order. E.g. if the output tensor has rank 5,
-`output_window_dims` is {`2`, `4`} then `output_gather_dims` is {`0`, `1`, `3`}
+`output_window_dims`, in ascending order. E.g. if the output tensor has rank
+`5`, `output_window_dims` is {`2`, `4`} then `output_gather_dims` is {`0`, `1`,
+`3`}
+
+If `index_vector_dim` is equal to `gather_indices.rank` we implicitly
+consider `gather_indices` to have a trailing `1` dimension (i.e. if
+`gather_indices` was of shape `[6,7]` and `index_vector_dim` is `2` then
+we implicitly consider the shape of `gather_indices` to be `[6,7,1]`).
The bounds for the output tensor along dimension `i` is computed as follows:
1. If `i` is present in `output_gather_dims` (i.e. is equal to
- `output_gather_dims[k]` for some `k`) then we pick the corresponding
- dimension bounds out of `gather_indices.shape` (i.e. pick
- `gather_indices.shape.dims[k]`).
+ `output_gather_dims[k]` for some `k`) then we pick the corresponding
+ dimension bounds out of `gather_indices.shape`, skipping
+ `index_vector_dim` (i.e. pick `gather_indices.shape.dims`[`k`] if `k`
+ < `index_vector_dim` and `gather_indices.shape.dims`[`k`+`1`]
+ otherwise).
2. If `i` is present in `output_window_dims` (i.e. equal to
- `output_window_dims[k]` for some `k`) then we pick the corresponding bound
- out of `window_bounds` after accounting for `elided_window_dims` (i.e. we
- pick `adjusted_window_bounds[k]` where `adjusted_window_bounds` is
- `window_bounds` with the bounds at indices `elided_window_dims` removed).
+ `output_window_dims`[`k`] for some `k`) then we pick the corresponding
+ bound out of `window_bounds` after accounting for `elided_window_dims`
+ (i.e. we pick `adjusted_window_bounds`[`k`] where `adjusted_window_bounds`
+ is `window_bounds` with the bounds at indices `elided_window_dims`
+ removed).
The operand index `In` corresponding to an output index `Out` is computed as
follows:
1. Let `G` = { `Out`[`k`] for `k` in `output_gather_dims` }. Use `G` to slice
- out vector `S` such that `S`[`i`] = `gather_indices`[`G`, `i`].
- 2. Create an index, `S`<sub>`in`</sub>, into `operand` using `S` by scattering
- `S` using the `gather_dims_to_operand_dims` map (`S`<sub>`in`</sub> is the
- starting indices for _operand slice_ mentioned above.). More precisely:
+ out vector `S` such that `S`[`i`] = `gather_indices`[Combine(`G`, `i`)]
+ where Combine(A, b) inserts b at position `index_vector_dim` into A.
+ Note that this is well defined even if `G` is empty -- if `G` is empty then
+ `S` = `gather_indices`.
+ 2. Create an index, `S`<sub>`in`</sub>, into `operand` using `S` by
+ scattering `S` using the `gather_dims_to_operand_dims` map
+ (`S`<sub>`in`</sub> is the starting indices for _operand slice_ mentioned
+ above). More precisely:
1. `S`<sub>`in`</sub>[`gather_dims_to_operand_dims`[`k`]] = `S`[`k`] if `k` <
`gather_dims_to_operand_dims.size`.
2. `S`<sub>`in`</sub>[`_`] = `0` otherwise.
`operand.rank` is `6` and `elided_window_dims` is {`0`, `2`} then
`window_dims_to_operand_dims` is {`0`→`1`, `1`→`3`, `2`→`4`, `3`→`5`}.
-### Informal Description
+### Informal Description and Examples
+
+`index_vector_dim` is set to `gather_indices.rank` - `1` in all of the
+examples that follow. More interesting values for `index_vector_dim`
+does not change the operation fundamentally, but makes the visual representation
+more cumbersome.
To get an intuition on how all of the above fits together, let's look at an
example that gathers 5 slices of shape `[8,6]` from a `[16,11]` tensor. The
projects / platform / upstream / tensorflow.git / commitdiff