Searched refs:source_dim (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | indexed_array_analysis.cc | 73 indexed_array->source_dim(), "->[", in ToString()
188 ScalarIndexedArray* source, Array* indices, int64 source_dim, in FoldGatherOfGather() argument
207 EraseAt(&simulated_index, source->source_dim()); in FoldGatherOfGather()
214 EraseAt(&simulated_index, source_dim); in FoldGatherOfGather()
221 FindIndex(source->output_dims(), source_dim); in FoldGatherOfGather()
248 return ConstructScalarIndexedArray(a, inner_indices, source->source_dim(), in FoldGatherOfGather()
290 int64 source_dim = dim_numbers.start_index_map(0); in ComputeArrayForGather() local
299 if (absl::c_linear_search(indexed->output_dims(), source_dim)) { in ComputeArrayForGather()
300 return FoldGatherOfGather(indexed, indices, source_dim, output_dims, in ComputeArrayForGather()
304 return Construct<ScalarIndexedConstantArray>(constant, indices, source_dim, in ComputeArrayForGather()
[all …]
|
| D | indexed_array_analysis.h | 186 int64 source_dim() const { return source_dim_; } in source_dim() function
194 explicit ScalarIndexedArray(Array* source, Array* indices, int64 source_dim, in ScalarIndexedArray() argument
198 source_dim_(source_dim), in ScalarIndexedArray()
225 int64 source_dim, in ScalarIndexedConstantArray() argument
228 : ScalarIndexedArray(source, indices, source_dim, in ScalarIndexedConstantArray()
305 ScalarIndexedArray* source, Array* indices, int64 source_dim,
337 Array* source, Array* indices, int64 source_dim, in ConstructScalarIndexedArray() argument
340 return Construct<ScalarIndexedConstantArray>(source, indices, source_dim, in ConstructScalarIndexedArray()
344 return Construct<ScalarIndexedArray>(source, indices, source_dim, in ConstructScalarIndexedArray()
|