Home
last modified time | relevance | path

Searched refs:gather_dim_numbers (Results 1 – 7 of 7) 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/
Dshape_inference.cc2959 const GatherDimensionNumbers& gather_dim_numbers, in InferGatherShape() argument
2977 gather_dim_numbers.index_vector_dim() || in InferGatherShape()
2978 gather_dim_numbers.index_vector_dim() < 0) { in InferGatherShape()
2984 gather_dim_numbers.index_vector_dim()); in InferGatherShape()
2992 gather_dim_numbers.index_vector_dim()) { in InferGatherShape()
2997 input_shape, expanded_start_indices_shape, gather_dim_numbers)); in InferGatherShape()
3007 gather_dim_numbers.offset_dims_size() + in InferGatherShape()
3008 gather_dim_numbers.collapsed_slice_dims_size()) { in InferGatherShape()
3013 slice_sizes.size(), StrJoin(gather_dim_numbers.offset_dims(), ","), in InferGatherShape()
3014 StrJoin(gather_dim_numbers.collapsed_slice_dims(), ",")); in InferGatherShape()
[all …]
Dhlo_instructions.cc2311 const GatherDimensionNumbers& gather_dim_numbers, in HloGatherInstruction() argument
2317 absl::make_unique<GatherDimensionNumbers>(gather_dim_numbers); in HloGatherInstruction()
2344 GatherDimensionNumbers gather_dim_numbers; in MakeGatherDimNumbers() local
2346 gather_dim_numbers.add_offset_dims(output_window_dim); in MakeGatherDimNumbers()
2349 gather_dim_numbers.add_collapsed_slice_dims(elided_window_dim); in MakeGatherDimNumbers()
2352 gather_dim_numbers.add_start_index_map(gather_dim_to_input_dim); in MakeGatherDimNumbers()
2355 gather_dim_numbers.set_index_vector_dim(index_vector_dim); in MakeGatherDimNumbers()
2356 return gather_dim_numbers; in MakeGatherDimNumbers()
Dshape_inference.h284 const GatherDimensionNumbers& gather_dim_numbers,
Dhlo_instruction.h728 const GatherDimensionNumbers& gather_dim_numbers,
Dhlo_instructions.h1364 const GatherDimensionNumbers& gather_dim_numbers,
Dhlo_instruction.cc1308 const GatherDimensionNumbers& gather_dim_numbers, in CreateGather() argument
1311 shape, operand, start_indices, gather_dim_numbers, slice_sizes); in CreateGather()