mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
36871622f1
This is follow-up of #164912 to mark unused C++ parameters to improve code readability. Pull Request resolved: https://github.com/pytorch/pytorch/pull/165121 Approved by: https://github.com/Skylion007
73 lines
2.2 KiB
C++
73 lines
2.2 KiB
C++
#pragma once
|
|
|
|
#include <ATen/Tensor.h>
|
|
#include <torch/csrc/Export.h>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace torch::dynamo::autograd {
|
|
class CompiledNodeArgs;
|
|
class SwapSavedVariables;
|
|
struct PackedArgs;
|
|
} // namespace torch::dynamo::autograd
|
|
|
|
// A hook that's called on gradients
|
|
|
|
namespace torch::autograd {
|
|
|
|
using Variable = at::Tensor;
|
|
using variable_list = std::vector<Variable>;
|
|
|
|
struct TORCH_API FunctionPreHook {
|
|
virtual ~FunctionPreHook() = default;
|
|
virtual variable_list operator()(const variable_list& grads) = 0;
|
|
// only implemented for python hooks, registers hook with compiled autograd
|
|
virtual void compiled_args(
|
|
torch::dynamo::autograd::CompiledNodeArgs& args) const {
|
|
TORCH_CHECK_NOT_IMPLEMENTED(
|
|
false,
|
|
std::string("compiled_args nyi, see [Note: Compiled Autograd] ") +
|
|
typeid(*this).name());
|
|
}
|
|
};
|
|
|
|
struct TORCH_API FunctionPostHook {
|
|
virtual ~FunctionPostHook() = default;
|
|
virtual variable_list operator()(
|
|
const variable_list& outputs /* grad_inputs */,
|
|
const variable_list& inputs /* grad_outputs */) = 0;
|
|
// only implemented for python hooks, registers hook with compiled autograd
|
|
virtual void compiled_args(
|
|
torch::dynamo::autograd::CompiledNodeArgs& args) const {
|
|
TORCH_CHECK_NOT_IMPLEMENTED(
|
|
false,
|
|
std::string("compiled_args nyi, see [Note: Compiled Autograd] ") +
|
|
typeid(*this).name());
|
|
}
|
|
};
|
|
|
|
struct TORCH_API PostAccumulateGradHook {
|
|
virtual ~PostAccumulateGradHook() = default;
|
|
virtual void operator()(const Variable& tensor) = 0;
|
|
// only implemented for python hooks on nodes, registers hook with compiled
|
|
// autograd
|
|
virtual void compiled_args(
|
|
torch::dynamo::autograd::CompiledNodeArgs& args) const {
|
|
TORCH_CHECK_NOT_IMPLEMENTED(
|
|
false,
|
|
std::string("compiled_args nyi, see [Note: Compiled Autograd] ") +
|
|
typeid(*this).name());
|
|
}
|
|
|
|
virtual void apply_with_saved(
|
|
Variable& /*unused*/,
|
|
torch::dynamo::autograd::SwapSavedVariables& /*unused*/) {
|
|
TORCH_CHECK_NOT_IMPLEMENTED(
|
|
false,
|
|
std::string("compiled_args nyi, see [Note: Compiled Autograd] ") +
|
|
typeid(*this).name());
|
|
}
|
|
};
|
|
|
|
} // namespace torch::autograd
|