Home
last modified time | relevance | path

Searched refs:broadcast_dims (Results 1 – 23 of 23) sorted by relevance

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dxla_broadcast_helper_op.cc45 std::vector<int64> broadcast_dims; in Compile() local
47 &broadcast_dims)); in Compile()
48 if (broadcast_dims.empty()) { in Compile()
64 context, broadcast_dims.size() == min_rank_shape->dims(), in Compile()
68 absl::StrJoin(broadcast_dims, ","), "]; argument shapes: ", in Compile()
70 std::vector<int64> sorted_broadcast_dims = broadcast_dims; in Compile()
72 std::set<int64> dims_set(broadcast_dims.begin(), broadcast_dims.end()); in Compile()
74 dims_set.size() == broadcast_dims.size() && in Compile()
75 broadcast_dims == sorted_broadcast_dims, in Compile()
79 absl::StrJoin(broadcast_dims, ","), "]")); in Compile()
[all …]
Ddiag_op.cc65 std::vector<int64> broadcast_dims(other_dims.begin(), other_dims.end()); in CreateDiagonal() local
66 broadcast_dims.push_back(1LL); in CreateDiagonal()
67 broadcast_dims.push_back(last_dim_size); in CreateDiagonal()
68 xla::XlaOp input_broadcast = xla::Reshape(input, broadcast_dims); in CreateDiagonal()
70 broadcast_dims[broadcast_dims.size() - 2] = last_dim_size; in CreateDiagonal()
72 xla::ShapeUtil::MakeShape(element_type, broadcast_dims); in CreateDiagonal()
Dresampler_ops.cc60 std::vector<int64> broadcast_dims(warp_dims.begin(), warp_dims.end() - 1); in BilinearWeights() local
61 broadcast_dims.push_back(4); in BilinearWeights()
62 broadcast_dims.push_back(2); in BilinearWeights()
64 const int64 broadcast_dims_size = broadcast_dims.size(); in BilinearWeights()
70 xla::BroadcastInDim(first_term, broadcast_dims, last_two_dims_indices); in BilinearWeights()
76 std::vector<int64> ratio_broadcast_indices(broadcast_dims.size()); in BilinearWeights()
81 xla::BroadcastInDim(ratio, broadcast_dims, ratio_broadcast_indices); in BilinearWeights()
92 xla::BroadcastInDim(sign_change, broadcast_dims, last_two_dims_indices); in BilinearWeights()
Dimage_ops.cc205 std::vector<int64> broadcast_dims(input_shape.dims() - 2); in Compile() local
206 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in Compile()
207 broadcast_dims.back() = channel_dim; in Compile()
211 broadcast_dims); in Compile()
/external/tensorflow/tensorflow/core/kernels/
Dadjust_contrast_op.h47 Eigen::array<int, 4> broadcast_dims; in operator() local
48 broadcast_dims[0] = 1; in operator()
49 broadcast_dims[1] = height; in operator()
50 broadcast_dims[2] = width; in operator()
51 broadcast_dims[3] = 1; in operator()
61 broadcast_dims; in operator() local
62 broadcast_dims.set(1, height); in operator()
63 broadcast_dims.set(2, width); in operator()
75 .broadcast(broadcast_dims); in operator()
109 Eigen::array<int, 4> broadcast_dims; in operator() local
[all …]
Dcwise_op_gpu_select.cu.cc49 Eigen::array<int, 1> broadcast_dims{size}; in operator ()() local
52 .broadcast(broadcast_dims) in operator ()()
68 Eigen::array<int, 2> broadcast_dims{{ 1, all_but_batch }}; in operator ()() local
71 Eigen::IndexList<Eigen::type2index<1>, int> broadcast_dims; in operator ()() local
72 broadcast_dims.set(1, all_but_batch); in operator ()()
87 .broadcast(broadcast_dims) in operator ()()
Dgather_nd_op_cpu_impl.h108 Eigen::array<Eigen::DenseIndex, 1> broadcast_dims{{ batch_size }};
111 Eigen::IndexList<Eigen::DenseIndex> broadcast_dims;
112 broadcast_dims.set(0, batch_size);
131 .broadcast(broadcast_dims)
Dcwise_op_select.cc269 Eigen::array<Eigen::DenseIndex, 2> broadcast_dims{{1, all_but_batch}}; in operator ()() local
272 Eigen::IndexList<Eigen::type2index<1>, Eigen::DenseIndex> broadcast_dims; in operator ()() local
273 broadcast_dims.set(1, all_but_batch); in operator ()()
280 .broadcast(broadcast_dims) in operator ()()
Dfill_functor.cc153 Eigen::array<int, 1> broadcast_dims{size}; in operator ()() local
155 To32Bit(out).device(d) = in.reshape(rank1).broadcast(broadcast_dims); in operator ()()
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dsvd.cc142 std::vector<int64> broadcast_dims(num_dims - 1); in HouseRow() local
143 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in HouseRow()
145 auto mu = Mul(sigma, Sqrt(Square(Div(x_0j, sigma, broadcast_dims)) + one), in HouseRow()
146 broadcast_dims); in HouseRow()
150 -Mul(sigma, Div(sigma, Add(x_0j, mu), broadcast_dims), broadcast_dims)); in HouseRow()
153 (Square(Div(sigma, v_0j, broadcast_dims)) + one)); in HouseRow()
156 BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), broadcast_dims), v, in HouseRow()
160 beta = Select(Lt(Add(sigma, ZerosLike(beta), broadcast_dims), eps), in HouseRow()
208 std::vector<int64> broadcast_dims(num_dims - 1); in HouseCol() local
209 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in HouseCol()
[all …]
Darithmetic.cc144 std::vector<int64> broadcast_dims(input_shape.rank() - 1); in ArgMinMax() local
145 std::iota(broadcast_dims.begin(), broadcast_dims.begin() + axis, 0); in ArgMinMax()
146 std::iota(broadcast_dims.begin() + axis, broadcast_dims.end(), axis + 1); in ArgMinMax()
149 ConvertElementType(Eq(input, input_max, broadcast_dims), output_type); in ArgMinMax()
Dself_adjoint_eig.cc165 std::vector<int64> broadcast_dims(batch_dims.size()); in Update() local
166 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in Update()
167 broadcast_dims.push_back(num_dims - 1); in Update()
178 broadcast_dims); in Update()
185 broadcast_dims); in Update()
346 std::vector<int64> broadcast_dims(num_dims - 1); in SortByEigenvalues() local
347 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in SortByEigenvalues()
348 broadcast_dims[num_dims - 2] = num_dims - 1; in SortByEigenvalues()
349 result.w = BroadcastInDim(result.w, dimensions, broadcast_dims); in SortByEigenvalues()
Dsvd_test.cc88 std::vector<int64> broadcast_dims(num_dims - 1); in ComputeMatmulUDVT() local
89 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in ComputeMatmulUDVT()
90 broadcast_dims[num_dims - 2] = num_dims - 1; in ComputeMatmulUDVT()
91 return BatchDot(Mul(u, d, broadcast_dims), TransposeInMinorDims(v), in ComputeMatmulUDVT()
Dself_adjoint_eig_test.cc106 std::vector<int64> broadcast_dims(shape.rank() - 1); in ComputeMatmulVWVt() local
107 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in ComputeMatmulVWVt()
109 broadcast_dims[shape.rank() - 2] = shape.rank() - 1; in ComputeMatmulVWVt()
110 auto vw = Mul(result.v, BroadcastInDim(result.w, out_dims, broadcast_dims)); in ComputeMatmulVWVt()
/external/tensorflow/tensorflow/compiler/tf2xla/lib/
Dbroadcast.cc46 std::vector<int64> broadcast_dims; in BroadcastTo() local
60 broadcast_dims.push_back(broadcast_shape.size()); in BroadcastTo()
77 absl::c_reverse(broadcast_dims); in BroadcastTo()
79 for (int64& broadcast_dim : broadcast_dims) { in BroadcastTo()
84 xla::BroadcastInDim(input, broadcast_shape, broadcast_dims); in BroadcastTo()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_helpers.cc90 std::vector<int64> broadcast_dims(indices_shape.dims()); in OneHot() local
91 std::iota(broadcast_dims.begin(), broadcast_dims.begin() + axis, 0); in OneHot()
92 std::iota(broadcast_dims.begin() + axis, broadcast_dims.end(), axis + 1); in OneHot()
102 xla::Eq(indices, xla::Iota(builder, iota_shape, axis), broadcast_dims), in OneHot()
/external/tensorflow/tensorflow/compiler/tf2xla/python/
Dxla.py114 def broadcasting_binary_op_wrapper(x, y, broadcast_dims=None, name=None): argument
116 broadcast_dims = broadcast_dims or []
117 broadcast_dims = ops.convert_to_tensor(broadcast_dims, dtypes.int64)
121 x, y = gen_xla_ops.xla_broadcast_helper(x, y, broadcast_dims)
/external/tensorflow/tensorflow/compiler/xla/service/
Dconvolution_group_converter.cc362 std::vector<int64> broadcast_dims; in HandleConvolution() local
367 broadcast_dims.push_back(i); in HandleConvolution()
370 expanded_filter_shape, reshaped_filter, broadcast_dims)); in HandleConvolution()
Dindexed_array_analysis.cc876 absl::Span<const int64> broadcast_dims = broadcast_instr->dimensions(); in ComputeArrayForElementwiseBinaryOp() local
878 return absl::c_find(broadcast_dims, output_dim) == broadcast_dims.end(); in ComputeArrayForElementwiseBinaryOp()
893 for (int64 broadcast_dim : broadcast_dims) { in ComputeArrayForElementwiseBinaryOp()
Dalgebraic_simplifier.cc1286 auto broadcast_dims = [](int64 rank, int64 non_broadcast_dim) { in HandleDotStrengthReduction() local
1342 broadcast_dims(rhs_rank, rhs_kept_dim)); in HandleDotStrengthReduction()
1353 broadcast_dims(lhs_rank, lhs_kept_dim)); in HandleDotStrengthReduction()
Dalgebraic_simplifier_test.cc2376 const std::vector<int64> broadcast_dims = in TEST_F() local
2378 EXPECT_EQ(1, broadcast_dims.size()); in TEST_F()
2379 EXPECT_THAT(broadcast_dims[0], ::testing::AnyOf(1, 2, 3)); in TEST_F()
/external/tensorflow/tensorflow/contrib/losses/python/losses/
Dloss_ops.py167 broadcast_dims = array_ops.slice(
169 num_to_broadcast = math_ops.cast(math_ops.reduce_prod(broadcast_dims),
/external/tensorflow/tensorflow/compiler/tests/
Dxla_ops_test.py60 lambda x, y: xla.add(x, y, broadcast_dims=(0,)),
66 lambda x, y: xla.add(x, y, broadcast_dims=(1,)),