大多数 pytorch 后端代码都是用 C++ 和/或 CUDA 实现的。要查看它，您需要在源代码中找到适当的入口点。有几种方法可以做到这一点，但我发现最简单的方法是在 github 上搜索关键字，而无需自己下载所有代码。
例如，如果您去github.com https://www.github.com以及寻找unique_dim repo:pytorch/pytorch，然后单击左侧的“代码”选项卡，您应该很快会找到以下内容。
From 火炬/jit/_builtins.py:103 https://github.com/pytorch/pytorch/blob/dc5cda0cca429a3080f63e7b30b3e87d18df8601/torch/jit/_builtins.py#L103
17: _builtin_ops = [
...
103: (torch._VF.unique_dim, "aten::unique_dim"),
从这个以及对代码的进一步分析我们可以得出结论：torch._VF.unique_dim实际上是调用aten::unique_dimATen 库中的函数。
像大多数功能一样ATen https://pytorch.org/cppdocs/#aten该函数有多种实现。大多数 ATen 函数都注册在aten/src/ATen/native/native_functions.yaml https://github.com/pytorch/pytorch/blob/b3bb234e16780ea3f3d749d2c8c156c9245eb797/aten/src/ATen/native/native_functions.yaml#L4457，一般这里的函数都会有一个_cpu and _cuda版本。
回到搜索结果我们可以发现CUDA实现其实就是调用函数unique_dim_cuda at aten/src/ATen/native/cuda/Unique.cu:197 https://github.com/pytorch/pytorch/blob/b3bb234e16780ea3f3d749d2c8c156c9245eb797/aten/src/ATen/native/cuda/Unique.cu#L197
196: std::tuple<Tensor, Tensor, Tensor>
197: unique_dim_cuda(const Tensor& self, const int64_t dim, const bool sorted, const bool return_inverse, const bool return_counts) {
198: return AT_DISPATCH_ALL_TYPES_AND2(kBool, kHalf, self.scalar_type(), "unique_dim", [&] {
199: return unique_dim_cuda_template<scalar_t>(self, dim, false, return_inverse, return_counts);
200: });
201: }
CPU 实现正在调用该函数unique_dim_cpu at aten/src/ATen/native/Unique.cpp:271 https://github.com/pytorch/pytorch/blob/7ee0712642492ef221a69d3fdf13b607f406bd78/aten/src/ATen/native/Unique.cpp#L271
270: std::tuple<Tensor, Tensor, Tensor>
271: unique_dim_cpu(const Tensor& self, const int64_t dim, const bool sorted, const bool return_inverse, const bool return_counts) {
272: return AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::BFloat16, at::ScalarType::Bool, self.scalar_type(), "unique_dim", [&] {
273: // The current implementation using `dim` always sorts due to unhashable tensors
274: return _unique_dim_cpu_template<scalar_t>(self, dim, false, return_inverse, return_counts);
275: });
276: }
从这一点开始，您应该能够进一步跟踪函数调用，以准确了解它们在做什么。
经过一系列类似的搜索，您应该会发现torch._unique2实施于aten/src/ATen/native/cuda/Unique.cu:188 https://github.com/pytorch/pytorch/blob/b3bb234e16780ea3f3d749d2c8c156c9245eb797/aten/src/ATen/native/cuda/Unique.cu#L188 and aten/src/ATen/native/Unique.cpp:264 https://github.com/pytorch/pytorch/blob/7ee0712642492ef221a69d3fdf13b607f406bd78/aten/src/ATen/native/Unique.cpp#L264分别针对 CUDA 和 CPU。
