/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/lib/ |
H A D | svd.cc | 119 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 D | lu_decomposition.cc | 34 const int num_batch_dims = a_shape.dimensions().size() - 2; in LuDecomposition() local
|
H A D | self_adjoint_eig.cc | 69 const int num_batch_dims = a_shape.dimensions().size() - 2; in SelfAdjointEig() local
|
/aosp_15_r20/external/pytorch/aten/src/ATen/ |
H A D | LegacyBatchingRegistrations.cpp | 279 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 D | LegacyBatchedFallback.cpp | 167 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 D | LegacyVmapTransforms.cpp | 205 int64_t num_batch_dims = collective_levels.count(); in logicalToPhysical() local
|
H A D | NamedTensorUtils.cpp | 298 static int64_t num_batch_dims(DimnameList names) { in num_batch_dims() function
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
H A D | qr_expander.cc | 213 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 D | dot_decomposer.cc | 44 const int64_t num_batch_dims = original_dnums.lhs_batch_dimensions_size(); in CanonicalizeDot() local
|
H A D | eigh_expander.cc | 459 const int64_t num_batch_dims = num_dims - 2; in BuildEigh() local
|
H A D | elemental_ir_emitter.cc | 2413 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 D | BatchedFallback.cpp | 172 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 D | LegacyBatchingRegistrations.cpp | 303 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 D | matmul_utils.cc | 176 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 D | gemm_broadcast_folding_rewriter.cc | 55 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 D | dot_op_emitter.cc | 1328 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 D | test_mps_linear.py | 213 num_batch_dims=1, argument
|
/aosp_15_r20/external/executorch/backends/xnnpack/test/ops/ |
H A D | linear.py | 465 num_batch_dims=1, argument
|
/aosp_15_r20/external/tensorflow/tensorflow/python/ops/ |
H A D | nn_ops.py | 342 num_batch_dims=1): argument 754 num_batch_dims=1): argument
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
H A D | legalize_tf.cc | 442 bool transpose_rhs, int64_t num_batch_dims, in BatchDot() 6739 const int64_t num_batch_dims = num_dims - 2; in QRBlock() local
|