Home
last modified time | relevance | path

Searched defs:num_batch_dims (Results 1 – 20 of 20) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/lib/
H A Dsvd.cc119 const int64_t num_batch_dims = num_dims - 2; in HouseRow() local
184 const int64_t num_batch_dims = num_dims - 2; in HouseCol() local
258 const int64_t num_batch_dims = num_dims - 2; in HouseHolderBidiagonalization() local
460 const int64_t num_batch_dims = num_dims - 2; in OneSidedJacobiUpdate() local
839 const int64_t num_batch_dims = num_dims - 2; in SVD() local
H A Dlu_decomposition.cc34 const int num_batch_dims = a_shape.dimensions().size() - 2; in LuDecomposition() local
H A Dself_adjoint_eig.cc69 const int num_batch_dims = a_shape.dimensions().size() - 2; in SelfAdjointEig() local
/aosp_15_r20/external/pytorch/aten/src/ATen/
H A DLegacyBatchingRegistrations.cpp279 int64_t num_batch_dims = self_physical.numBatchDims(); in squeeze_batching_rule() local
368 static int64_t getGradInputPhysicalDim(int64_t dim, IntArrayRef input_sizes, int64_t num_batch_dims in getGradInputPhysicalDim()
491 static void checkBatchDimsAtFrontInLayout(IntArrayRef physical_strides, int64_t num_batch_dims) { in checkBatchDimsAtFrontInLayout()
526 int64_t num_batch_dims, in checkBasicAsStridedValidForSlice()
576 int64_t num_batch_dims = self_physical_view.numBatchDims(); in _new_zeros_with_same_feature_meta_batching_rule() local
619 auto num_batch_dims = physical_view.numBatchDims(); in as_strided_batching_rule() local
H A DLegacyBatchedFallback.cpp167 auto num_batch_dims = input_physical_views.front().numBatchDims(); in batchedTensorInplaceForLoopFallback() local
301 auto num_batch_dims = input_physical_views.front().numBatchDims(); in batchedTensorForLoopFallback() local
H A DLegacyVmapTransforms.cpp205 int64_t num_batch_dims = collective_levels.count(); in logicalToPhysical() local
H A DNamedTensorUtils.cpp298 static int64_t num_batch_dims(DimnameList names) { in num_batch_dims() function
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dqr_expander.cc213 const int64_t num_batch_dims = num_dims - 2; in QrBlock() local
393 const int64_t num_batch_dims = num_dims - 2; in BuildQrDecomposition() local
464 const int64_t num_batch_dims = num_dims - 2; in ProductOfElementaryHouseholderReflectors() local
H A Ddot_decomposer.cc44 const int64_t num_batch_dims = original_dnums.lhs_batch_dimensions_size(); in CanonicalizeDot() local
H A Deigh_expander.cc459 const int64_t num_batch_dims = num_dims - 2; in BuildEigh() local
H A Delemental_ir_emitter.cc2413 int64_t num_batch_dims = dim_numbers.rhs_batch_dimensions_size(); in EmitElementalDot() local
/aosp_15_r20/external/pytorch/aten/src/ATen/functorch/
H A DBatchedFallback.cpp172 auto num_batch_dims = input_physical_views.front().numBatchDims(); in batchedTensorInplaceForLoopFallback() local
317 auto num_batch_dims = input_physical_views.front().numBatchDims(); in batchedTensorForLoopFallback() local
H A DLegacyBatchingRegistrations.cpp303 int64_t num_batch_dims, in checkBasicAsStridedValidForSlice()
370 auto num_batch_dims = physical_view.numBatchDims(); in as_strided_batching_rule() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dmatmul_utils.cc176 size_t num_batch_dims = std::max(lhs_num_batch_dims, rhs_num_batch_dims); in For() local
258 int64_t num_batch_dims = in For() local
H A Dgemm_broadcast_folding_rewriter.cc55 int num_batch_dims = dim_nums->lhs_batch_dimensions_size(); in HandleCustomCall() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/cpu/
H A Ddot_op_emitter.cc1328 int64_t num_batch_dims = in PotentiallyImplementedAsEigenMatmul() local
1403 int64_t num_batch_dims = in EmitBatchDotOperation() local
/aosp_15_r20/external/executorch/backends/apple/mps/test/
H A Dtest_mps_linear.py213 num_batch_dims=1, argument
/aosp_15_r20/external/executorch/backends/xnnpack/test/ops/
H A Dlinear.py465 num_batch_dims=1, argument
/aosp_15_r20/external/tensorflow/tensorflow/python/ops/
H A Dnn_ops.py342 num_batch_dims=1): argument
754 num_batch_dims=1): argument
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
H A Dlegalize_tf.cc442 bool transpose_rhs, int64_t num_batch_dims, in BatchDot()
6739 const int64_t num_batch_dims = num_dims - 2; in QRBlock() local