kRank Class — pytorch Architecture

Source Code

aten/src/ATen/native/transformers/hip/aotriton_adapter.h lines 128–146

template<int kRank, bool kRequireZeros>
struct LazyTensorFunctions : public LazyTensorContext {
  static aotriton::TensorView<kRank> acquire(void* cookie) {
    auto ctx = (LazyTensorContext*)cookie;
    if (!ctx->tensor.defined()) {
      auto q = ctx->like_tensor;
      if constexpr (kRequireZeros) {
        ctx->tensor = at::zeros(q.sizes(),
                                q.options().dtype(at::kFloat));
      } else {
        ctx->tensor = at::empty_like(q);
      }
    }
    return mk_aotensor<kRank>(ctx->tensor, ctx->tensor_name);
  }

  static void dispose(void* cookie) {
  }
};