Home
last modified time | relevance | path

Searched refs:scatter_dim_numbers (Results 1 – 7 of 7) sorted by relevance

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dresampler_ops.cc173 xla::ScatterDimensionNumbers scatter_dim_numbers; in ScatterToGradData() local
177 scatter_dim_numbers.add_update_window_dims(neighbor_data_dimensions - 3); in ScatterToGradData()
178 scatter_dim_numbers.add_update_window_dims(neighbor_data_dimensions - 2); in ScatterToGradData()
179 scatter_dim_numbers.add_update_window_dims(neighbor_data_dimensions - 1); in ScatterToGradData()
180 scatter_dim_numbers.set_index_vector_dim(warp_dims - 1); in ScatterToGradData()
182 scatter_dim_numbers.add_inserted_window_dims(0); in ScatterToGradData()
183 scatter_dim_numbers.add_scatter_dims_to_operand_dims(0); in ScatterToGradData()
187 scatter_dim_numbers.add_scatter_dims_to_operand_dims(2); in ScatterToGradData()
188 scatter_dim_numbers.add_scatter_dims_to_operand_dims(1); in ScatterToGradData()
192 scatter_dim_numbers); in ScatterToGradData()
/external/tensorflow/tensorflow/compiler/xla/service/
Dshape_inference.cc3660 const ScatterDimensionNumbers& scatter_dim_numbers) { in InferScatterShape() argument
3677 scatter_dim_numbers.index_vector_dim() || in InferScatterShape()
3678 scatter_dim_numbers.index_vector_dim() < 0) { in InferScatterShape()
3684 scatter_dim_numbers.index_vector_dim()); in InferScatterShape()
3690 scatter_dim_numbers.index_vector_dim()) { in InferScatterShape()
3710 scatter_dim_numbers.inserted_window_dims_size() && in InferScatterShape()
3711 scatter_dim_numbers.inserted_window_dims(inserted_dims_seen) == i) { in InferScatterShape()
3719 scatter_dim_numbers.update_window_dims_size(); in InferScatterShape()
3727 scatter_dim_numbers)); in InferScatterShape()
3729 for (int i = 0; i < scatter_dim_numbers.update_window_dims_size(); ++i) { in InferScatterShape()
[all …]
Dshape_inference.h355 const ScatterDimensionNumbers& scatter_dim_numbers);
Dhlo_instructions.cc3143 const ScatterDimensionNumbers& scatter_dim_numbers, bool indices_are_sorted, in HloScatterInstruction() argument
3154 std::make_unique<ScatterDimensionNumbers>(scatter_dim_numbers); in HloScatterInstruction()
3184 ScatterDimensionNumbers scatter_dim_numbers; in MakeScatterDimNumbers() local
3186 scatter_dim_numbers.add_update_window_dims(update_window_dim); in MakeScatterDimNumbers()
3189 scatter_dim_numbers.add_inserted_window_dims(inserted_window_dim); in MakeScatterDimNumbers()
3192 scatter_dim_numbers.add_scatter_dims_to_operand_dims( in MakeScatterDimNumbers()
3195 scatter_dim_numbers.set_index_vector_dim(index_vector_dim); in MakeScatterDimNumbers()
3196 return scatter_dim_numbers; in MakeScatterDimNumbers()
Dhlo_instruction.h1069 const ScatterDimensionNumbers& scatter_dim_numbers,
1076 const ScatterDimensionNumbers& scatter_dim_numbers,
Dhlo_instruction.cc1939 const ScatterDimensionNumbers& scatter_dim_numbers, bool indices_are_sorted, in CreateScatter() argument
1943 scatter_dim_numbers, indices_are_sorted, unique_indices)); in CreateScatter()
1950 const ScatterDimensionNumbers& scatter_dim_numbers, bool indices_are_sorted, in CreateScatter() argument
1958 shape, args, update_computation, scatter_dim_numbers, indices_are_sorted, in CreateScatter()
Dhlo_instructions.h2154 const ScatterDimensionNumbers& scatter_dim_numbers,