| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | shape_partition.cc | 25 std::vector<int64> outer_dims; in Run() local
31 outer_dims.push_back(dimension); in Run()
48 static_cast<double>(target_partition_count), 1.0 / outer_dims.size()); in Run()
52 std::vector<int64> dimension_partition_counts(outer_dims.size()); in Run()
53 for (int64 i = 0; i < outer_dims.size(); ++i) { in Run()
55 std::min(static_cast<int64>(shape_.dimensions(outer_dims[i])), in Run()
79 shape_.dimensions(outer_dims[i]) - dimension_partition_counts[i], in Run()
|
| D | ir_emitter.cc | 1907 std::vector<int64> outer_dims; in HandleSlice() local
1909 outer_dims.push_back(LayoutUtil::Major(layout, i)); in HandleSlice()
1915 outer_dims.push_back(memcpy_dim); in HandleSlice()
1920 const int64 num_outer_loops = outer_dims.size(); in HandleSlice()
1923 loops.AddLoopsForShapeOnDimensions(slice->shape(), outer_dims, "slice"); in HandleSlice()
2472 std::vector<int64> outer_dims(std::next(concat_dim_layout_itr), in EmitFastConcatenate() local
2482 loops.AddLoopsForShapeOnDimensions(output_shape, outer_dims, "concat"); in EmitFastConcatenate()
2489 if (!outer_dims.empty()) { in EmitFastConcatenate()
2534 if (!outer_dims.empty()) { in EmitFastConcatenate()
|
| /external/tensorflow/tensorflow/lite/kernels/ |
| D | scatter_nd.cc | 58 const int outer_dims = indices.DimensionsCount() - 1; in CheckShapes() local
59 for (int i = 0; i < outer_dims; ++i) { in CheckShapes()
63 const int ix = indices.Dims(outer_dims); in CheckShapes()
64 TF_LITE_ENSURE_EQ(context, updates.DimensionsCount() - outer_dims, in CheckShapes()
66 for (int i = 0; i + outer_dims < updates.DimensionsCount(); ++i) { in CheckShapes()
67 TF_LITE_ENSURE_EQ(context, updates.Dims(i + outer_dims), in CheckShapes()
|
| /external/tensorflow/tensorflow/core/kernels/ |
| D | scatter_nd_op.cc | 95 const int64 outer_dims = indices.shape().dims() - 1; in Compute() local
97 for (int i = 0; i < outer_dims; ++i) { in Compute()
101 "Dimensions [0,", outer_dims, in Compute()
103 "] must match dimensions [0,", outer_dims, in Compute()
107 const int64 ix = indices.shape().dim_size(outer_dims); in Compute()
108 OP_REQUIRES(c, updates.shape().dims() - outer_dims == shape.dims() - ix, in Compute()
112 outer_dims, ",", updates.shape().dims(), in Compute()
115 for (int i = 0; i + outer_dims < updates.shape().dims(); ++i) { in Compute()
117 c, updates.shape().dim_size(i + outer_dims) == shape.dim_size(ix + i), in Compute()
120 "] must match dimensions [", outer_dims, ",", in Compute()
[all …]
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/ |
| D | lower_general_dot.cc | 112 llvm::SmallVector<int64_t, 5> outer_dims; in ProcessDotArg() local
115 outer_dims.push_back(it.index()); in ProcessDotArg()
120 return TransposeReshape(arg, loc, outer_dims, contract_dims, shape, in ProcessDotArg()
124 return TransposeReshape(arg, loc, contract_dims, outer_dims, shape, rewriter); in ProcessDotArg()
|
| /external/tensorflow/tensorflow/python/kernel_tests/array_ops/ |
| D | gather_op_test.py | 160 outer_dims = axis
165 dest_slice = ((slice(None),) * outer_dims + (dest_index,) +
167 source_slice = ((slice(None),) * outer_dims + (source_index,) +
229 outer_dims = axis
234 dest_slice = ((slice(None),) * outer_dims + (dest_index,) +
236 source_slice = ((slice(None),) * outer_dims + (source_index,) +
|
| /external/tensorflow/tensorflow/python/ops/signal/ |
| D | fft_ops.py | 99 outer_dims = _math_ops.maximum(0, input_rank - fft_rank)
100 outer_paddings = _array_ops.zeros([outer_dims], fft_length.dtype)
|
| /external/tensorflow/tensorflow/compiler/mlir/lite/ir/ |
| D | tfl_ops.cc | 704 const auto outer_dims = output_type.getShape().take_front(axis); in ConstFoldConcatenateOpDense() local
706 outer_dims.begin(), outer_dims.end(), 1, std::multiplies<int64_t>()); in ConstFoldConcatenateOpDense()
1045 auto outer_dims = indices_type.getRank() - 1; in Verify() local
1046 auto outermost_dim = indices_type.getDimSize(outer_dims); in Verify()
1049 for (auto i = 0; i < outer_dims; i++) { in Verify()
1062 auto output_rank = outermost_dim + updates_type.getRank() - outer_dims; in Verify()
1090 auto corresponding_dim = i - outermost_dim + outer_dims; in Verify()
|
| /external/tensorflow/tensorflow/core/framework/ |
| D | common_shape_fns.cc | 2335 const int64 outer_dims = c->Rank(indices_shape) - 1; in ScatterNdShapeHelper() local
2345 c->Subshape(indices_shape, 0, outer_dims, &prefix_indices)); in ScatterNdShapeHelper()
2348 c->Subshape(updates_shape, 0, outer_dims, &prefix_updates)); in ScatterNdShapeHelper()
2353 "Dimensions [0,", outer_dims, in ScatterNdShapeHelper()
2356 " must match dimensions [0,", outer_dims, in ScatterNdShapeHelper()
2365 c->Subshape(updates_shape, outer_dims, &suffix_updates)); in ScatterNdShapeHelper()
2372 outer_dims, ",", c->Rank(updates_shape), in ScatterNdShapeHelper()
|
| /external/tensorflow/tensorflow/lite/kernels/internal/reference/ |
| D | reference_ops.h | 1147 const int outer_dims = indices_shape.DimensionsCount() - 1; in ScatterNd() local
1148 const int indices_nd = indices_shape.Dims(outer_dims); in ScatterNd()
1150 for (int i = 0; i < outer_dims; ++i) { in ScatterNd()
1153 for (int i = outer_dims; i < updates_dims; ++i) { in ScatterNd()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | algebraic_simplifier.cc | 2443 const int64 outer_dims = std::max(rhs_outer_dims, lhs_outer_dims); in HandleDot() local
2446 outer_dims + dot->dot_dimension_numbers().lhs_batch_dimensions_size()); in HandleDot()
|