Home
last modified time | relevance | path

Searched defs:featureDim (Results 1 – 3 of 3) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/native/sparse/cuda/
H A DSparseCUDAApplyUtils.cuh222 int featureDim = startFeature + ii * C10_WARP_SIZE; in C10_LAUNCH_BOUNDS_1() local
232 int featureDim = startFeature + ii * C10_WARP_SIZE; in C10_LAUNCH_BOUNDS_1() local
272 int featureDim = startFeature + ii * C10_WARP_SIZE; in coalesceValuesKernel() local
282 int featureDim = startFeature + ii * C10_WARP_SIZE; in coalesceValuesKernel() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/
H A DEmbeddingBag.cu77 int64_t featureDim = (chunk % chunksPerBag) * blockDim.x + threadIdx.x; in EmbeddingBag_updateOutputKernel_max() local
129 int64_t featureDim = (chunk % chunksPerBag) * blockDim.x + threadIdx.x; in EmbeddingBag_updateOutputKernel_sum_mean() local
253 int64_t featureDim = (chunk % chunksPerBag) * blockDim.x + threadIdx.x; in EmbeddingBag_accGradParametersKernel_max() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/lib/Dialect/mhlo/transforms/
H A Dunfuse_batch_norm.cc40 Value value1d, Value shapeValue, int64_t featureDim, in broadcastToFeatureDim()
120 int64_t featureDim = bnOp.feature_index(); in matchAndRewrite() local