| /external/tensorflow/tensorflow/compiler/tf2xla/lib/ |
| scatter.cc | 138 xla::ScatterDimensionNumbers dim_numbers; local
139 dim_numbers.set_index_vector_dim(indices_are_vectors
165 dim_numbers.add_update_window_dims(i);
170 dim_numbers.add_inserted_window_dims(i);
171 dim_numbers.add_scatter_dims_to_operand_dims(i);
193 VLOG(3) << " index_vector_dim: " << dim_numbers.index_vector_dim();
195 << absl::StrJoin(dim_numbers.update_window_dims(), ",") << "]";
197 << absl::StrJoin(dim_numbers.inserted_window_dims(), ",") << "]";
199 << absl::StrJoin(dim_numbers.scatter_dims_to_operand_dims(), ",")
203 dim_numbers);
[all...] |
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| batch_dot_simplification.cc | 26 const DotDimensionNumbers& dim_numbers = batch_dot->dot_dimension_numbers(); local
34 if (dim_numbers.lhs_contracting_dimensions_size() != 1) {
39 for (int64 batch_dim : dim_numbers.lhs_batch_dimensions()) {
54 DotDimensionNumbers new_dim_numbers = dim_numbers;
58 for (int64 i = 0, e = dim_numbers.lhs_batch_dimensions_size() -
|
| gather_expander.cc | 111 HloInstruction* index_vector, const GatherDimensionNumbers& dim_numbers,
132 int64 index_vector_dim_index = FindIndex(dim_numbers.start_index_map(), i);
133 if (index_vector_dim_index != dim_numbers.start_index_map_size()) {
153 const GatherDimensionNumbers& dim_numbers = gather.gather_dimension_numbers(); local
161 dim_numbers.index_vector_dim() ==
197 ExpandIndexVectorIntoOperandSpace(index_vector, dim_numbers,
207 AsInt64Slice(dim_numbers.collapsed_slice_dims())));
235 const GatherDimensionNumbers& dim_numbers) {
240 if (!absl::c_binary_search(dim_numbers.collapsed_slice_dims(), i)) {
318 const GatherDimensionNumbers& dim_numbers local
[all...] |
| scatter_expander.cc | 133 HloInstruction* index_vector, const ScatterDimensionNumbers& dim_numbers,
154 FindIndex(dim_numbers.scatter_dims_to_operand_dims(), i);
156 dim_numbers.scatter_dims_to_operand_dims_size()) {
222 const ScatterDimensionNumbers& dim_numbers = local
258 ExpandIndexVectorIntoOperandSpace(index_vector, dim_numbers,
279 AsInt64Slice(dim_numbers.inserted_window_dims())));
350 const ScatterDimensionNumbers& dim_numbers = local
364 if (i != dim_numbers.index_vector_dim()) {
379 scatter_indices, dim_numbers.index_vector_dim()));
388 updates, AsInt64Slice(dim_numbers.update_window_dims())))
[all...] |
| hlo_cost_analysis_test.cc | 650 GatherDimensionNumbers dim_numbers; local
651 dim_numbers.add_offset_dims(1);
652 dim_numbers.add_collapsed_slice_dims(0);
653 dim_numbers.add_start_index_map(0);
654 dim_numbers.set_index_vector_dim(1);
655 Gather(operand, indices, dim_numbers, {1, 3});
677 ScatterDimensionNumbers dim_numbers; local
678 dim_numbers.set_index_vector_dim(1);
679 dim_numbers.add_update_window_dims(1);
680 dim_numbers.add_inserted_window_dims(0)
[all...] |
| triangular_solve_expander.cc | 74 GatherDimensionNumbers dim_numbers; local
76 dim_numbers.add_offset_dims(i);
77 dim_numbers.add_start_index_map(i);
81 dim_numbers.add_offset_dims(ndims - 1);
82 dim_numbers.add_offset_dims(ndims);
83 dim_numbers.add_start_index_map(ndims - 2);
84 dim_numbers.add_start_index_map(ndims - 1);
85 dim_numbers.set_index_vector_dim(1);
86 diag_blocks = Gather(a, start_indices, dim_numbers, slice_sizes);
|
| convolution_group_converter.cc | 202 auto dim_numbers = convolution->convolution_dimension_numbers(); local
218 int64 input_batch_dimension = dim_numbers.input_batch_dimension();
219 int64 output_batch_dimension = dim_numbers.output_batch_dimension();
220 int64 output_feature_dimension = dim_numbers.output_feature_dimension();
243 convolution->window(), dim_numbers, convolution->precision_config()));
334 auto dim_numbers = convolution->convolution_dimension_numbers(); local
336 int64 kernel_input_feature_dim = dim_numbers.kernel_input_feature_dimension();
339 dim_numbers.kernel_output_feature_dimension();
383 convolution->window(), dim_numbers, convolution->precision_config());
387 int64 activation_input_feature_dim = dim_numbers.input_feature_dimension()
[all...] |
| hlo_evaluator.cc | 364 const DotDimensionNumbers& dim_numbers,
374 ShapeInference::InferDotOpShape(lhs.shape(), rhs.shape(), dim_numbers));
378 dim_numbers, precision_config);
782 const Shape& output_shape, const GatherDimensionNumbers& dim_numbers) {
789 !absl::c_binary_search(dim_numbers.offset_dims(), i);
801 const GatherDimensionNumbers& dim_numbers) {
807 absl::c_binary_search(dim_numbers.offset_dims(), i);
809 while (absl::c_binary_search(dim_numbers.collapsed_slice_dims(),
831 const GatherDimensionNumbers* dim_numbers, const Shape& input_shape,
833 : dim_numbers_(*dim_numbers), start_indices_(*start_indices)
1050 const GatherDimensionNumbers& dim_numbers = local
[all...] |
| elemental_ir_emitter.cc | 1867 const GatherDimensionNumbers& dim_numbers = hlo->gather_dimension_numbers(); local
2118 const DotDimensionNumbers& dim_numbers = hlo->dot_dimension_numbers(); local
[all...] |
| hlo_parser.cc | 1636 GatherDimensionNumbers dim_numbers = local
1672 ScatterDimensionNumbers dim_numbers = local
[all...] |
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| gather_operation_test.cc | 646 GatherDimensionNumbers dim_numbers; local
647 dim_numbers.add_offset_dims(1);
648 dim_numbers.add_collapsed_slice_dims(0);
649 dim_numbers.add_start_index_map(0);
650 dim_numbers.set_index_vector_dim(1);
651 Gather(operand, indices, dim_numbers, {1, 3});
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| gather_op.cc | 116 xla::GatherDimensionNumbers dim_numbers; local
122 dim_numbers.add_collapsed_slice_dims(i);
131 dim_numbers.add_offset_dims(i);
135 dim_numbers.add_offset_dims(i + indices_rank - num_index_dims);
139 dim_numbers.set_index_vector_dim(indices_are_nd ? (indices_shape.dims() - 1)
142 dim_numbers.add_start_index_map(i);
145 *gather_output = xla::Gather(input, indices, dim_numbers, slice_sizes);
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| ir_emission_utils.cc | 69 const DotDimensionNumbers& dim_numbers = dot.dot_dimension_numbers(); local
74 dim_numbers.lhs_batch_dimensions_size())) {
78 CHECK_EQ(lhs_shape.dimensions(dim_numbers.lhs_contracting_dimensions(0)),
79 rhs_shape.dimensions(dim_numbers.rhs_contracting_dimensions(0)));
|
| ir_emitter_unnested.cc | 1081 const ScatterDimensionNumbers& dim_numbers = local
[all...] |
