Searched refs:gather_dim_numbers (Results 1 – 7 of 7) sorted by relevance
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
D | resampler_ops.cc | 140 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/ |
D | shape_inference.cc | 2959 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 …]
|
D | hlo_instructions.cc | 2311 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()
|
D | shape_inference.h | 284 const GatherDimensionNumbers& gather_dim_numbers,
|
D | hlo_instruction.h | 728 const GatherDimensionNumbers& gather_dim_numbers,
|
D | hlo_instructions.h | 1364 const GatherDimensionNumbers& gather_dim_numbers,
|
D | hlo_instruction.cc | 1308 const GatherDimensionNumbers& gather_dim_numbers, in CreateGather() argument 1311 shape, operand, start_indices, gather_dim_numbers, slice_sizes); in CreateGather()
|