Home
last modified time | relevance | path

Searched refs:gather_dim_numbers (Results 1 – 8 of 8) sorted by relevance

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dresampler_ops.cc140 xla::GatherDimensionNumbers gather_dim_numbers; in Gather2by2Neighbors() local
144 gather_dim_numbers.add_offset_dims(neighbor_data_dimensions - 3); in Gather2by2Neighbors()
145 gather_dim_numbers.add_offset_dims(neighbor_data_dimensions - 2); in Gather2by2Neighbors()
146 gather_dim_numbers.add_offset_dims(neighbor_data_dimensions - 1); in Gather2by2Neighbors()
148 gather_dim_numbers.set_index_vector_dim(warp_dims - 1); in Gather2by2Neighbors()
149 gather_dim_numbers.add_collapsed_slice_dims(0); in Gather2by2Neighbors()
150 gather_dim_numbers.add_start_index_map(0); in Gather2by2Neighbors()
154 gather_dim_numbers.add_start_index_map(2); in Gather2by2Neighbors()
155 gather_dim_numbers.add_start_index_map(1); in Gather2by2Neighbors()
158 auto neighbors_data = xla::Gather(data, gather_indices, gather_dim_numbers, in Gather2by2Neighbors()
/external/tensorflow/tensorflow/compiler/xla/service/
Ddynamic_padder.cc429 GatherDimensionNumbers gather_dim_numbers; in RewriteDynamicReshapeSplitInput() local
435 gather_dim_numbers.add_offset_dims(i); in RewriteDynamicReshapeSplitInput()
439 gather_dim_numbers.add_start_index_map(input_dim); in RewriteDynamicReshapeSplitInput()
440 gather_dim_numbers.set_index_vector_dim(1); in RewriteDynamicReshapeSplitInput()
441 gather_dim_numbers.add_collapsed_slice_dims(input_dim); in RewriteDynamicReshapeSplitInput()
461 operand_static, cumsum, gather_dim_numbers, slice_sizes, true)); in RewriteDynamicReshapeSplitInput()
650 GatherDimensionNumbers gather_dim_numbers; in RewriteDynamicReshapeCombineInput() local
656 gather_dim_numbers.add_offset_dims(i); in RewriteDynamicReshapeCombineInput()
660 gather_dim_numbers.add_start_index_map(output_dim); in RewriteDynamicReshapeCombineInput()
661 gather_dim_numbers.set_index_vector_dim(1); in RewriteDynamicReshapeCombineInput()
[all …]
Dshape_inference.cc3304 const GatherDimensionNumbers& gather_dim_numbers, in InferGatherShape() argument
3322 gather_dim_numbers.index_vector_dim() || in InferGatherShape()
3323 gather_dim_numbers.index_vector_dim() < 0) { in InferGatherShape()
3329 gather_dim_numbers.index_vector_dim()); in InferGatherShape()
3344 gather_dim_numbers.index_vector_dim()) { in InferGatherShape()
3350 input_shape, expanded_start_indices_shape, gather_dim_numbers)); in InferGatherShape()
3360 gather_dim_numbers.offset_dims_size() + in InferGatherShape()
3361 gather_dim_numbers.collapsed_slice_dims_size()) { in InferGatherShape()
3366 slice_sizes.size(), StrJoin(gather_dim_numbers.offset_dims(), ","), in InferGatherShape()
3367 StrJoin(gather_dim_numbers.collapsed_slice_dims(), ",")); in InferGatherShape()
[all …]
Dhlo_instructions.cc2725 const GatherDimensionNumbers& gather_dim_numbers, in HloGatherInstruction() argument
2732 absl::make_unique<GatherDimensionNumbers>(gather_dim_numbers); in HloGatherInstruction()
2759 GatherDimensionNumbers gather_dim_numbers; in MakeGatherDimNumbers() local
2761 gather_dim_numbers.add_offset_dims(output_window_dim); in MakeGatherDimNumbers()
2764 gather_dim_numbers.add_collapsed_slice_dims(elided_window_dim); in MakeGatherDimNumbers()
2767 gather_dim_numbers.add_start_index_map(gather_dim_to_input_dim); in MakeGatherDimNumbers()
2770 gather_dim_numbers.set_index_vector_dim(index_vector_dim); in MakeGatherDimNumbers()
2771 return gather_dim_numbers; in MakeGatherDimNumbers()
Dshape_inference.h315 const GatherDimensionNumbers& gather_dim_numbers,
Dhlo_instruction.h946 const GatherDimensionNumbers& gather_dim_numbers,
Dhlo_instructions.h1657 const GatherDimensionNumbers& gather_dim_numbers,
Dhlo_instruction.cc1617 const GatherDimensionNumbers& gather_dim_numbers, in CreateGather() argument
1620 shape, operand, start_indices, gather_dim_numbers, slice_sizes, in CreateGather()