Home
last modified time | relevance | path

Searched refs:BroadcastInDim (Results 1 – 25 of 29) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dresampler_ops.cc72 xla::BroadcastInDim(first_term, broadcast_dims, last_two_dims_indices); in BilinearWeights()
83 xla::BroadcastInDim(ratio, broadcast_dims, ratio_broadcast_indices); in BilinearWeights()
94 xla::BroadcastInDim(sign_change, broadcast_dims, last_two_dims_indices); in BilinearWeights()
228 xla::BroadcastInDim(is_in_bound, result_dims, broadcasted_dims); in BoundSamples()
292 auto broadcast_reshaped_weights = xla::BroadcastInDim( in CalculateGradData()
298 XlaOp broadcast_grad_output = xla::BroadcastInDim( in CalculateGradData()
414 xla::BroadcastInDim( in CalculateGradWarp()
422 xla::BroadcastInDim( in CalculateGradWarp()
430 xla::BroadcastInDim( in CalculateGradWarp()
438 xla::BroadcastInDim( in CalculateGradWarp()
[all …]
Dxla_broadcast_helper_op.cc93 lhs = xla::BroadcastInDim(lhs, broadcast_shape, broadcast_dims); in Compile()
95 rhs = xla::BroadcastInDim(rhs, broadcast_shape, broadcast_dims); in Compile()
Dimage_resize_ops.cc191 xla::BroadcastInDim(make_kernel_func(builder, type, kernel_size[1]), in MakeGeneralResizeKernel()
210 return xla::BroadcastInDim(make_kernel_func(builder, type, kernel_size[dim]), in MakeGeneralResizeKernelInDim()
233 return xla::BroadcastInDim(input, broadcast_shape.dimensions(), in BroadcastSpatialDimensions()
559 xla::BroadcastInDim(h_span_start, {out_size[0], out_size[1], 1}, {0}); in GeneralCompile()
595 xla::BroadcastInDim(w_span_start, {out_size[0], out_size[1], 1}, {1}); in GeneralCompile()
617 w_sub = xla::BroadcastInDim(w_sub, {out_size[1], w_span_size}, {0}); in GeneralCompile()
636 h_sub = xla::BroadcastInDim(h_sub, {out_size[0], h_span_size}, {0}); in GeneralCompile()
Ddiag_op.cc72 xla::BroadcastInDim(input, out_dim_sizes, broadcast_dimensions); in CreateDiagonal()
Dselect_op.cc71 cond_handle = xla::BroadcastInDim(cond_handle, then_shape.dim_sizes(), in Compile()
Dfake_quantize_ops.cc267 return xla::BroadcastInDim(op, input_dimensions, in Compile()
328 return xla::BroadcastInDim(op, input_dimensions, in Compile()
Ddynamic_partition_op.cc139 partitions = xla::BroadcastInDim(partitions, data_shape.dimensions(), in Compile()
Dquantize_and_dequantize_op.cc169 return xla::BroadcastInDim(op, input_dimensions, {axis_}); in Compile()
Dmatrix_diag_ops.cc223 xla::XlaOp diag_broadcast = xla::BroadcastInDim( in SetMatrixDiag()
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dslicing.cc191 BroadcastInDim(index, sizes, index_broadcast_dims), in TorchGather()
195 mask, BroadcastInDim(input, sizes, input_broadcast_dims), in TorchGather()
249 Eq(BroadcastInDim(index, sizes, index_broadcast_dims), in TorchScatterDense()
253 Select(mask, BroadcastInDim(src, sizes, index_broadcast_dims), in TorchScatterDense()
Dbroadcast.cc82 XlaOp output = BroadcastInDim(input, broadcast_shape, broadcast_dims); in BroadcastTo()
Dsvd.cc152 BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), broadcast_dims), v, in HouseRow()
218 BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), broadcast_dims), v, in HouseCol()
578 BroadcastInDim(diag, shape.dimensions(), broadcasted_dims); in ComputeToleranceComparison()
581 BroadcastInDim(diag, shape.dimensions(), broadcasted_dims); in ComputeToleranceComparison()
757 d = BroadcastInDim(d, dimensions, broadcast_dims); in SortBySingularValuesAndPostProcessing()
Dmatrix.cc131 XlaOp diag_base_indices = BroadcastInDim(Iota(builder, S32, diag_len), in GetMatrixDiagonalViaGather()
207 BroadcastInDim(diag, shape.dimensions(), broadcast_dims), in SetMatrixDiagonal()
351 x = BroadcastInDim(x, broadcast_sizes, labels->at(2)); in EinsumInverseDiagonal()
Dself_adjoint_eig_test.cc128 BroadcastInDim(ConvertElementType(result.w, shape.element_type()), in ComputeMatmulVWVt()
/external/tensorflow/tensorflow/compiler/xla/service/
Deigh_expander.cc152 auto c = BroadcastInDim(rotation.c, shape.dimensions(), broadcast_dims); in ApplyJacobiRotationOverRows()
153 auto s = BroadcastInDim(rotation.s, shape.dimensions(), broadcast_dims); in ApplyJacobiRotationOverRows()
173 auto c = BroadcastInDim(rotation.c, shape.dimensions(), broadcast_dims); in ApplyJacobiRotationOverCols()
174 auto s = BroadcastInDim(rotation.s, shape.dimensions(), broadcast_dims); in ApplyJacobiRotationOverCols()
364 w = BroadcastInDim(w, dimensions, broadcast_dims); in SortByEigenvalues()
Dqr_expander.cc269 new_x = BroadcastInDim(new_x, ConcatVectors(batch_dims, {m, n}), in QrBlock()
344 auto tau_scale = BroadcastInDim(-taus, ConcatVectors(batch_dims, {1, n}), in CompactWYRepresentation()
Dcholesky_expander.cc211 BroadcastInDim(seen_error, a_shape.dimensions(), error_dim_indices), in BuildCholesky()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dbroadcast_simple_test.cc162 BroadcastInDim(ConstantR1<float>(&b, {1, 2}), {2, 2}, {1});
175 BroadcastInDim(ConstantR1<float>(&b, {1, 2}), {2, 2}, {0});
188 BroadcastInDim(ConstantR2<float>(&b, {{1.0, 5.0}, {2.0, 6.0}}), {2, 2, 2},
206 BroadcastInDim(ConstantR2<float>(&b, {{1.0, 5.0}, {2.0, 6.0}}), {2, 2, 2},
224 BroadcastInDim(ConstantR1<float>(&b, {1, 2}), {3, 2}, {1});
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder_test.cc316 TEST_F(XlaBuilderTest, BroadcastInDim) { in TEST_F() argument
319 BroadcastInDim(x, {2, 4, 3}, in TEST_F()
329 BroadcastInDim(x, {2, 3, 4}, in TEST_F()
339 BroadcastInDim(x, {-3, 3, 4}, in TEST_F()
752 BroadcastInDim(gte, /*out_dim_size=*/{3, 5, 4}, in TEST_F()
Dxla_builder.h458 XlaOp BroadcastInDim(XlaOp operand,
1105 friend XlaOp BroadcastInDim(
1749 XlaOp BroadcastInDim(XlaOp operand,
Dxla_builder.cc956 XlaOp XlaBuilder::BroadcastInDim( in BroadcastInDim() function in xla::XlaBuilder
4058 XlaOp BroadcastInDim(const XlaOp operand, in BroadcastInDim() function
4061 return operand.builder()->BroadcastInDim(operand, out_dim_size, in BroadcastInDim()
/external/tensorflow/tensorflow/compiler/xla/python/xla_extension/
Dops.pyi124 def BroadcastInDim(operand: XlaOp,
/external/tensorflow/tensorflow/compiler/xla/python/
Dops.cc116 ops.def("BroadcastInDim", &BroadcastInDim, py::arg("operand"), in BuildOpsSubmodule()
Dxla_client_test.py1441 ops.BroadcastInDim(ops.Constant(c, NumpyArrayS32([1, 2])), [2, 2], [0])
1443 ops.BroadcastInDim(ops.Constant(c, NumpyArrayS32([1, 2])), [2, 2], [1])
/external/tensorflow/tensorflow/compiler/xla/g3doc/
Doperation_semantics.md426 ## BroadcastInDim section in Operation Semantics
429 [`XlaBuilder::BroadcastInDim`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_b…
433 <b> `BroadcastInDim(operand, out_dim_size, broadcast_dimensions)` </b>

12