Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dresampler_ops.cc70 xla::BroadcastInDim(first_term, broadcast_dims, last_two_dims_indices); 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()
225 xla::BroadcastInDim(is_in_bound, result_dims, broadcasted_dims); in BoundSamples()
289 auto broadcast_reshaped_weights = xla::BroadcastInDim( in CalculateGradData()
295 XlaOp broadcast_grad_output = xla::BroadcastInDim( in CalculateGradData()
411 xla::BroadcastInDim( in CalculateGradWarp()
419 xla::BroadcastInDim( in CalculateGradWarp()
427 xla::BroadcastInDim( in CalculateGradWarp()
435 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()
Dselect_op.cc70 cond_handle = xla::BroadcastInDim(cond_handle, then_shape.dim_sizes(), in Compile()
Ddiag_op.cc71 xla::BroadcastInDim(input, out_dim_sizes, broadcast_dimensions); in CreateDiagonal()
Dimage_resize_ops.cc185 xla::BroadcastInDim(make_kernel_func(builder, type, kernel_size[1]), in MakeGeneralResizeKernel()
204 return xla::BroadcastInDim(make_kernel_func(builder, type, kernel_size[dim]), in MakeGeneralResizeKernelInDim()
227 return xla::BroadcastInDim(input, broadcast_shape.dimensions(), in BroadcastSpatialDimensions()
Dquantize_and_dequantize_op.cc169 return xla::BroadcastInDim(op, input_dimensions, {axis_}); in Compile()
Dmatrix_diag_ops.cc220 xla::XlaOp diag_broadcast = xla::BroadcastInDim( in SetMatrixDiag()
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dslicing.cc190 BroadcastInDim(index, sizes, index_broadcast_dims), in TorchGather()
194 mask, BroadcastInDim(input, sizes, input_broadcast_dims), in TorchGather()
246 Eq(BroadcastInDim(index, sizes, index_broadcast_dims), in TorchScatterDense()
250 Select(mask, BroadcastInDim(src, sizes, index_broadcast_dims), in TorchScatterDense()
Dsvd.cc151 BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), broadcast_dims), v, in HouseRow()
217 BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), broadcast_dims), v, in HouseCol()
577 BroadcastInDim(diag, shape.dimensions(), broadcasted_dims); in ComputeToleranceComparison()
580 BroadcastInDim(diag, shape.dimensions(), broadcasted_dims); in ComputeToleranceComparison()
742 d = BroadcastInDim(d, dimensions, broadcast_dims); in SortBySingularValuesAndPostProcessing()
Dmatrix.cc130 XlaOp diag_base_indices = BroadcastInDim(Iota(builder, S32, diag_len), in GetMatrixDiagonalViaGather()
206 BroadcastInDim(diag, shape.dimensions(), broadcast_dims), in SetMatrixDiagonal()
320 x = BroadcastInDim(x, broadcast_sizes, labels->at(2)); in EinsumInverseDiagonal()
Dself_adjoint_eig_test.cc110 auto vw = Mul(result.v, BroadcastInDim(result.w, out_dims, broadcast_dims)); in ComputeMatmulVWVt()
Dself_adjoint_eig.cc349 result.w = BroadcastInDim(result.w, dimensions, broadcast_dims); in SortByEigenvalues()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dbroadcast_simple_test.cc164 BroadcastInDim(ConstantR1<float>(&b, {1, 2}), {2, 2}, {1});
177 BroadcastInDim(ConstantR1<float>(&b, {1, 2}), {2, 2}, {0});
190 BroadcastInDim(ConstantR2<float>(&b, {{1.0, 5.0}, {2.0, 6.0}}), {2, 2, 2},
208 BroadcastInDim(ConstantR2<float>(&b, {{1.0, 5.0}, {2.0, 6.0}}), {2, 2, 2},
226 BroadcastInDim(ConstantR1<float>(&b, {1, 2}), {3, 2}, {1});
/external/tensorflow/tensorflow/compiler/tf2xla/lib/
Dbroadcast.cc87 xla::BroadcastInDim(input, broadcast_shape, broadcast_dims); in BroadcastTo()
/external/tensorflow/tensorflow/compiler/xla/service/
Dqr_expander.cc245 new_x = BroadcastInDim(new_x, ConcatVectors(batch_dims, {m, n}), in QrBlock()
320 auto tau_scale = BroadcastInDim(-taus, ConcatVectors(batch_dims, {1, n}), in CompactWYRepresentation()
Dcholesky_expander.cc213 BroadcastInDim(seen_error, a_shape.dimensions(), error_dim_indices), in BuildCholesky()
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder_test.cc314 TEST_F(XlaBuilderTest, BroadcastInDim) { in TEST_F() argument
317 BroadcastInDim(x, {2, 4, 3}, in TEST_F()
327 BroadcastInDim(x, {2, 3, 4}, in TEST_F()
337 BroadcastInDim(x, {-3, 3, 4}, in TEST_F()
682 BroadcastInDim(gte, /*out_dim_size=*/{3, 5, 4}, in TEST_F()
Dxla_builder.h455 XlaOp BroadcastInDim(XlaOp operand,
1061 friend XlaOp BroadcastInDim(
1651 XlaOp BroadcastInDim(XlaOp operand, const absl::Span<const int64> out_dim_size,
Dxla_builder.cc890 XlaOp XlaBuilder::BroadcastInDim( in BroadcastInDim() function in xla::XlaBuilder
4102 XlaOp BroadcastInDim(const XlaOp operand, in BroadcastInDim() function
4105 return operand.builder()->BroadcastInDim(operand, out_dim_size, in BroadcastInDim()
/external/tensorflow/tensorflow/compiler/xla/python/
Dops.cc89 ops.def("BroadcastInDim", &BroadcastInDim, py::arg("operand"), in BuildOpsSubmodule()
Dxla_client_test.py1201 ops.BroadcastInDim(ops.Constant(c, NumpyArrayS32([1, 2])), [2, 2], [0])
1203 ops.BroadcastInDim(ops.Constant(c, NumpyArrayS32([1, 2])), [2, 2], [1])
/external/tensorflow/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/
Dhlo_ops_base.td900 and exposed in the XLA client BroadcastInDim method.
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dmlir_hlo_to_hlo.cc612 BroadcastInDim(operand, Convert_ArrayRef(type.getShape()), in ExportXlaOp()
/external/tensorflow/tensorflow/compiler/xla/g3doc/
Doperation_semantics.md392 ## BroadcastInDim section in Operation Semantics
395 [`XlaBuilder::BroadcastInDim`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_b…
399 <b> `BroadcastInDim(operand, out_dim_size, broadcast_dimensions)` </b>