specialize_autogradzero.h

#pragma once

#include <torch/csrc/jit/ir.h>

namespace torch {
namespace jit {

// propagate autograd zero information through a gradient graph and
// remove grad_of blocks if present.
// Note: this is a very limited pass. It only propagates autograd zeros for
// operations generated by the symbolic autodiff code and cleans up
// AutogradAdds when possible. Outputs of other nodes are conservatively
// marked Unknown and not optimized.
TORCH_API void specializeAutogradZero(Graph& g);

} // namespace jit
} // namespace torch