Home
last modified time | relevance | path

Searched defs:broadcast_dims (Results 1 – 25 of 25) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/
H A Dcwise_op_gpu_select.cu.cc66 Eigen::array<int, 1> broadcast_dims{size}; in operator ()() local
88 Eigen::IndexList<Eigen::type2index<1>, int> broadcast_dims; in operator ()() local
H A Dcwise_op_select.cc361 Eigen::IndexList<Eigen::type2index<1>, Eigen::DenseIndex> broadcast_dims; in operator ()() local
/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/image/
H A Dadjust_contrast_op.h47 broadcast_dims; in operator() local
95 broadcast_dims; in operator() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/lib/
H A Dsvd.cc138 std::vector<int64_t> broadcast_dims(num_dims - 1); in HouseRow() local
203 std::vector<int64_t> broadcast_dims(num_dims - 1); in HouseCol() local
480 std::vector<int64_t> broadcast_dims(batch_dims.size()); in OneSidedJacobiUpdate() local
742 std::vector<int64_t> broadcast_dims(num_dims - 1); in SortBySingularValuesAndPostProcessing() local
H A Dbroadcast.cc45 std::vector<int64_t> broadcast_dims; in BroadcastTo() local
H A Dmatrix.cc183 std::vector<int64_t> broadcast_dims(n_dims - 1); in SetMatrixDiagonal() local
274 std::vector<int64_t> broadcast_dims; in EinsumDiagonalLabels() local
H A Dsvd_test.cc92 std::vector<int64_t> broadcast_dims(num_dims - 1); in ComputeMatmulUDVT() local
H A Dself_adjoint_eig_test.cc122 std::vector<int64_t> broadcast_dims(shape.rank() - 1); in ComputeMatmulVWVt() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Deigh_expander.cc150 std::vector<int64_t> broadcast_dims(shape.dimensions().size() - 1); in ApplyJacobiRotationOverRows() local
170 std::vector<int64_t> broadcast_dims(shape.dimensions().size() - 1); in ApplyJacobiRotationOverCols() local
361 std::vector<int64_t> broadcast_dims(num_dims - 1); in SortByEigenvalues() local
H A Dconvolution_group_converter.cc468 std::vector<int64_t> broadcast_dims; in HandleConvolution() local
H A Dspace_to_batch_converter.cc1444 std::vector<int64_t> broadcast_dims; in PropagateOnBroadcast() local
1519 auto broadcast_dims = broadcast->dimensions(); in IsBroadcastPropagatable() local
2351 std::vector<int64_t> broadcast_dims(new_space_dims.begin(), in SelectValidPortion() local
H A Dwhile_loop_concat_code_motion.cc872 std::vector<int64_t> broadcast_dims; in RewriteLoopWithConcatGroups() local
H A Dindexed_array_analysis.cc880 absl::Span<const int64_t> broadcast_dims = broadcast_instr->dimensions(); in ComputeArrayForElementwiseBinaryOp() local
H A Ddynamic_padder.cc1481 std::vector<int64_t> broadcast_dims; in RewriteDynamicBinaryOp() local
H A Dalgebraic_simplifier.cc1525 DimensionVector broadcast_dims; in HandleConcatenate() local
5707 DimensionVector broadcast_dims(nested_root->shape().dimensions_size()); in HandleReduceWindow() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
H A Dxla_broadcast_helper_op.cc45 std::vector<int64_t> broadcast_dims; in Compile() local
H A Dresampler_ops.cc62 std::vector<int64_t> broadcast_dims(warp_dims.begin(), warp_dims.end() - 1); in BilinearWeights() local
H A Dimage_ops.cc214 std::vector<int64_t> broadcast_dims(input_shape.dims() - 2); in Compile() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dgemm_rewriter.cc327 absl::Span<const int64_t> broadcast_dims = broadcast_bias->dimensions(); in FuseVectorBiasAdd() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/
H A Dxla_helpers.cc96 std::vector<int64_t> broadcast_dims(indices_shape.dims()); in OneHot() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
H A Dlegalize_tf_collective.cc167 auto broadcast_dims = GetI64ElementsAttr({}, &builder); in ConvertAllReduce() local
H A Dlegalize_tf.cc354 DenseIntElementsAttr broadcast_dims, in StaticBinaryBroadcast()
410 auto broadcast_dims = GetI64ElementsAttr({feature_dim}, &builder); in Broadcast1DToFeatureDim() local
435 auto broadcast_dims = GetI64ElementsAttrForSeq(0, rank, &builder); in BroadcastToShapeOf() local
4182 auto broadcast_dims = GetI64ElementsAttr({}, &rewriter); in matchAndRewrite() local
5311 llvm::SmallVector<int64_t, 4> broadcast_dims(indices_shape.size()); in matchAndRewrite() local
6144 auto broadcast_dims = GetI64ElementsAttr({}, &rewriter); in matchAndRewrite() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/python/
H A Dxla.py122 def broadcasting_binary_op_wrapper(x, y, broadcast_dims=None, name=None): argument
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dspmd_partitioner.cc261 absl::flat_hash_set<int64_t> broadcast_dims; in AddInstruction() local
1184 std::vector<int64_t> broadcast_dims; in ReplicatePartial() local
1954 if (auto broadcast_dims = state_.b->BroadcastDimsForCreatedHlo(hlo())) { in ReshardWithCollectivePermute() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/
H A DTensorShape.cpp492 std::vector<int64_t> broadcast_dims; in sparse_broadcast_to() local