Searched refs:dim_numbers_ (Results 1 – 2 of 2) sorted by relevance
833 : dim_numbers_(*dim_numbers), start_indices_(*start_indices) { in OutputBatchIndexToInputIndex()836 !absl::c_binary_search(dim_numbers_.offset_dims(), i)); in OutputBatchIndexToInputIndex()841 std::distance(dim_numbers_.start_index_map().begin(), in OutputBatchIndexToInputIndex()842 absl::c_find(dim_numbers_.start_index_map(), i)); in OutputBatchIndexToInputIndex()844 dim_numbers_.start_index_map_size()) { in OutputBatchIndexToInputIndex()855 start_indices_.shape().dimensions(dim_numbers_.index_vector_dim()); in OutputBatchIndexToInputIndex()894 if (index_vector_index_i == dim_numbers_.index_vector_dim()) { in PropagateOutputIndexGatherDimsToIndexVectorIndex()905 int64 index_vector_dim = dim_numbers_.index_vector_dim(); in FetchIndexVector()946 const GatherDimensionNumbers& dim_numbers_; member in xla::OutputBatchIndexToInputIndex
2108 : dim_numbers_(*dim_numbers), scatter_indices_(*scatter_indices) { in UpdateScatterIndexToInputIndex()2111 !absl::c_binary_search(dim_numbers_.update_window_dims(), i)); in UpdateScatterIndexToInputIndex()2116 FindIndex(dim_numbers_.scatter_dims_to_operand_dims(), i); in UpdateScatterIndexToInputIndex()2118 dim_numbers_.scatter_dims_to_operand_dims_size()) { in UpdateScatterIndexToInputIndex()2129 scatter_indices_.shape().dimensions(dim_numbers_.index_vector_dim()); in UpdateScatterIndexToInputIndex()2168 if (index_vector_index_i == dim_numbers_.index_vector_dim()) { in PropagateUpdateIndexScatterDimsToIndexVectorIndex()2179 int64 index_vector_dim = dim_numbers_.index_vector_dim(); in FetchIndexVector()2220 const ScatterDimensionNumbers& dim_numbers_; variable