Нет (единого) источника для bmm
как такового. Из файла Readme ATen :
ATen «собственные» функции - это современный механизм для добавления операторов и функций в ATen (они являются «собственными» в отличие от устаревших функций, которые связаны через метаданные TH / TH C cwrap). Собственные функции объявлены в native_functions.yaml и имеют реализации, определенные в одном из файлов cpp в этом каталоге.
bmm
объявлено в aten \ src \ ATen \ native \ native_functions.yaml :
- func: bmm(Tensor self, Tensor mat2) -> Tensor
use_c10_dispatcher: full
variants: function, method
dispatch:
CPU: bmm_cpu
CUDA: bmm_cuda
SparseCPU: bmm_sparse_cpu
SparseCUDA: bmm_sparse_cuda
supports_named_tensor: True
Реализации (например, bmm_cpu
) находятся в aten\src\ATen\native\LinearAlgebra.cpp
.