diff --git a/oneflow/core/autograd/gradient_funcs/activation.cpp b/oneflow/core/autograd/gradient_funcs/activation.cpp
index 1cb4643902f..42d8ba4bf2a 100644
--- a/oneflow/core/autograd/gradient_funcs/activation.cpp
+++ b/oneflow/core/autograd/gradient_funcs/activation.cpp
@@ -19,15 +19,15 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct BaseActivationInterpState : public OpExprInterpState {
+struct BaseActivationCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class BaseActivation : public OpExprGradFunction<BaseActivationInterpState> {
+class BaseActivation : public OpExprGradFunction<BaseActivationCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(BaseActivationInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(BaseActivationCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -39,7 +39,7 @@ class BaseActivation : public OpExprGradFunction<BaseActivationInterpState> {
 
 class Silu : public BaseActivation {
  public:
-  Maybe<void> Apply(const BaseActivationInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const BaseActivationCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -53,7 +53,7 @@ class Silu : public BaseActivation {
 
 class Mish : public BaseActivation {
  public:
-  Maybe<void> Apply(const BaseActivationInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const BaseActivationCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -67,7 +67,7 @@ class Mish : public BaseActivation {
 
 class Selu : public BaseActivation {
  public:
-  Maybe<void> Apply(const BaseActivationInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const BaseActivationCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -81,7 +81,7 @@ class Selu : public BaseActivation {
 
 class Softsign : public BaseActivation {
  public:
-  Maybe<void> Apply(const BaseActivationInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const BaseActivationCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -95,7 +95,7 @@ class Softsign : public BaseActivation {
 
 class GeLU : public BaseActivation {
  public:
-  Maybe<void> Apply(const BaseActivationInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const BaseActivationCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -109,7 +109,7 @@ class GeLU : public BaseActivation {
 
 class HardSigmoid : public BaseActivation {
  public:
-  Maybe<void> Apply(const BaseActivationInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const BaseActivationCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -123,7 +123,7 @@ class HardSigmoid : public BaseActivation {
 
 class HardSwish : public BaseActivation {
  public:
-  Maybe<void> Apply(const BaseActivationInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const BaseActivationCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -136,15 +136,15 @@ class HardSwish : public BaseActivation {
 };
 
 // ===== Activation with parms ====
-struct ReLUInterpState : public OpExprInterpState {
+struct ReLUCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class ReLU : public OpExprGradFunction<ReLUInterpState> {
+class ReLU : public OpExprGradFunction<ReLUCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(ReLUInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+  Maybe<void> Capture(ReLUCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
                       const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -153,7 +153,7 @@ class ReLU : public OpExprGradFunction<ReLUInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const ReLUInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ReLUCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -165,12 +165,13 @@ class ReLU : public OpExprGradFunction<ReLUInterpState> {
   }
 };
 
-struct LeakyReluInterpState : public OpExprInterpState {
+// ===== Activation with parms ====
+struct LeakyReluCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   float alpha;
 };
 
-class LeakyRelu : public OpExprGradFunction<LeakyReluInterpState> {
+class LeakyRelu : public OpExprGradFunction<LeakyReluCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -179,7 +180,7 @@ class LeakyRelu : public OpExprGradFunction<LeakyReluInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(LeakyReluInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(LeakyReluCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     ctx->requires_grad = inputs.at(0)->requires_grad();
@@ -191,7 +192,7 @@ class LeakyRelu : public OpExprGradFunction<LeakyReluInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const LeakyReluInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const LeakyReluCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -206,13 +207,13 @@ class LeakyRelu : public OpExprGradFunction<LeakyReluInterpState> {
   AttrMap base_attrs_;
 };
 
-struct HardTanhInterpState : public OpExprInterpState {
+struct HardTanhCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   double min_val;
   double max_val;
 };
 
-class HardTanh : public OpExprGradFunction<HardTanhInterpState> {
+class HardTanh : public OpExprGradFunction<HardTanhCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -221,7 +222,7 @@ class HardTanh : public OpExprGradFunction<HardTanhInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(HardTanhInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(HardTanhCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
     ctx->requires_grad = inputs.at(0)->requires_grad();
@@ -234,7 +235,7 @@ class HardTanh : public OpExprGradFunction<HardTanhInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const HardTanhInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const HardTanhCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -250,12 +251,12 @@ class HardTanh : public OpExprGradFunction<HardTanhInterpState> {
   AttrMap base_attrs_;
 };
 
-struct EluInterpState : public OpExprInterpState {
+struct EluCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   double alpha;
 };
 
-class Elu : public OpExprGradFunction<EluInterpState> {
+class Elu : public OpExprGradFunction<EluCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -264,7 +265,7 @@ class Elu : public OpExprGradFunction<EluInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(EluInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+  Maybe<void> Capture(EluCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
                       const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     ctx->requires_grad = inputs.at(0)->requires_grad();
@@ -276,7 +277,7 @@ class Elu : public OpExprGradFunction<EluInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const EluInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const EluCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -291,16 +292,16 @@ class Elu : public OpExprGradFunction<EluInterpState> {
   AttrMap base_attrs_;
 };
 
-struct PReLUInterpState : public OpExprInterpState {
+struct PReLUCaptureState : public AutoGradCaptureState {
   bool input_requires_grad;
   bool alpha_requires_grad;
 };
 
-class PReLU : public OpExprGradFunction<PReLUInterpState> {
+class PReLU : public OpExprGradFunction<PReLUCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(PReLUInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+  Maybe<void> Capture(PReLUCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
                       const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 2);
     ctx->input_requires_grad = inputs.at(0)->requires_grad();  // input
@@ -311,7 +312,7 @@ class PReLU : public OpExprGradFunction<PReLUInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const PReLUInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const PReLUCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     const auto& dy = out_grads.at(0);
diff --git a/oneflow/core/autograd/gradient_funcs/adaptive_pool.cpp b/oneflow/core/autograd/gradient_funcs/adaptive_pool.cpp
index cc4009a8098..2bef644efd2 100644
--- a/oneflow/core/autograd/gradient_funcs/adaptive_pool.cpp
+++ b/oneflow/core/autograd/gradient_funcs/adaptive_pool.cpp
@@ -23,18 +23,18 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct AdaptivePoolInterpState : public OpExprInterpState {
+struct AdaptivePoolCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class AdaptivePoolNdGrad : public OpExprGradFunction<AdaptivePoolInterpState> {
+class AdaptivePoolNdGrad : public OpExprGradFunction<AdaptivePoolCaptureState> {
  public:
-  using OpExprGradFunction<AdaptivePoolInterpState>::Init;
+  using OpExprGradFunction<AdaptivePoolCaptureState>::Init;
 
   Maybe<void> Init(const OpExpr& op, std::string mode, const int& ndims);
-  Maybe<void> Capture(AdaptivePoolInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(AdaptivePoolCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const AdaptivePoolInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const AdaptivePoolCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -52,7 +52,7 @@ Maybe<void> AdaptivePoolNdGrad::Init(const OpExpr& op, std::string mode, const i
   return Maybe<void>::Ok();
 }
 
-Maybe<void> AdaptivePoolNdGrad::Capture(AdaptivePoolInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> AdaptivePoolNdGrad::Capture(AdaptivePoolCaptureState* ctx, const TensorTuple& inputs,
                                         const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -61,7 +61,7 @@ Maybe<void> AdaptivePoolNdGrad::Capture(AdaptivePoolInterpState* ctx, const Tens
   return Maybe<void>::Ok();
 }
 
-Maybe<void> AdaptivePoolNdGrad::Apply(const AdaptivePoolInterpState* ctx,
+Maybe<void> AdaptivePoolNdGrad::Apply(const AdaptivePoolCaptureState* ctx,
                                       const TensorTuple& out_grads, TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/add_n.cpp b/oneflow/core/autograd/gradient_funcs/add_n.cpp
index 5c083ffd587..2748de7a063 100644
--- a/oneflow/core/autograd/gradient_funcs/add_n.cpp
+++ b/oneflow/core/autograd/gradient_funcs/add_n.cpp
@@ -18,16 +18,16 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct AddNInterpState : public OpExprInterpState {
+struct AddNCaptureState : public AutoGradCaptureState {
   int32_t input_num;
   std::vector<bool> requires_grad;
 };
 
-class AddN : public OpExprGradFunction<AddNInterpState> {
+class AddN : public OpExprGradFunction<AddNCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(AddNInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+  Maybe<void> Capture(AddNCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
                       const AttrMap& attrs) const override {
     ctx->input_num = inputs.size();
     ctx->requires_grad.resize(inputs.size());
@@ -37,7 +37,7 @@ class AddN : public OpExprGradFunction<AddNInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const AddNInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const AddNCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(ctx->input_num);
diff --git a/oneflow/core/autograd/gradient_funcs/avg_pooling.cpp b/oneflow/core/autograd/gradient_funcs/avg_pooling.cpp
index f0bd56a2d8e..526b4458e65 100644
--- a/oneflow/core/autograd/gradient_funcs/avg_pooling.cpp
+++ b/oneflow/core/autograd/gradient_funcs/avg_pooling.cpp
@@ -26,7 +26,7 @@ namespace one {
 
 namespace {
 
-struct AvgPoolingInterpState : public OpExprInterpState {
+struct AvgPoolingCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   size_t input_index;
   size_t output_index;
@@ -40,13 +40,13 @@ struct AvgPoolingInterpState : public OpExprInterpState {
   int64_t divisor_override;
 };
 
-class AvgPoolingNdGrad : public OpExprGradFunction<AvgPoolingInterpState> {
+class AvgPoolingNdGrad : public OpExprGradFunction<AvgPoolingCaptureState> {
  public:
   virtual ~AvgPoolingNdGrad() = default;
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(AvgPoolingInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(AvgPoolingCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const AvgPoolingInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const AvgPoolingCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -60,7 +60,7 @@ Maybe<void> AvgPoolingNdGrad::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> AvgPoolingNdGrad::Capture(AvgPoolingInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> AvgPoolingNdGrad::Capture(AvgPoolingCaptureState* ctx, const TensorTuple& inputs,
                                       const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -80,7 +80,7 @@ Maybe<void> AvgPoolingNdGrad::Capture(AvgPoolingInterpState* ctx, const TensorTu
   return Maybe<void>::Ok();
 }
 
-Maybe<void> AvgPoolingNdGrad::Apply(const AvgPoolingInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> AvgPoolingNdGrad::Apply(const AvgPoolingCaptureState* ctx, const TensorTuple& out_grads,
                                     TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/batch_gather.cpp b/oneflow/core/autograd/gradient_funcs/batch_gather.cpp
index bfba2bd0c6c..b21fc693481 100644
--- a/oneflow/core/autograd/gradient_funcs/batch_gather.cpp
+++ b/oneflow/core/autograd/gradient_funcs/batch_gather.cpp
@@ -22,17 +22,17 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct BatchGatherInterpState : public OpExprInterpState {
+struct BatchGatherCaptureState : public AutoGradCaptureState {
   int64_t num_segments;
   bool requires_grad;
 };
 
-class BatchGather : public OpExprGradFunction<BatchGatherInterpState> {
+class BatchGather : public OpExprGradFunction<BatchGatherCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(BatchGatherInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(BatchGatherCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const BatchGatherInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const BatchGatherCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -48,7 +48,7 @@ Maybe<void> BatchGather::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> BatchGather::Capture(BatchGatherInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> BatchGather::Capture(BatchGatherCaptureState* ctx, const TensorTuple& inputs,
                                  const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -59,7 +59,7 @@ Maybe<void> BatchGather::Capture(BatchGatherInterpState* ctx, const TensorTuple&
   return Maybe<void>::Ok();
 }
 
-Maybe<void> BatchGather::Apply(const BatchGatherInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> BatchGather::Apply(const BatchGatherCaptureState* ctx, const TensorTuple& out_grads,
                                TensorTuple* in_grads) const {
   in_grads->resize(2);
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
diff --git a/oneflow/core/autograd/gradient_funcs/bias_add.cpp b/oneflow/core/autograd/gradient_funcs/bias_add.cpp
index 04378ee553e..6c2a52c1959 100644
--- a/oneflow/core/autograd/gradient_funcs/bias_add.cpp
+++ b/oneflow/core/autograd/gradient_funcs/bias_add.cpp
@@ -23,13 +23,13 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct BiasAddInterpState : public OpExprInterpState {
+struct BiasAddCaptureState : public AutoGradCaptureState {
   bool input_requires_grad;
   bool bias_requires_grad;
   int32_t axis;
 };
 
-class BiasAdd : public OpExprGradFunction<BiasAddInterpState> {
+class BiasAdd : public OpExprGradFunction<BiasAddCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -42,7 +42,7 @@ class BiasAdd : public OpExprGradFunction<BiasAddInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(BiasAddInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(BiasAddCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 2);
     ctx->input_requires_grad = inputs.at(0)->requires_grad();
@@ -52,7 +52,7 @@ class BiasAdd : public OpExprGradFunction<BiasAddInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const BiasAddInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const BiasAddCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const int64_t num_axes = out_grads.at(0)->shape()->NumAxes();
     in_grads->resize(2);
diff --git a/oneflow/core/autograd/gradient_funcs/broadcast_binary_ops.cpp b/oneflow/core/autograd/gradient_funcs/broadcast_binary_ops.cpp
index b5c0ef859f9..d3d7eb28588 100644
--- a/oneflow/core/autograd/gradient_funcs/broadcast_binary_ops.cpp
+++ b/oneflow/core/autograd/gradient_funcs/broadcast_binary_ops.cpp
@@ -21,15 +21,15 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-class BroadcastBinaryGrad : public OpExprGradFunction<OpExprInterpState> {
+class BroadcastBinaryGrad : public OpExprGradFunction<AutoGradCaptureState> {
  public:
   BroadcastBinaryGrad() = default;
   virtual ~BroadcastBinaryGrad() = default;
 
   virtual Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(OpExprInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
-                      const AttrMap& attrs) const override {
+  Maybe<void> Capture(AutoGradCaptureState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 2);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
     ctx->SaveTensorForBackward(inputs.at(0));
@@ -41,7 +41,7 @@ class BroadcastBinaryGrad : public OpExprGradFunction<OpExprInterpState> {
 
 class BroadcastAdd : public BroadcastBinaryGrad {
  public:
-  Maybe<void> Apply(const OpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const AutoGradCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const auto& x = ctx->SavedTensors().at(0);
     const auto& y = ctx->SavedTensors().at(1);
@@ -60,7 +60,7 @@ REGISTER_OP_EXPR_GRAD_FUNCTION("broadcast_add", BroadcastAdd);
 
 class BroadcastSub : public BroadcastBinaryGrad {
  public:
-  Maybe<void> Apply(const OpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const AutoGradCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const auto& x = ctx->SavedTensors().at(0);
     const auto& y = ctx->SavedTensors().at(1);
@@ -80,7 +80,7 @@ REGISTER_OP_EXPR_GRAD_FUNCTION("broadcast_sub", BroadcastSub);
 
 class BroadcastMul : public BroadcastBinaryGrad {
  public:
-  Maybe<void> Apply(const OpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const AutoGradCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const auto& x = ctx->SavedTensors().at(0);
     const auto& y = ctx->SavedTensors().at(1);
@@ -101,7 +101,7 @@ REGISTER_OP_EXPR_GRAD_FUNCTION("broadcast_mul", BroadcastMul);
 
 class BroadcastDiv : public BroadcastBinaryGrad {
  public:
-  Maybe<void> Apply(const OpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const AutoGradCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const auto& x = ctx->SavedTensors().at(0);
     const auto& y = ctx->SavedTensors().at(1);
@@ -122,7 +122,7 @@ REGISTER_OP_EXPR_GRAD_FUNCTION("broadcast_div", BroadcastDiv);
 
 class BroadcastMinMax : public BroadcastBinaryGrad {
  public:
-  Maybe<void> Apply(const OpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const AutoGradCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const auto& x = ctx->SavedTensors().at(0);
     const auto& y = ctx->SavedTensors().at(1);
diff --git a/oneflow/core/autograd/gradient_funcs/broadcast_floor_mod.cpp b/oneflow/core/autograd/gradient_funcs/broadcast_floor_mod.cpp
index 47c9217b6b6..b1a58cc3403 100644
--- a/oneflow/core/autograd/gradient_funcs/broadcast_floor_mod.cpp
+++ b/oneflow/core/autograd/gradient_funcs/broadcast_floor_mod.cpp
@@ -19,22 +19,22 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct BroadcastFModInterpState : public OpExprInterpState {
+struct BroadcastFModCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class BroadcastFMod : public OpExprGradFunction<BroadcastFModInterpState> {
+class BroadcastFMod : public OpExprGradFunction<BroadcastFModCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(BroadcastFModInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(BroadcastFModCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 2);
     ctx->requires_grad = inputs.at(0)->requires_grad();
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const BroadcastFModInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const BroadcastFModCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(2);
diff --git a/oneflow/core/autograd/gradient_funcs/broadcast_like.cpp b/oneflow/core/autograd/gradient_funcs/broadcast_like.cpp
index b1571e5413f..7968fae8ad7 100644
--- a/oneflow/core/autograd/gradient_funcs/broadcast_like.cpp
+++ b/oneflow/core/autograd/gradient_funcs/broadcast_like.cpp
@@ -22,17 +22,17 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct BroadCastLikeInterpState : public OpExprInterpState {
+struct BroadCastLikeCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   std::vector<int32_t> broadcast_axes;
 };
 
-class BroadCastLike : public OpExprGradFunction<BroadCastLikeInterpState> {
+class BroadCastLike : public OpExprGradFunction<BroadCastLikeCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(BroadCastLikeInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(BroadCastLikeCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const BroadCastLikeInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const BroadCastLikeCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -50,7 +50,7 @@ Maybe<void> BroadCastLike::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> BroadCastLike::Capture(BroadCastLikeInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> BroadCastLike::Capture(BroadCastLikeCaptureState* ctx, const TensorTuple& inputs,
                                    const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -61,7 +61,7 @@ Maybe<void> BroadCastLike::Capture(BroadCastLikeInterpState* ctx, const TensorTu
   return Maybe<void>::Ok();
 }
 
-Maybe<void> BroadCastLike::Apply(const BroadCastLikeInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> BroadCastLike::Apply(const BroadCastLikeCaptureState* ctx, const TensorTuple& out_grads,
                                  TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/cast.cpp b/oneflow/core/autograd/gradient_funcs/cast.cpp
index b4f60b19028..7a8934b0133 100644
--- a/oneflow/core/autograd/gradient_funcs/cast.cpp
+++ b/oneflow/core/autograd/gradient_funcs/cast.cpp
@@ -24,11 +24,11 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct CastOpExprInterpState : public OpExprInterpState {
+struct CastCaptureState : public AutoGradCaptureState {
   DataType data_type;
 };
 
-class Cast : public OpExprGradFunction<CastOpExprInterpState> {
+class Cast : public OpExprGradFunction<CastCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -38,13 +38,13 @@ class Cast : public OpExprGradFunction<CastOpExprInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(CastOpExprInterpState* ctx, const TensorTuple& inputs,
-                      const TensorTuple& outputs, const AttrMap& attrs) const override {
+  Maybe<void> Capture(CastCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+                      const AttrMap& attrs) const override {
     ctx->data_type = inputs.at(0)->dtype()->data_type();
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const CastOpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const CastCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     in_grads->resize(1);
     MutableAttrMap attrs;
diff --git a/oneflow/core/autograd/gradient_funcs/clip_by_scalar.cpp b/oneflow/core/autograd/gradient_funcs/clip_by_scalar.cpp
index 7fa5963c770..f87dbac7ddd 100644
--- a/oneflow/core/autograd/gradient_funcs/clip_by_scalar.cpp
+++ b/oneflow/core/autograd/gradient_funcs/clip_by_scalar.cpp
@@ -19,13 +19,13 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ClipByScalarInterpState : public OpExprInterpState {
+struct ClipByScalarCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   functional::Scalar min;
   functional::Scalar max;
 };
 
-class ClipByScalar : public OpExprGradFunction<ClipByScalarInterpState> {
+class ClipByScalar : public OpExprGradFunction<ClipByScalarCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -34,7 +34,7 @@ class ClipByScalar : public OpExprGradFunction<ClipByScalarInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(ClipByScalarInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ClipByScalarCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     ctx->requires_grad = inputs.at(0)->requires_grad();
@@ -54,7 +54,7 @@ class ClipByScalar : public OpExprGradFunction<ClipByScalarInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const ClipByScalarInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ClipByScalarCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/clip_by_scalar_max.cpp b/oneflow/core/autograd/gradient_funcs/clip_by_scalar_max.cpp
index 587740aef97..291d56d81fd 100644
--- a/oneflow/core/autograd/gradient_funcs/clip_by_scalar_max.cpp
+++ b/oneflow/core/autograd/gradient_funcs/clip_by_scalar_max.cpp
@@ -19,12 +19,12 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ClipByScalarMaxInterpState : public OpExprInterpState {
+struct ClipByScalarMaxCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   functional::Scalar max;
 };
 
-class ClipByScalarMax : public OpExprGradFunction<ClipByScalarMaxInterpState> {
+class ClipByScalarMax : public OpExprGradFunction<ClipByScalarMaxCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -33,7 +33,7 @@ class ClipByScalarMax : public OpExprGradFunction<ClipByScalarMaxInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(ClipByScalarMaxInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ClipByScalarMaxCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     ctx->requires_grad = inputs.at(0)->requires_grad();
@@ -51,7 +51,7 @@ class ClipByScalarMax : public OpExprGradFunction<ClipByScalarMaxInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const ClipByScalarMaxInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ClipByScalarMaxCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/clip_by_scalar_min.cpp b/oneflow/core/autograd/gradient_funcs/clip_by_scalar_min.cpp
index f73470dfccd..35ebc620b51 100644
--- a/oneflow/core/autograd/gradient_funcs/clip_by_scalar_min.cpp
+++ b/oneflow/core/autograd/gradient_funcs/clip_by_scalar_min.cpp
@@ -19,12 +19,12 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ClipByScalarMinInterpState : public OpExprInterpState {
+struct ClipByScalarMinCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   functional::Scalar min;
 };
 
-class ClipByScalarMin : public OpExprGradFunction<ClipByScalarMinInterpState> {
+class ClipByScalarMin : public OpExprGradFunction<ClipByScalarMinCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -33,7 +33,7 @@ class ClipByScalarMin : public OpExprGradFunction<ClipByScalarMinInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(ClipByScalarMinInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ClipByScalarMinCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     ctx->requires_grad = inputs.at(0)->requires_grad();
@@ -51,7 +51,7 @@ class ClipByScalarMin : public OpExprGradFunction<ClipByScalarMinInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const ClipByScalarMinInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ClipByScalarMinCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/combined_margin_loss.cpp b/oneflow/core/autograd/gradient_funcs/combined_margin_loss.cpp
index 3f22cc5b7ca..9aea3f43512 100644
--- a/oneflow/core/autograd/gradient_funcs/combined_margin_loss.cpp
+++ b/oneflow/core/autograd/gradient_funcs/combined_margin_loss.cpp
@@ -21,7 +21,7 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct CombinedMarginLossInterpState : public OpExprInterpState {
+struct CombinedMarginLossCaptureState : public AutoGradCaptureState {
   float m1;
   float m2;
   float m3;
@@ -31,7 +31,7 @@ struct CombinedMarginLossInterpState : public OpExprInterpState {
   bool requires_grad;
 };
 
-class CombinedMarginLoss : public OpExprGradFunction<CombinedMarginLossInterpState> {
+class CombinedMarginLoss : public OpExprGradFunction<CombinedMarginLossCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -40,7 +40,7 @@ class CombinedMarginLoss : public OpExprGradFunction<CombinedMarginLossInterpSta
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(CombinedMarginLossInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(CombinedMarginLossCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 2);
     ctx->requires_grad = inputs.at(0)->requires_grad();  // x
@@ -57,7 +57,7 @@ class CombinedMarginLoss : public OpExprGradFunction<CombinedMarginLossInterpSta
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const CombinedMarginLossInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const CombinedMarginLossCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 2);
     in_grads->resize(2);
diff --git a/oneflow/core/autograd/gradient_funcs/concat.cpp b/oneflow/core/autograd/gradient_funcs/concat.cpp
index 9bbb13aa50d..3d71b118f6f 100644
--- a/oneflow/core/autograd/gradient_funcs/concat.cpp
+++ b/oneflow/core/autograd/gradient_funcs/concat.cpp
@@ -23,18 +23,18 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ConcatInterpState : public OpExprInterpState {
+struct ConcatCaptureState : public AutoGradCaptureState {
   std::vector<bool> requires_grad;
   int64_t axis;
   int64_t input_num;
 };
 
-class Concat : public OpExprGradFunction<ConcatInterpState> {
+class Concat : public OpExprGradFunction<ConcatCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(ConcatInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
-                      const AttrMap& attrs) const override;
-  Maybe<void> Apply(const ConcatInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Capture(ConcatCaptureState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override;
+  Maybe<void> Apply(const ConcatCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -55,7 +55,7 @@ Maybe<void> Concat::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Concat::Capture(ConcatInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Concat::Capture(ConcatCaptureState* ctx, const TensorTuple& inputs,
                             const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad.resize(inputs.size());
   for (int i = 0; i < inputs.size(); ++i) { ctx->requires_grad[i] = inputs.at(i)->requires_grad(); }
@@ -67,7 +67,7 @@ Maybe<void> Concat::Capture(ConcatInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Concat::Apply(const ConcatInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Concat::Apply(const ConcatCaptureState* ctx, const TensorTuple& out_grads,
                           TensorTuple* in_grads) const {
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
   in_grads->resize(ctx->input_num);
diff --git a/oneflow/core/autograd/gradient_funcs/consistent_cast.cpp b/oneflow/core/autograd/gradient_funcs/consistent_cast.cpp
index 284b8f824a8..87f7eacdd5b 100644
--- a/oneflow/core/autograd/gradient_funcs/consistent_cast.cpp
+++ b/oneflow/core/autograd/gradient_funcs/consistent_cast.cpp
@@ -22,13 +22,13 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct CastConsistentOpExprInterpState : public OpExprInterpState {
+struct CastConsistentCaptureState : public AutoGradCaptureState {
   Symbol<ParallelDesc> parallel_desc;
   Symbol<cfg::NdSbp> nd_sbp;
   std::shared_ptr<const Shape> shape;
 };
 
-class CastToConsistent : public OpExprGradFunction<CastConsistentOpExprInterpState> {
+class CastToConsistent : public OpExprGradFunction<CastConsistentCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const CastToConsistentOpExpr*>(&op);
@@ -38,7 +38,7 @@ class CastToConsistent : public OpExprGradFunction<CastConsistentOpExprInterpSta
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(CastConsistentOpExprInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(CastConsistentCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs,
                       const OpExprInterpContext& interp_ctx) const override {
     ctx->parallel_desc = JUST(interp_ctx.parallel_desc.value());
@@ -46,7 +46,7 @@ class CastToConsistent : public OpExprGradFunction<CastConsistentOpExprInterpSta
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const CastConsistentOpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const CastConsistentCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const auto& out_grad = out_grads.at(0);
     CHECK_OR_RETURN(out_grad->is_consistent());
@@ -63,7 +63,7 @@ class CastToConsistent : public OpExprGradFunction<CastConsistentOpExprInterpSta
 
 REGISTER_OP_EXPR_GRAD_FUNCTION("cast_to_consistent", CastToConsistent);
 
-class CastFromConsistent : public OpExprGradFunction<CastConsistentOpExprInterpState> {
+class CastFromConsistent : public OpExprGradFunction<CastConsistentCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const CastFromConsistentOpExpr*>(&op);
@@ -73,7 +73,7 @@ class CastFromConsistent : public OpExprGradFunction<CastConsistentOpExprInterpS
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(CastConsistentOpExprInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(CastConsistentCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     const auto& input = inputs.at(0);
     CHECK_OR_RETURN(input->is_consistent());
@@ -83,7 +83,7 @@ class CastFromConsistent : public OpExprGradFunction<CastConsistentOpExprInterpS
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const CastConsistentOpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const CastConsistentCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const auto& dual_nd_sbp = JUST(GetDualNdSbp(ctx->nd_sbp));
     MutableAttrMap attrs;
diff --git a/oneflow/core/autograd/gradient_funcs/conv.cpp b/oneflow/core/autograd/gradient_funcs/conv.cpp
index 8080a931618..cbe84ba5f15 100644
--- a/oneflow/core/autograd/gradient_funcs/conv.cpp
+++ b/oneflow/core/autograd/gradient_funcs/conv.cpp
@@ -23,7 +23,7 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ConvolutionNdInterpState : public OpExprInterpState {
+struct ConvolutionNdCaptureState : public AutoGradCaptureState {
   bool input_requires_grad = false;
   bool weight_requires_grad = false;
   size_t input_index;
@@ -37,12 +37,12 @@ struct ConvolutionNdInterpState : public OpExprInterpState {
   int32_t groups;
 };
 
-class ConvolutionNd : public OpExprGradFunction<ConvolutionNdInterpState> {
+class ConvolutionNd : public OpExprGradFunction<ConvolutionNdCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(ConvolutionNdInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ConvolutionNdCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const ConvolutionNdInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ConvolutionNdCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -56,7 +56,7 @@ Maybe<void> ConvolutionNd::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> ConvolutionNd::Capture(ConvolutionNdInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> ConvolutionNd::Capture(ConvolutionNdCaptureState* ctx, const TensorTuple& inputs,
                                    const TensorTuple& outputs, const AttrMap& attrs) const {
   CHECK_EQ_OR_RETURN(inputs.size(), 2);
   ctx->input_requires_grad = inputs.at(0)->requires_grad();
@@ -77,7 +77,7 @@ Maybe<void> ConvolutionNd::Capture(ConvolutionNdInterpState* ctx, const TensorTu
   return Maybe<void>::Ok();
 }
 
-Maybe<void> ConvolutionNd::Apply(const ConvolutionNdInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> ConvolutionNd::Apply(const ConvolutionNdCaptureState* ctx, const TensorTuple& out_grads,
                                  TensorTuple* in_grads) const {
   in_grads->resize(2);
   size_t num_spatial_dims = ctx->kernel_size.size();
diff --git a/oneflow/core/autograd/gradient_funcs/copy.cpp b/oneflow/core/autograd/gradient_funcs/copy.cpp
index d84665db3ae..4798ec05e52 100644
--- a/oneflow/core/autograd/gradient_funcs/copy.cpp
+++ b/oneflow/core/autograd/gradient_funcs/copy.cpp
@@ -23,12 +23,12 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct CopyOpExprInterpState : public OpExprInterpState {
+struct CopyCaptureState : public AutoGradCaptureState {
   std::string device_type;
   int64_t device_id;
 };
 
-class Copy : public OpExprGradFunction<CopyOpExprInterpState> {
+class Copy : public OpExprGradFunction<CopyCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -38,14 +38,14 @@ class Copy : public OpExprGradFunction<CopyOpExprInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(CopyOpExprInterpState* ctx, const TensorTuple& inputs,
-                      const TensorTuple& outputs, const AttrMap& attrs) const override {
+  Maybe<void> Capture(CopyCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+                      const AttrMap& attrs) const override {
     ctx->device_type = JUST(inputs.at(0)->device())->type();
     ctx->device_id = JUST(inputs.at(0)->device())->device_id();
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const CopyOpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const CopyCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     in_grads->resize(1);
     MutableAttrMap attrs;
diff --git a/oneflow/core/autograd/gradient_funcs/ctc_loss.cpp b/oneflow/core/autograd/gradient_funcs/ctc_loss.cpp
index 252a74946e0..76166b85dda 100644
--- a/oneflow/core/autograd/gradient_funcs/ctc_loss.cpp
+++ b/oneflow/core/autograd/gradient_funcs/ctc_loss.cpp
@@ -23,18 +23,18 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct CTCLossInterpState : public OpExprInterpState {
+struct CTCLossCaptureState : public AutoGradCaptureState {
   int32_t blank;
   bool zero_infinity;
   bool requires_grad;
 };
 
-class CTCLoss : public OpExprGradFunction<CTCLossInterpState> {
+class CTCLoss : public OpExprGradFunction<CTCLossCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(CTCLossInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(CTCLossCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const CTCLossInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const CTCLossCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -51,7 +51,7 @@ Maybe<void> CTCLoss::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> CTCLoss::Capture(CTCLossInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> CTCLoss::Capture(CTCLossCaptureState* ctx, const TensorTuple& inputs,
                              const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -71,7 +71,7 @@ Maybe<void> CTCLoss::Capture(CTCLossInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> CTCLoss::Apply(const CTCLossInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> CTCLoss::Apply(const CTCLossCaptureState* ctx, const TensorTuple& out_grads,
                            TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 2);
diff --git a/oneflow/core/autograd/gradient_funcs/deconv.cpp b/oneflow/core/autograd/gradient_funcs/deconv.cpp
index fa2e92551f4..c79ad41db51 100644
--- a/oneflow/core/autograd/gradient_funcs/deconv.cpp
+++ b/oneflow/core/autograd/gradient_funcs/deconv.cpp
@@ -23,17 +23,17 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct DeConvolutionNdInterpState : public OpExprInterpState {
+struct DeConvolutionNdCaptureState : public AutoGradCaptureState {
   bool weight_requires_grad = false;
   bool activation_requires_grad = false;
 };
 
-class DeConvolutionNd : public OpExprGradFunction<DeConvolutionNdInterpState> {
+class DeConvolutionNd : public OpExprGradFunction<DeConvolutionNdCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(DeConvolutionNdInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(DeConvolutionNdCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const DeConvolutionNdInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const DeConvolutionNdCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -71,7 +71,7 @@ Maybe<void> DeConvolutionNd::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> DeConvolutionNd::Capture(DeConvolutionNdInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> DeConvolutionNd::Capture(DeConvolutionNdCaptureState* ctx, const TensorTuple& inputs,
                                      const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->activation_requires_grad = inputs.at(0)->requires_grad();
   ctx->weight_requires_grad = inputs.at(1)->requires_grad();
@@ -84,7 +84,7 @@ Maybe<void> DeConvolutionNd::Capture(DeConvolutionNdInterpState* ctx, const Tens
   return Maybe<void>::Ok();
 }
 
-Maybe<void> DeConvolutionNd::Apply(const DeConvolutionNdInterpState* ctx,
+Maybe<void> DeConvolutionNd::Apply(const DeConvolutionNdCaptureState* ctx,
                                    const TensorTuple& out_grads, TensorTuple* in_grads) const {
   in_grads->resize(2);
   if (ctx->activation_requires_grad) {
diff --git a/oneflow/core/autograd/gradient_funcs/diag.cpp b/oneflow/core/autograd/gradient_funcs/diag.cpp
index cfd0aee9daf..ccee46a3c5b 100644
--- a/oneflow/core/autograd/gradient_funcs/diag.cpp
+++ b/oneflow/core/autograd/gradient_funcs/diag.cpp
@@ -20,17 +20,17 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct DiagInterpState : public OpExprInterpState {
+struct DiagCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   int32_t diagonal;
 };
 
-class Diag : public OpExprGradFunction<DiagInterpState> {
+class Diag : public OpExprGradFunction<DiagCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(DiagInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+  Maybe<void> Capture(DiagCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
                       const AttrMap& attrs) const override;
-  Maybe<void> Apply(const DiagInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const DiagCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -44,7 +44,7 @@ Maybe<void> Diag::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Diag::Capture(DiagInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Diag::Capture(DiagCaptureState* ctx, const TensorTuple& inputs,
                           const TensorTuple& outputs, const AttrMap& attrs) const {
   CHECK_EQ_OR_RETURN(outputs.size(), 1);
   ctx->requires_grad = inputs.at(0)->requires_grad();
@@ -55,7 +55,7 @@ Maybe<void> Diag::Capture(DiagInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Diag::Apply(const DiagInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Diag::Apply(const DiagCaptureState* ctx, const TensorTuple& out_grads,
                         TensorTuple* in_grads) const {
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
   in_grads->resize(2);
diff --git a/oneflow/core/autograd/gradient_funcs/dim_gather.cpp b/oneflow/core/autograd/gradient_funcs/dim_gather.cpp
index 4ae5b63e960..b6a39930b41 100644
--- a/oneflow/core/autograd/gradient_funcs/dim_gather.cpp
+++ b/oneflow/core/autograd/gradient_funcs/dim_gather.cpp
@@ -22,17 +22,17 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct DimGatherInterpState : public OpExprInterpState {
+struct DimGatherCaptureState : public AutoGradCaptureState {
   int32_t dim;
   bool requires_grad;
 };
 
-class DimGather : public OpExprGradFunction<DimGatherInterpState> {
+class DimGather : public OpExprGradFunction<DimGatherCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(DimGatherInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(DimGatherCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const DimGatherInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const DimGatherCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -49,7 +49,7 @@ Maybe<void> DimGather::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> DimGather::Capture(DimGatherInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> DimGather::Capture(DimGatherCaptureState* ctx, const TensorTuple& inputs,
                                const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -62,7 +62,7 @@ Maybe<void> DimGather::Capture(DimGatherInterpState* ctx, const TensorTuple& inp
   return Maybe<void>::Ok();
 }
 
-Maybe<void> DimGather::Apply(const DimGatherInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> DimGather::Apply(const DimGatherCaptureState* ctx, const TensorTuple& out_grads,
                              TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/dim_scatter.cpp b/oneflow/core/autograd/gradient_funcs/dim_scatter.cpp
index 6bda00e3abc..f443de1c3d1 100644
--- a/oneflow/core/autograd/gradient_funcs/dim_scatter.cpp
+++ b/oneflow/core/autograd/gradient_funcs/dim_scatter.cpp
@@ -23,7 +23,7 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct DimScatterInterpState : public OpExprInterpState {
+struct DimScatterCaptureState : public AutoGradCaptureState {
   int32_t dim;
   bool input_requires_grad;
   bool src_requires_grad;
@@ -32,14 +32,14 @@ struct DimScatterInterpState : public OpExprInterpState {
 enum SCATTER_TYPE { SCATTER_UPDATE, SCATTER_ADD };
 
 template<SCATTER_TYPE T>
-class DimScatter : public OpExprGradFunction<DimScatterInterpState> {
+class DimScatter : public OpExprGradFunction<DimScatterCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(DimScatterInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(DimScatterCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const DimScatterInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const DimScatterCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
-  Maybe<void> ApplyCommon(const DimScatterInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> ApplyCommon(const DimScatterCaptureState* ctx, const TensorTuple& out_grads,
                           TensorTuple* in_grads) const;
 
  private:
@@ -55,7 +55,7 @@ Maybe<void> DimScatter<T>::Init(const OpExpr& op) {
 }
 
 template<SCATTER_TYPE T>
-Maybe<void> DimScatter<T>::Capture(DimScatterInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> DimScatter<T>::Capture(DimScatterCaptureState* ctx, const TensorTuple& inputs,
                                    const TensorTuple& outputs, const AttrMap& attrs) const {
   CHECK_EQ_OR_RETURN(inputs.size(), 3);
   CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -72,7 +72,7 @@ Maybe<void> DimScatter<T>::Capture(DimScatterInterpState* ctx, const TensorTuple
 }
 
 template<SCATTER_TYPE T>
-Maybe<void> DimScatter<T>::ApplyCommon(const DimScatterInterpState* ctx,
+Maybe<void> DimScatter<T>::ApplyCommon(const DimScatterCaptureState* ctx,
                                        const TensorTuple& out_grads, TensorTuple* in_grads) const {
   const std::shared_ptr<oneflow::one::Tensor>& index = ctx->SavedTensors().at(0);
 
@@ -85,7 +85,7 @@ Maybe<void> DimScatter<T>::ApplyCommon(const DimScatterInterpState* ctx,
 }
 
 template<>
-Maybe<void> DimScatter<SCATTER_TYPE::SCATTER_UPDATE>::Apply(const DimScatterInterpState* ctx,
+Maybe<void> DimScatter<SCATTER_TYPE::SCATTER_UPDATE>::Apply(const DimScatterCaptureState* ctx,
                                                             const TensorTuple& out_grads,
                                                             TensorTuple* in_grads) const {
   if ((!ctx->input_requires_grad) && (!ctx->src_requires_grad)) { return Maybe<void>::Ok(); }
@@ -101,7 +101,7 @@ Maybe<void> DimScatter<SCATTER_TYPE::SCATTER_UPDATE>::Apply(const DimScatterInte
 }
 
 template<>
-Maybe<void> DimScatter<SCATTER_TYPE::SCATTER_ADD>::Apply(const DimScatterInterpState* ctx,
+Maybe<void> DimScatter<SCATTER_TYPE::SCATTER_ADD>::Apply(const DimScatterCaptureState* ctx,
                                                          const TensorTuple& out_grads,
                                                          TensorTuple* in_grads) const {
   if ((!ctx->input_requires_grad) && (!ctx->src_requires_grad)) { return Maybe<void>::Ok(); }
@@ -114,12 +114,12 @@ Maybe<void> DimScatter<SCATTER_TYPE::SCATTER_ADD>::Apply(const DimScatterInterpS
   return Maybe<void>::Ok();
 }
 
-class DimScatterUpdateScalar : public OpExprGradFunction<DimScatterInterpState> {
+class DimScatterUpdateScalar : public OpExprGradFunction<DimScatterCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(DimScatterInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(DimScatterCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const DimScatterInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const DimScatterCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -134,7 +134,7 @@ Maybe<void> DimScatterUpdateScalar::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> DimScatterUpdateScalar::Capture(DimScatterInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> DimScatterUpdateScalar::Capture(DimScatterCaptureState* ctx, const TensorTuple& inputs,
                                             const TensorTuple& outputs,
                                             const AttrMap& attrs) const {
   CHECK_EQ_OR_RETURN(inputs.size(), 2);
@@ -150,7 +150,7 @@ Maybe<void> DimScatterUpdateScalar::Capture(DimScatterInterpState* ctx, const Te
   return Maybe<void>::Ok();
 }
 
-Maybe<void> DimScatterUpdateScalar::Apply(const DimScatterInterpState* ctx,
+Maybe<void> DimScatterUpdateScalar::Apply(const DimScatterCaptureState* ctx,
                                           const TensorTuple& out_grads,
                                           TensorTuple* in_grads) const {
   if (!ctx->input_requires_grad) { return Maybe<void>::Ok(); }
diff --git a/oneflow/core/autograd/gradient_funcs/dropout.cpp b/oneflow/core/autograd/gradient_funcs/dropout.cpp
index c1f6e76914e..d79cf0d8aef 100644
--- a/oneflow/core/autograd/gradient_funcs/dropout.cpp
+++ b/oneflow/core/autograd/gradient_funcs/dropout.cpp
@@ -22,17 +22,17 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct DropoutInterpState : public OpExprInterpState {
+struct DropoutCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   float scale;
 };
 
-class Dropout : public OpExprGradFunction<DropoutInterpState> {
+class Dropout : public OpExprGradFunction<DropoutCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(DropoutInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(DropoutCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const DropoutInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const DropoutCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -49,7 +49,7 @@ Maybe<void> Dropout::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Dropout::Capture(DropoutInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Dropout::Capture(DropoutCaptureState* ctx, const TensorTuple& inputs,
                              const TensorTuple& outputs, const AttrMap& attrs) const {
   ComposedAttrMap composed_attrs(attrs, base_attrs_);
   ctx->requires_grad = inputs.at(0)->requires_grad();
@@ -62,7 +62,7 @@ Maybe<void> Dropout::Capture(DropoutInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Dropout::Apply(const DropoutInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Dropout::Apply(const DropoutCaptureState* ctx, const TensorTuple& out_grads,
                            TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/eager_nccl_broadcast.cpp b/oneflow/core/autograd/gradient_funcs/eager_nccl_broadcast.cpp
index bb6587f70fc..49bde68c7b8 100644
--- a/oneflow/core/autograd/gradient_funcs/eager_nccl_broadcast.cpp
+++ b/oneflow/core/autograd/gradient_funcs/eager_nccl_broadcast.cpp
@@ -52,12 +52,12 @@ Maybe<one::UserOpExpr> FindOrCreatEagerNcclReduceOpExpr(Symbol<ParallelDesc> par
 
 }  // namespace
 
-struct EagerNcclBroadcastOpExprInterpState : public OpExprInterpState {
+struct EagerNcclBroadcastCaptureState : public AutoGradCaptureState {
   Symbol<ParallelDesc> parallel_desc;
   int64_t root;
 };
 
-class EagerNcclBroadcast : public OpExprGradFunction<EagerNcclBroadcastOpExprInterpState> {
+class EagerNcclBroadcast : public OpExprGradFunction<EagerNcclBroadcastCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -65,7 +65,7 @@ class EagerNcclBroadcast : public OpExprGradFunction<EagerNcclBroadcastOpExprInt
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(EagerNcclBroadcastOpExprInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(EagerNcclBroadcastCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs,
                       const OpExprInterpContext& interp_ctx) const override {
     ctx->root = JUST(interp_ctx.attrs.GetAttr<int64_t>("root"));
@@ -73,7 +73,7 @@ class EagerNcclBroadcast : public OpExprGradFunction<EagerNcclBroadcastOpExprInt
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const EagerNcclBroadcastOpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const EagerNcclBroadcastCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const auto& grad_op = JUST(FindOrCreatEagerNcclReduceOpExpr(ctx->parallel_desc, ctx->root));
     in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/elementwise_minimum_maximum.cpp b/oneflow/core/autograd/gradient_funcs/elementwise_minimum_maximum.cpp
index e155c6310b9..5d95292ed85 100644
--- a/oneflow/core/autograd/gradient_funcs/elementwise_minimum_maximum.cpp
+++ b/oneflow/core/autograd/gradient_funcs/elementwise_minimum_maximum.cpp
@@ -22,14 +22,14 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ElementwiseXimumOpExprInterpState : public OpExprInterpState {
+struct ElementwiseXimumCaptureState : public AutoGradCaptureState {
   bool x_requires_grad;
   bool y_requires_grad;
 };
 
-class ElementwiseXimumOp : public OpExprGradFunction<ElementwiseXimumOpExprInterpState> {
+class ElementwiseXimumOp : public OpExprGradFunction<ElementwiseXimumCaptureState> {
  public:
-  Maybe<void> Capture(ElementwiseXimumOpExprInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ElementwiseXimumCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     ctx->x_requires_grad = inputs.at(0)->requires_grad();
     ctx->y_requires_grad = inputs.at(1)->requires_grad();
@@ -38,7 +38,7 @@ class ElementwiseXimumOp : public OpExprGradFunction<ElementwiseXimumOpExprInter
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const ElementwiseXimumOpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ElementwiseXimumCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (!(ctx->x_requires_grad || ctx->y_requires_grad)) { return Maybe<void>::Ok(); }
 
diff --git a/oneflow/core/autograd/gradient_funcs/expand.cpp b/oneflow/core/autograd/gradient_funcs/expand.cpp
index f735c1d0c20..b5e8fa7f25b 100644
--- a/oneflow/core/autograd/gradient_funcs/expand.cpp
+++ b/oneflow/core/autograd/gradient_funcs/expand.cpp
@@ -22,18 +22,18 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ExpandInterpState : public OpExprInterpState {
+struct ExpandCaptureState : public AutoGradCaptureState {
   std::vector<int32_t> out_shape;
   std::vector<int32_t> stride;
   bool requires_grad;
 };
 
-class Expand : public OpExprGradFunction<ExpandInterpState> {
+class Expand : public OpExprGradFunction<ExpandCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(ExpandInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
-                      const AttrMap& attrs) const override;
-  Maybe<void> Apply(const ExpandInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Capture(ExpandCaptureState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override;
+  Maybe<void> Apply(const ExpandCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -52,7 +52,7 @@ Maybe<void> Expand::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Expand::Capture(ExpandInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Expand::Capture(ExpandCaptureState* ctx, const TensorTuple& inputs,
                             const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -63,7 +63,7 @@ Maybe<void> Expand::Capture(ExpandInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Expand::Apply(const ExpandInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Expand::Apply(const ExpandCaptureState* ctx, const TensorTuple& out_grads,
                           TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/fake_quantization.cpp b/oneflow/core/autograd/gradient_funcs/fake_quantization.cpp
index 064b4875e23..006d9c7c24c 100644
--- a/oneflow/core/autograd/gradient_funcs/fake_quantization.cpp
+++ b/oneflow/core/autograd/gradient_funcs/fake_quantization.cpp
@@ -18,22 +18,22 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct FakeQuantizationInterpState : public OpExprInterpState {
+struct FakeQuantizationCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class FakeQuantization : public OpExprGradFunction<FakeQuantizationInterpState> {
+class FakeQuantization : public OpExprGradFunction<FakeQuantizationCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(FakeQuantizationInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(FakeQuantizationCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 3);
     ctx->requires_grad = inputs.at(0)->requires_grad();
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const FakeQuantizationInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const FakeQuantizationCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(3);
diff --git a/oneflow/core/autograd/gradient_funcs/flatten.cpp b/oneflow/core/autograd/gradient_funcs/flatten.cpp
index e8effbccd17..4b0f46bb7bd 100644
--- a/oneflow/core/autograd/gradient_funcs/flatten.cpp
+++ b/oneflow/core/autograd/gradient_funcs/flatten.cpp
@@ -23,16 +23,16 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct FlattenInterpState : public OpExprInterpState {
+struct FlattenCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class Flatten : public OpExprGradFunction<FlattenInterpState> {
+class Flatten : public OpExprGradFunction<FlattenCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(FlattenInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(FlattenCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const FlattenInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const FlattenCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -47,7 +47,7 @@ Maybe<void> Flatten::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Flatten::Capture(FlattenInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Flatten::Capture(FlattenCaptureState* ctx, const TensorTuple& inputs,
                              const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -55,7 +55,7 @@ Maybe<void> Flatten::Capture(FlattenInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Flatten::Apply(const FlattenInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Flatten::Apply(const FlattenCaptureState* ctx, const TensorTuple& out_grads,
                            TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/flip.cpp b/oneflow/core/autograd/gradient_funcs/flip.cpp
index 1022e4900a2..655821f4bd0 100644
--- a/oneflow/core/autograd/gradient_funcs/flip.cpp
+++ b/oneflow/core/autograd/gradient_funcs/flip.cpp
@@ -20,17 +20,17 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct FlipInterpState : public OpExprInterpState {
+struct FlipCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   std::vector<int32_t> dims;
 };
 
-class Flip : public OpExprGradFunction<FlipInterpState> {
+class Flip : public OpExprGradFunction<FlipCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(FlipInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+  Maybe<void> Capture(FlipCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
                       const AttrMap& attrs) const override;
-  Maybe<void> Apply(const FlipInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const FlipCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -44,7 +44,7 @@ Maybe<void> Flip::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Flip::Capture(FlipInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Flip::Capture(FlipCaptureState* ctx, const TensorTuple& inputs,
                           const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -53,7 +53,7 @@ Maybe<void> Flip::Capture(FlipInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Flip::Apply(const FlipInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Flip::Apply(const FlipCaptureState* ctx, const TensorTuple& out_grads,
                         TensorTuple* in_grads) const {
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
   in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/gather.cpp b/oneflow/core/autograd/gradient_funcs/gather.cpp
index 0f83aceb71f..042f193f3c0 100644
--- a/oneflow/core/autograd/gradient_funcs/gather.cpp
+++ b/oneflow/core/autograd/gradient_funcs/gather.cpp
@@ -23,17 +23,17 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct GatherInterpState : public OpExprInterpState {
+struct GatherCaptureState : public AutoGradCaptureState {
   int64_t axis;
   bool requires_grad;
 };
 
-class Gather : public OpExprGradFunction<GatherInterpState> {
+class Gather : public OpExprGradFunction<GatherCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(GatherInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
-                      const AttrMap& attrs) const override;
-  Maybe<void> Apply(const GatherInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Capture(GatherCaptureState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override;
+  Maybe<void> Apply(const GatherCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -47,7 +47,7 @@ Maybe<void> Gather::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Gather::Capture(GatherInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Gather::Capture(GatherCaptureState* ctx, const TensorTuple& inputs,
                             const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -60,7 +60,7 @@ Maybe<void> Gather::Capture(GatherInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Gather::Apply(const GatherInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Gather::Apply(const GatherCaptureState* ctx, const TensorTuple& out_grads,
                           TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/gather_nd.cpp b/oneflow/core/autograd/gradient_funcs/gather_nd.cpp
index 3ba5b07a21e..84764cb953c 100644
--- a/oneflow/core/autograd/gradient_funcs/gather_nd.cpp
+++ b/oneflow/core/autograd/gradient_funcs/gather_nd.cpp
@@ -19,15 +19,15 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct GatherNdInterpState : public OpExprInterpState {
+struct GatherNdCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class GatherNd : public OpExprGradFunction<GatherNdInterpState> {
+class GatherNd : public OpExprGradFunction<GatherNdCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(GatherNdInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(GatherNdCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 2);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -39,7 +39,7 @@ class GatherNd : public OpExprGradFunction<GatherNdInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const GatherNdInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const GatherNdCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(2);
diff --git a/oneflow/core/autograd/gradient_funcs/hierarchical_parallel_cast.cpp b/oneflow/core/autograd/gradient_funcs/hierarchical_parallel_cast.cpp
index 8dc2bce6207..c591adb49f6 100644
--- a/oneflow/core/autograd/gradient_funcs/hierarchical_parallel_cast.cpp
+++ b/oneflow/core/autograd/gradient_funcs/hierarchical_parallel_cast.cpp
@@ -49,12 +49,11 @@ Maybe<one::UserOpExpr> FindOrCreatHierarchicalParallelCastOpExpr(Symbol<cfg::NdS
 
 }  // namespace
 
-struct HerarchicalParallelCastOpExprInterpState : public OpExprInterpState {
+struct HerarchicalParallelCastCaptureState : public AutoGradCaptureState {
   Symbol<cfg::NdSbp> nd_sbp;
 };
 
-class HerarchicalParallelCast
-    : public OpExprGradFunction<HerarchicalParallelCastOpExprInterpState> {
+class HerarchicalParallelCast : public OpExprGradFunction<HerarchicalParallelCastCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -62,14 +61,14 @@ class HerarchicalParallelCast
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(HerarchicalParallelCastOpExprInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(HerarchicalParallelCastCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     ctx->nd_sbp = JUST(inputs.at(0)->nd_sbp());
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const HerarchicalParallelCastOpExprInterpState* ctx,
-                    const TensorTuple& out_grads, TensorTuple* in_grads) const override {
+  Maybe<void> Apply(const HerarchicalParallelCastCaptureState* ctx, const TensorTuple& out_grads,
+                    TensorTuple* in_grads) const override {
     const auto& grad_op = JUST(FindOrCreatHierarchicalParallelCastOpExpr(ctx->nd_sbp));
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/identity.cpp b/oneflow/core/autograd/gradient_funcs/identity.cpp
index c47a14c5f51..0c929f0284b 100644
--- a/oneflow/core/autograd/gradient_funcs/identity.cpp
+++ b/oneflow/core/autograd/gradient_funcs/identity.cpp
@@ -18,22 +18,22 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct IdentityInterpState : public OpExprInterpState {
+struct IdentityCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class Identity : public OpExprGradFunction<IdentityInterpState> {
+class Identity : public OpExprGradFunction<IdentityCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(IdentityInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(IdentityCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     ctx->requires_grad = inputs.at(0)->requires_grad();
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const IdentityInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const IdentityCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/l2_normalize.cpp b/oneflow/core/autograd/gradient_funcs/l2_normalize.cpp
index 5dc22de5b8f..bfa7d5687e6 100644
--- a/oneflow/core/autograd/gradient_funcs/l2_normalize.cpp
+++ b/oneflow/core/autograd/gradient_funcs/l2_normalize.cpp
@@ -22,18 +22,18 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct L2NormalizeInterpState : public OpExprInterpState {
+struct L2NormalizeCaptureState : public AutoGradCaptureState {
   int64_t axis;
   float epsilon;
   bool requires_grad;
 };
 
-class L2Normalize : public OpExprGradFunction<L2NormalizeInterpState> {
+class L2Normalize : public OpExprGradFunction<L2NormalizeCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(L2NormalizeInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(L2NormalizeCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const L2NormalizeInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const L2NormalizeCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -47,7 +47,7 @@ Maybe<void> L2Normalize::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> L2Normalize::Capture(L2NormalizeInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> L2Normalize::Capture(L2NormalizeCaptureState* ctx, const TensorTuple& inputs,
                                  const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -61,7 +61,7 @@ Maybe<void> L2Normalize::Capture(L2NormalizeInterpState* ctx, const TensorTuple&
   return Maybe<void>::Ok();
 }
 
-Maybe<void> L2Normalize::Apply(const L2NormalizeInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> L2Normalize::Apply(const L2NormalizeCaptureState* ctx, const TensorTuple& out_grads,
                                TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/layer_norm.cpp b/oneflow/core/autograd/gradient_funcs/layer_norm.cpp
index fcaa5402bb3..fc950643853 100644
--- a/oneflow/core/autograd/gradient_funcs/layer_norm.cpp
+++ b/oneflow/core/autograd/gradient_funcs/layer_norm.cpp
@@ -23,7 +23,7 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct LayerNormInterpState : public OpExprInterpState {
+struct LayerNormCaptureState : public AutoGradCaptureState {
   bool center;
   bool scale;
 
@@ -47,14 +47,14 @@ struct LayerNormInterpState : public OpExprInterpState {
 // y, mean, inv_variance, [normalized] =
 //   layer_norm(x, [beta], [gamma], center=False, scale=False, begin_norm_axis=1,
 //              begin_params_axis=-1, epsilon=1e-5)
-class LayerNorm : public OpExprGradFunction<LayerNormInterpState> {
+class LayerNorm : public OpExprGradFunction<LayerNormCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
 
-  Maybe<void> Capture(LayerNormInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(LayerNormCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
 
-  Maybe<void> Apply(const LayerNormInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const LayerNormCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -73,7 +73,7 @@ Maybe<void> LayerNorm::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> LayerNorm::Capture(LayerNormInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> LayerNorm::Capture(LayerNormCaptureState* ctx, const TensorTuple& inputs,
                                const TensorTuple& outputs, const AttrMap& attrs) const {
   ComposedAttrMap composed_attrs(attrs, base_attrs_);
   ctx->center = JUST(composed_attrs.GetAttr<bool>("center"));
@@ -101,7 +101,7 @@ Maybe<void> LayerNorm::Capture(LayerNormInterpState* ctx, const TensorTuple& inp
   return Maybe<void>::Ok();
 }
 
-Maybe<void> LayerNorm::Apply(const LayerNormInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> LayerNorm::Apply(const LayerNormCaptureState* ctx, const TensorTuple& out_grads,
                              TensorTuple* in_grads) const {
   const auto& saved_tensors = ctx->SavedTensors();
   in_grads->resize(ctx->center + ctx->scale + 1);
diff --git a/oneflow/core/autograd/gradient_funcs/logsoftmax.cpp b/oneflow/core/autograd/gradient_funcs/logsoftmax.cpp
index fc247d18883..ef50dbff2d1 100644
--- a/oneflow/core/autograd/gradient_funcs/logsoftmax.cpp
+++ b/oneflow/core/autograd/gradient_funcs/logsoftmax.cpp
@@ -23,16 +23,16 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct LogSoftmaxInterpState : public OpExprInterpState {
+struct LogSoftmaxCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class LogSoftmax : public OpExprGradFunction<LogSoftmaxInterpState> {
+class LogSoftmax : public OpExprGradFunction<LogSoftmaxCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(LogSoftmaxInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(LogSoftmaxCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const LogSoftmaxInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const LogSoftmaxCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -53,7 +53,7 @@ Maybe<void> LogSoftmax::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> LogSoftmax::Capture(LogSoftmaxInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> LogSoftmax::Capture(LogSoftmaxCaptureState* ctx, const TensorTuple& inputs,
                                 const TensorTuple& outputs, const AttrMap& attrs) const {
   ComposedAttrMap composed_attrs(attrs, base_attrs_);
   CHECK_EQ_OR_RETURN(inputs.size(), 1);
@@ -65,7 +65,7 @@ Maybe<void> LogSoftmax::Capture(LogSoftmaxInterpState* ctx, const TensorTuple& i
   return Maybe<void>::Ok();
 }
 
-Maybe<void> LogSoftmax::Apply(const LogSoftmaxInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> LogSoftmax::Apply(const LogSoftmaxCaptureState* ctx, const TensorTuple& out_grads,
                               TensorTuple* in_grads) const {
   if (!ctx->requires_grad) return Maybe<void>::Ok();
   CHECK_EQ_OR_RETURN(out_grads.size(), 2);
diff --git a/oneflow/core/autograd/gradient_funcs/math_binary_op.cpp b/oneflow/core/autograd/gradient_funcs/math_binary_op.cpp
index df653f3de16..bdb19599c66 100644
--- a/oneflow/core/autograd/gradient_funcs/math_binary_op.cpp
+++ b/oneflow/core/autograd/gradient_funcs/math_binary_op.cpp
@@ -23,13 +23,13 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct BinaryMathOpExprInterpState : public OpExprInterpState {
+struct BinaryMathCaptureState : public AutoGradCaptureState {
   bool x_requires_grad;
   bool y_requires_grad;
 };
 
-class BinaryMathOp : public OpExprGradFunction<BinaryMathOpExprInterpState> {
-  Maybe<void> Capture(BinaryMathOpExprInterpState* ctx, const TensorTuple& inputs,
+class BinaryMathOp : public OpExprGradFunction<BinaryMathCaptureState> {
+  Maybe<void> Capture(BinaryMathCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     ctx->x_requires_grad = inputs.at(0)->requires_grad();
     ctx->y_requires_grad = inputs.at(1)->requires_grad();
@@ -38,7 +38,7 @@ class BinaryMathOp : public OpExprGradFunction<BinaryMathOpExprInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const BinaryMathOpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const BinaryMathCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (!(ctx->x_requires_grad || ctx->y_requires_grad)) { return Maybe<void>::Ok(); }
 
diff --git a/oneflow/core/autograd/gradient_funcs/math_unary_op.cpp b/oneflow/core/autograd/gradient_funcs/math_unary_op.cpp
index 797f9112b9d..5a383cbecf0 100644
--- a/oneflow/core/autograd/gradient_funcs/math_unary_op.cpp
+++ b/oneflow/core/autograd/gradient_funcs/math_unary_op.cpp
@@ -23,19 +23,19 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct UnaryMathOpExprInterpState : public OpExprInterpState {
+struct UnaryMathCaptureState : public AutoGradCaptureState {
   bool x_requires_grad;
 };
 
-class UnaryMathOp : public OpExprGradFunction<UnaryMathOpExprInterpState> {
-  Maybe<void> Capture(UnaryMathOpExprInterpState* ctx, const TensorTuple& inputs,
+class UnaryMathOp : public OpExprGradFunction<UnaryMathCaptureState> {
+  Maybe<void> Capture(UnaryMathCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     ctx->x_requires_grad = inputs.at(0)->requires_grad();
     ctx->SaveTensorForBackward(inputs.at(0));
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const UnaryMathOpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const UnaryMathCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (!ctx->x_requires_grad) { return Maybe<void>::Ok(); }
     const auto& x = ctx->SavedTensors().at(0);
diff --git a/oneflow/core/autograd/gradient_funcs/matmul.cpp b/oneflow/core/autograd/gradient_funcs/matmul.cpp
index b684c2739c7..283768fb666 100644
--- a/oneflow/core/autograd/gradient_funcs/matmul.cpp
+++ b/oneflow/core/autograd/gradient_funcs/matmul.cpp
@@ -22,7 +22,7 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct MatmulInterpState : public OpExprInterpState {
+struct MatmulCaptureState : public AutoGradCaptureState {
   bool transpose_a;
   bool transpose_b;
   double alpha;
@@ -32,11 +32,11 @@ struct MatmulInterpState : public OpExprInterpState {
   size_t b_index;
 };
 
-class MatmulBase : public OpExprGradFunction<MatmulInterpState> {
+class MatmulBase : public OpExprGradFunction<MatmulCaptureState> {
  public:
-  Maybe<void> Capture(MatmulInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
-                      const AttrMap& attrs) const override;
-  Maybe<void> Apply(const MatmulInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Capture(MatmulCaptureState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override;
+  Maybe<void> Apply(const MatmulCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  protected:
@@ -45,7 +45,7 @@ class MatmulBase : public OpExprGradFunction<MatmulInterpState> {
   std::shared_ptr<OpExpr> grad_b_op_;
 };
 
-Maybe<void> MatmulBase::Capture(MatmulInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> MatmulBase::Capture(MatmulCaptureState* ctx, const TensorTuple& inputs,
                                 const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad_a = inputs.at(0)->requires_grad();
   ctx->requires_grad_b = inputs.at(1)->requires_grad();
@@ -64,7 +64,7 @@ Maybe<void> MatmulBase::Capture(MatmulInterpState* ctx, const TensorTuple& input
   return Maybe<void>::Ok();
 }
 
-Maybe<void> MatmulBase::Apply(const MatmulInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> MatmulBase::Apply(const MatmulCaptureState* ctx, const TensorTuple& out_grads,
                               TensorTuple* in_grads) const {
   if (!ctx->requires_grad_a && !ctx->requires_grad_b) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
@@ -150,7 +150,7 @@ REGISTER_OP_EXPR_GRAD_FUNCTION("batch_matmul", BatchMatmul);
 class BroadcastMatmul : public MatmulBase {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Apply(const MatmulInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const MatmulCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 };
 
@@ -168,7 +168,7 @@ Maybe<void> BroadcastMatmul::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> BroadcastMatmul::Apply(const MatmulInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> BroadcastMatmul::Apply(const MatmulCaptureState* ctx, const TensorTuple& out_grads,
                                    TensorTuple* in_grads) const {
   if (!ctx->requires_grad_a && !ctx->requires_grad_b) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/multiply.cpp b/oneflow/core/autograd/gradient_funcs/multiply.cpp
index ff37dd5a40e..6ac80541b5e 100644
--- a/oneflow/core/autograd/gradient_funcs/multiply.cpp
+++ b/oneflow/core/autograd/gradient_funcs/multiply.cpp
@@ -19,19 +19,19 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct MultiplyInterpState : public OpExprInterpState {
+struct MultiplyCaptureState : public AutoGradCaptureState {
   bool requires_grad_x;
   bool requires_grad_y;
   int32_t index_x;
   int32_t index_y;
 };
 
-class Multiply : public OpExprGradFunction<MultiplyInterpState> {
+class Multiply : public OpExprGradFunction<MultiplyCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(MultiplyInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(MultiplyCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const MultiplyInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const MultiplyCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -45,7 +45,7 @@ Maybe<void> Multiply::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Multiply::Capture(MultiplyInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Multiply::Capture(MultiplyCaptureState* ctx, const TensorTuple& inputs,
                               const TensorTuple& outputs, const AttrMap& attrs) const {
   CHECK_EQ_OR_RETURN(inputs.size(), 2);
   ctx->requires_grad_x = inputs.at(0)->requires_grad();
@@ -55,7 +55,7 @@ Maybe<void> Multiply::Capture(MultiplyInterpState* ctx, const TensorTuple& input
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Multiply::Apply(const MultiplyInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Multiply::Apply(const MultiplyCaptureState* ctx, const TensorTuple& out_grads,
                             TensorTuple* in_grads) const {
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
   in_grads->resize(2);
diff --git a/oneflow/core/autograd/gradient_funcs/narrow.cpp b/oneflow/core/autograd/gradient_funcs/narrow.cpp
index ddaccd47ab2..c5e2a427856 100644
--- a/oneflow/core/autograd/gradient_funcs/narrow.cpp
+++ b/oneflow/core/autograd/gradient_funcs/narrow.cpp
@@ -21,14 +21,14 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct NarrowOpInterpState : public OpExprInterpState {
+struct NarrowCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   int64_t dim;
   int64_t start;
   int64_t length;
 };
 
-class NarrowOp : public OpExprGradFunction<NarrowOpInterpState> {
+class Narrow : public OpExprGradFunction<NarrowCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -37,7 +37,7 @@ class NarrowOp : public OpExprGradFunction<NarrowOpInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(NarrowOpInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(NarrowCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -52,7 +52,7 @@ class NarrowOp : public OpExprGradFunction<NarrowOpInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const NarrowOpInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const NarrowCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (ctx->requires_grad) {
       const auto& like = ctx->SavedTensors().at(0);
@@ -67,7 +67,7 @@ class NarrowOp : public OpExprGradFunction<NarrowOpInterpState> {
   AttrMap base_attrs_;
 };
 
-REGISTER_OP_EXPR_GRAD_FUNCTION("narrow", NarrowOp);
+REGISTER_OP_EXPR_GRAD_FUNCTION("narrow", Narrow);
 
 }  // namespace one
 }  // namespace oneflow
diff --git a/oneflow/core/autograd/gradient_funcs/normalization.cpp b/oneflow/core/autograd/gradient_funcs/normalization.cpp
index 53b6973841a..b3c71194f02 100644
--- a/oneflow/core/autograd/gradient_funcs/normalization.cpp
+++ b/oneflow/core/autograd/gradient_funcs/normalization.cpp
@@ -21,7 +21,7 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct NormalizationGradInterpState : public OpExprInterpState {
+struct NormalizationGradCaptureState : public AutoGradCaptureState {
   int32_t axis;
   float epsilon;
   bool track_running_stats;
@@ -39,7 +39,7 @@ struct NormalizationGradInterpState : public OpExprInterpState {
 // inference:
 // y = normalization(x, moving_mean, moving_variance, gamma, beta, axis=1, epsilon=0.01,
 // momentum=0.9)
-class NormalizationGrad : public OpExprGradFunction<NormalizationGradInterpState> {
+class NormalizationGrad : public OpExprGradFunction<NormalizationGradCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -48,7 +48,7 @@ class NormalizationGrad : public OpExprGradFunction<NormalizationGradInterpState
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(NormalizationGradInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(NormalizationGradCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     ctx->x_requires_grad = inputs.at(0)->requires_grad();
     std::shared_ptr<Tensor> gamma, beta;
@@ -81,7 +81,7 @@ class NormalizationGrad : public OpExprGradFunction<NormalizationGradInterpState
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const NormalizationGradInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const NormalizationGradCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const auto& x = ctx->SavedTensors().at(0);      // x
     const auto& gamma = ctx->SavedTensors().at(1);  // gamma
diff --git a/oneflow/core/autograd/gradient_funcs/padding.cpp b/oneflow/core/autograd/gradient_funcs/padding.cpp
index 3d54edadeec..040549730e5 100644
--- a/oneflow/core/autograd/gradient_funcs/padding.cpp
+++ b/oneflow/core/autograd/gradient_funcs/padding.cpp
@@ -19,12 +19,12 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct Pad2dInterpState : public OpExprInterpState {
+struct Pad2dCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   std::vector<int64_t> paddings;
 };
 
-class Pad2d : public OpExprGradFunction<Pad2dInterpState> {
+class Pad2d : public OpExprGradFunction<Pad2dCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const UserOpExpr* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -33,7 +33,7 @@ class Pad2d : public OpExprGradFunction<Pad2dInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(Pad2dInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+  Maybe<void> Capture(Pad2dCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
                       const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -51,7 +51,7 @@ class Pad2d : public OpExprGradFunction<Pad2dInterpState> {
 
 class ReflectionPad2d : public Pad2d {
  public:
-  Maybe<void> Apply(const Pad2dInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const Pad2dCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -64,7 +64,7 @@ class ReflectionPad2d : public Pad2d {
 
 class ReplicationPad2d : public Pad2d {
  public:
-  Maybe<void> Apply(const Pad2dInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const Pad2dCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
@@ -78,13 +78,13 @@ class ReplicationPad2d : public Pad2d {
 REGISTER_OP_EXPR_GRAD_FUNCTION("reflection_pad2d", ReflectionPad2d);
 REGISTER_OP_EXPR_GRAD_FUNCTION("replication_pad2d", ReplicationPad2d);
 
-struct ConstantPadNdInterpState : public OpExprInterpState {
+struct ConstantPadNdCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   std::vector<int64_t> paddings;
   functional::Scalar padding_value;
 };
 
-class ConstantPadNd : public OpExprGradFunction<ConstantPadNdInterpState> {
+class ConstantPadNd : public OpExprGradFunction<ConstantPadNdCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const UserOpExpr* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -93,7 +93,7 @@ class ConstantPadNd : public OpExprGradFunction<ConstantPadNdInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(ConstantPadNdInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ConstantPadNdCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -112,7 +112,7 @@ class ConstantPadNd : public OpExprGradFunction<ConstantPadNdInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const ConstantPadNdInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ConstantPadNdCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/pool.cpp b/oneflow/core/autograd/gradient_funcs/pool.cpp
index 49d9880a79f..e1ae9dc278d 100644
--- a/oneflow/core/autograd/gradient_funcs/pool.cpp
+++ b/oneflow/core/autograd/gradient_funcs/pool.cpp
@@ -26,7 +26,7 @@ namespace one {
 
 namespace {
 
-struct PoolInterpState : public OpExprInterpState {
+struct PoolCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   size_t input_index;
   size_t output_index;
@@ -40,16 +40,16 @@ struct PoolInterpState : public OpExprInterpState {
   bool ceil_mode;
 };
 
-class PoolNdGrad : public OpExprGradFunction<PoolInterpState> {
+class PoolNdGrad : public OpExprGradFunction<PoolCaptureState> {
  public:
   virtual ~PoolNdGrad() = default;
 
-  using OpExprGradFunction<PoolInterpState>::Init;
+  using OpExprGradFunction<PoolCaptureState>::Init;
 
   Maybe<void> Init(const OpExpr& op, const std::string& mode);
-  Maybe<void> Capture(PoolInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+  Maybe<void> Capture(PoolCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
                       const AttrMap& attrs) const override;
-  Maybe<void> Apply(const PoolInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const PoolCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -65,7 +65,7 @@ Maybe<void> PoolNdGrad::Init(const OpExpr& op, const std::string& mode) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> PoolNdGrad::Capture(PoolInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> PoolNdGrad::Capture(PoolCaptureState* ctx, const TensorTuple& inputs,
                                 const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -84,7 +84,7 @@ Maybe<void> PoolNdGrad::Capture(PoolInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> PoolNdGrad::Apply(const PoolInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> PoolNdGrad::Apply(const PoolCaptureState* ctx, const TensorTuple& out_grads,
                               TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/pooling.cpp b/oneflow/core/autograd/gradient_funcs/pooling.cpp
index 7111d91195e..e0c811870ca 100644
--- a/oneflow/core/autograd/gradient_funcs/pooling.cpp
+++ b/oneflow/core/autograd/gradient_funcs/pooling.cpp
@@ -27,7 +27,7 @@ namespace one {
 
 namespace {
 
-struct PoolingInterpState : public OpExprInterpState {
+struct PoolingCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   size_t input_index;
   size_t output_index;
@@ -42,16 +42,16 @@ struct PoolingInterpState : public OpExprInterpState {
   bool ceil_mode;
 };
 
-class PoolingNdGrad : public OpExprGradFunction<PoolingInterpState> {
+class PoolingNdGrad : public OpExprGradFunction<PoolingCaptureState> {
  public:
   virtual ~PoolingNdGrad() = default;
 
-  using OpExprGradFunction<PoolingInterpState>::Init;
+  using OpExprGradFunction<PoolingCaptureState>::Init;
 
   Maybe<void> Init(const OpExpr& op, const std::string& mode);
-  Maybe<void> Capture(PoolingInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(PoolingCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const PoolingInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const PoolingCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -67,7 +67,7 @@ Maybe<void> PoolingNdGrad::Init(const OpExpr& op, const std::string& mode) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> PoolingNdGrad::Capture(PoolingInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> PoolingNdGrad::Capture(PoolingCaptureState* ctx, const TensorTuple& inputs,
                                    const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -87,7 +87,7 @@ Maybe<void> PoolingNdGrad::Capture(PoolingInterpState* ctx, const TensorTuple& i
   return Maybe<void>::Ok();
 }
 
-Maybe<void> PoolingNdGrad::Apply(const PoolingInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> PoolingNdGrad::Apply(const PoolingCaptureState* ctx, const TensorTuple& out_grads,
                                  TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_LE_OR_RETURN(out_grads.size(), 2);
diff --git a/oneflow/core/autograd/gradient_funcs/reduce_ops.cpp b/oneflow/core/autograd/gradient_funcs/reduce_ops.cpp
index 987d1f8f710..0a28fb336e3 100644
--- a/oneflow/core/autograd/gradient_funcs/reduce_ops.cpp
+++ b/oneflow/core/autograd/gradient_funcs/reduce_ops.cpp
@@ -22,39 +22,39 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ReduceSumOpInterpState : public OpExprInterpState {
+struct ReduceSumCaptureState : public AutoGradCaptureState {
   std::vector<int32_t> axis;
 };
 
-class ReduceSumOp : public OpExprGradFunction<ReduceSumOpInterpState> {
+class ReduceSum : public OpExprGradFunction<ReduceSumCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(ReduceSumOpInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ReduceSumCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const ReduceSumOpInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ReduceSumCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
   AttrMap base_attrs_;
 };
 
-Maybe<void> ReduceSumOp::Init(const OpExpr& op) {
+Maybe<void> ReduceSum::Init(const OpExpr& op) {
   const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
   CHECK_NOTNULL_OR_RETURN(fw_op_expr);
   base_attrs_ = MakeAttrMapFromUserOpConf(fw_op_expr->proto());
   return Maybe<void>::Ok();
 }
 
-Maybe<void> ReduceSumOp::Capture(ReduceSumOpInterpState* ctx, const TensorTuple& inputs,
-                                 const TensorTuple& outputs, const AttrMap& attrs) const {
+Maybe<void> ReduceSum::Capture(ReduceSumCaptureState* ctx, const TensorTuple& inputs,
+                               const TensorTuple& outputs, const AttrMap& attrs) const {
   ComposedAttrMap composed_attrs(attrs, base_attrs_);
   ctx->axis = JUST(composed_attrs.GetAttr<std::vector<int32_t>>("axis"));
   ctx->SaveTensorForBackward(inputs.at(0));
   return Maybe<void>::Ok();
 }
 
-Maybe<void> ReduceSumOp::Apply(const ReduceSumOpInterpState* ctx, const TensorTuple& out_grads,
-                               TensorTuple* in_grads) const {
+Maybe<void> ReduceSum::Apply(const ReduceSumCaptureState* ctx, const TensorTuple& out_grads,
+                             TensorTuple* in_grads) const {
   const auto& input = ctx->SavedTensors().at(0);
   const auto& dy = out_grads.at(0);
   in_grads->resize(1);
@@ -62,34 +62,34 @@ Maybe<void> ReduceSumOp::Apply(const ReduceSumOpInterpState* ctx, const TensorTu
   return Maybe<void>::Ok();
 }
 
-REGISTER_OP_EXPR_GRAD_FUNCTION("reduce_sum", ReduceSumOp);
+REGISTER_OP_EXPR_GRAD_FUNCTION("reduce_sum", ReduceSum);
 
-struct ReduceMaxOrMinOpInterpState : public OpExprInterpState {
+struct ReduceMaxOrMinCaptureState : public AutoGradCaptureState {
   std::vector<int32_t> axis;
   bool keepdims;
 };
 
-class ReduceMaxOrMinOp : public OpExprGradFunction<ReduceMaxOrMinOpInterpState> {
+class ReduceMaxOrMin : public OpExprGradFunction<ReduceMaxOrMinCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(ReduceMaxOrMinOpInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ReduceMaxOrMinCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const ReduceMaxOrMinOpInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ReduceMaxOrMinCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
   AttrMap base_attrs_;
 };
 
-Maybe<void> ReduceMaxOrMinOp::Init(const OpExpr& op) {
+Maybe<void> ReduceMaxOrMin::Init(const OpExpr& op) {
   const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
   CHECK_NOTNULL_OR_RETURN(fw_op_expr);
   base_attrs_ = MakeAttrMapFromUserOpConf(fw_op_expr->proto());
   return Maybe<void>::Ok();
 }
 
-Maybe<void> ReduceMaxOrMinOp::Capture(ReduceMaxOrMinOpInterpState* ctx, const TensorTuple& inputs,
-                                      const TensorTuple& outputs, const AttrMap& attrs) const {
+Maybe<void> ReduceMaxOrMin::Capture(ReduceMaxOrMinCaptureState* ctx, const TensorTuple& inputs,
+                                    const TensorTuple& outputs, const AttrMap& attrs) const {
   ComposedAttrMap composed_attrs(attrs, base_attrs_);
   ctx->axis = JUST(composed_attrs.GetAttr<std::vector<int32_t>>("axis"));
   ctx->keepdims = JUST(composed_attrs.GetAttr<bool>("keepdims"));
@@ -98,8 +98,8 @@ Maybe<void> ReduceMaxOrMinOp::Capture(ReduceMaxOrMinOpInterpState* ctx, const Te
   return Maybe<void>::Ok();
 }
 
-Maybe<void> ReduceMaxOrMinOp::Apply(const ReduceMaxOrMinOpInterpState* ctx,
-                                    const TensorTuple& out_grads, TensorTuple* in_grads) const {
+Maybe<void> ReduceMaxOrMin::Apply(const ReduceMaxOrMinCaptureState* ctx,
+                                  const TensorTuple& out_grads, TensorTuple* in_grads) const {
   const auto& input = ctx->SavedTensors().at(0);
   const auto& output = ctx->SavedTensors().at(1);
   const auto& dy = out_grads.at(0);
@@ -116,8 +116,8 @@ Maybe<void> ReduceMaxOrMinOp::Apply(const ReduceMaxOrMinOpInterpState* ctx,
   return Maybe<void>::Ok();
 }
 
-REGISTER_OP_EXPR_GRAD_FUNCTION("reduce_min", ReduceMaxOrMinOp);
-REGISTER_OP_EXPR_GRAD_FUNCTION("reduce_max", ReduceMaxOrMinOp);
+REGISTER_OP_EXPR_GRAD_FUNCTION("reduce_min", ReduceMaxOrMin);
+REGISTER_OP_EXPR_GRAD_FUNCTION("reduce_max", ReduceMaxOrMin);
 
 }  // namespace one
 }  // namespace oneflow
diff --git a/oneflow/core/autograd/gradient_funcs/reshape.cpp b/oneflow/core/autograd/gradient_funcs/reshape.cpp
index dcc99995288..c743e4cb6e9 100644
--- a/oneflow/core/autograd/gradient_funcs/reshape.cpp
+++ b/oneflow/core/autograd/gradient_funcs/reshape.cpp
@@ -24,7 +24,7 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-class ReshapeOpExprGrad : public OpExprGradFunction<OpExprInterpState> {
+class ReshapeOpExprGrad : public OpExprGradFunction<AutoGradCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -33,13 +33,13 @@ class ReshapeOpExprGrad : public OpExprGradFunction<OpExprInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(OpExprInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
-                      const AttrMap& attrs) const override {
+  Maybe<void> Capture(AutoGradCaptureState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override {
     ctx->SaveTensorForBackward(inputs.at(0));
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const OpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const AutoGradCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const auto& saved_tensors = ctx->SavedTensors();
     in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/scalar_add.cpp b/oneflow/core/autograd/gradient_funcs/scalar_add.cpp
index b8db59221a9..3cfaaa7c0ce 100644
--- a/oneflow/core/autograd/gradient_funcs/scalar_add.cpp
+++ b/oneflow/core/autograd/gradient_funcs/scalar_add.cpp
@@ -19,22 +19,22 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ScalarAddInterpState : public OpExprInterpState {
+struct ScalarAddCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class ScalarAdd : public OpExprGradFunction<ScalarAddInterpState> {
+class ScalarAdd : public OpExprGradFunction<ScalarAddCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(ScalarAddInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ScalarAddCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     ctx->requires_grad = inputs.at(0)->requires_grad();
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const ScalarAddInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ScalarAddCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/scalar_fmod.cpp b/oneflow/core/autograd/gradient_funcs/scalar_fmod.cpp
index d5d1b45ece9..0922f391054 100644
--- a/oneflow/core/autograd/gradient_funcs/scalar_fmod.cpp
+++ b/oneflow/core/autograd/gradient_funcs/scalar_fmod.cpp
@@ -20,22 +20,22 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ScalarFModGradInterpState : public OpExprInterpState {
+struct ScalarFModGradCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class ScalarFModGrad : public OpExprGradFunction<ScalarFModGradInterpState> {
+class ScalarFModGrad : public OpExprGradFunction<ScalarFModGradCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(ScalarFModGradInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ScalarFModGradCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     ctx->requires_grad = inputs.at(0)->requires_grad();
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const ScalarFModGradInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ScalarFModGradCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/scalar_mul.cpp b/oneflow/core/autograd/gradient_funcs/scalar_mul.cpp
index bc64ff10126..6f9942bb56e 100644
--- a/oneflow/core/autograd/gradient_funcs/scalar_mul.cpp
+++ b/oneflow/core/autograd/gradient_funcs/scalar_mul.cpp
@@ -20,12 +20,12 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ScalarMulInterpState : public OpExprInterpState {
+struct ScalarMulCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   functional::Scalar operand;
 };
 
-class ScalarMul : public OpExprGradFunction<ScalarMulInterpState> {
+class ScalarMul : public OpExprGradFunction<ScalarMulCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -34,7 +34,7 @@ class ScalarMul : public OpExprGradFunction<ScalarMulInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(ScalarMulInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ScalarMulCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     ctx->requires_grad = inputs.at(0)->requires_grad();
@@ -49,7 +49,7 @@ class ScalarMul : public OpExprGradFunction<ScalarMulInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const ScalarMulInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ScalarMulCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/scalar_pow.cpp b/oneflow/core/autograd/gradient_funcs/scalar_pow.cpp
index 03f021aaad5..19946f3d228 100644
--- a/oneflow/core/autograd/gradient_funcs/scalar_pow.cpp
+++ b/oneflow/core/autograd/gradient_funcs/scalar_pow.cpp
@@ -22,12 +22,12 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ScalarPowInterpState : public OpExprInterpState {
+struct ScalarPowCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   double exponent;
 };
 
-class ScalarPow : public OpExprGradFunction<ScalarPowInterpState> {
+class ScalarPow : public OpExprGradFunction<ScalarPowCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -38,7 +38,7 @@ class ScalarPow : public OpExprGradFunction<ScalarPowInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(ScalarPowInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ScalarPowCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -51,7 +51,7 @@ class ScalarPow : public OpExprGradFunction<ScalarPowInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const ScalarPowInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ScalarPowCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const auto& x = ctx->SavedTensors().at(0);
     MutableAttrMap attrs;
diff --git a/oneflow/core/autograd/gradient_funcs/scatter_nd.cpp b/oneflow/core/autograd/gradient_funcs/scatter_nd.cpp
index 6f49d9df750..6f8119b236b 100644
--- a/oneflow/core/autograd/gradient_funcs/scatter_nd.cpp
+++ b/oneflow/core/autograd/gradient_funcs/scatter_nd.cpp
@@ -19,15 +19,15 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct ScatterNdInterpState : public OpExprInterpState {
+struct ScatterNdCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class ScatterNd : public OpExprGradFunction<ScatterNdInterpState> {
+class ScatterNd : public OpExprGradFunction<ScatterNdCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(ScatterNdInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(ScatterNdCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 2);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -38,7 +38,7 @@ class ScatterNd : public OpExprGradFunction<ScatterNdInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const ScatterNdInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const ScatterNdCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(2);
diff --git a/oneflow/core/autograd/gradient_funcs/select_first.cpp b/oneflow/core/autograd/gradient_funcs/select_first.cpp
index f14741fccb9..95ca051a336 100644
--- a/oneflow/core/autograd/gradient_funcs/select_first.cpp
+++ b/oneflow/core/autograd/gradient_funcs/select_first.cpp
@@ -23,22 +23,22 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct SelectFirstExprInterpState : public OpExprInterpState {
+struct SelectFirstCaptureState : public AutoGradCaptureState {
   TensorTuple inputs;
   bool requires_grad;
 };
 
-class SelectFirst : public OpExprGradFunction<SelectFirstExprInterpState> {
+class SelectFirst : public OpExprGradFunction<SelectFirstCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(SelectFirstExprInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(SelectFirstCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     ctx->inputs = inputs;
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const SelectFirstExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const SelectFirstCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     in_grads->at(0) = out_grads.at(0);
     for (int i = 1; i < ctx->inputs.size(); i++) {
diff --git a/oneflow/core/autograd/gradient_funcs/slice.cpp b/oneflow/core/autograd/gradient_funcs/slice.cpp
index 85b84562479..e5c9e226d79 100644
--- a/oneflow/core/autograd/gradient_funcs/slice.cpp
+++ b/oneflow/core/autograd/gradient_funcs/slice.cpp
@@ -22,14 +22,14 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct SliceOpExprInterpState : public OpExprInterpState {
+struct SliceCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   std::vector<int64_t> start;
   std::vector<int64_t> stop;
   std::vector<int64_t> step;
 };
 
-class Slice : public OpExprGradFunction<SliceOpExprInterpState> {
+class Slice : public OpExprGradFunction<SliceCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -38,8 +38,8 @@ class Slice : public OpExprGradFunction<SliceOpExprInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(SliceOpExprInterpState* ctx, const TensorTuple& inputs,
-                      const TensorTuple& outputs, const AttrMap& attrs) const override {
+  Maybe<void> Capture(SliceCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+                      const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
     ctx->requires_grad = inputs.at(0)->requires_grad();
@@ -53,7 +53,7 @@ class Slice : public OpExprGradFunction<SliceOpExprInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const SliceOpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const SliceCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     const auto& like = ctx->SavedTensors().at(0);
 
@@ -67,7 +67,7 @@ class Slice : public OpExprGradFunction<SliceOpExprInterpState> {
   AttrMap base_attrs_;
 };
 
-struct SliceUpdateOpExprInterpState : public OpExprInterpState {
+struct SliceUpdateCaptureState : public AutoGradCaptureState {
   bool requires_grad_x;
   bool requires_grad_update;
   std::vector<int64_t> start;
@@ -75,7 +75,7 @@ struct SliceUpdateOpExprInterpState : public OpExprInterpState {
   std::vector<int64_t> step;
 };
 
-class SliceUpdate : public OpExprGradFunction<SliceUpdateOpExprInterpState> {
+class SliceUpdate : public OpExprGradFunction<SliceUpdateCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -85,7 +85,7 @@ class SliceUpdate : public OpExprGradFunction<SliceUpdateOpExprInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(SliceUpdateOpExprInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(SliceUpdateCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 2);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -102,7 +102,7 @@ class SliceUpdate : public OpExprGradFunction<SliceUpdateOpExprInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const SliceUpdateOpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const SliceUpdateCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     in_grads->resize(2);
 
diff --git a/oneflow/core/autograd/gradient_funcs/smoothl1_loss.cpp b/oneflow/core/autograd/gradient_funcs/smoothl1_loss.cpp
index 05db4f22012..fcf6aa3e801 100644
--- a/oneflow/core/autograd/gradient_funcs/smoothl1_loss.cpp
+++ b/oneflow/core/autograd/gradient_funcs/smoothl1_loss.cpp
@@ -21,7 +21,7 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct SmoothL1LossInterpState : public OpExprInterpState {
+struct SmoothL1LossCaptureState : public AutoGradCaptureState {
   std::string reduction;
   float beta;
   size_t prediction_index;
@@ -29,7 +29,7 @@ struct SmoothL1LossInterpState : public OpExprInterpState {
   bool requires_grad;
 };
 
-class SmoothL1Loss : public OpExprGradFunction<SmoothL1LossInterpState> {
+class SmoothL1Loss : public OpExprGradFunction<SmoothL1LossCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
@@ -38,7 +38,7 @@ class SmoothL1Loss : public OpExprGradFunction<SmoothL1LossInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Capture(SmoothL1LossInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(SmoothL1LossCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 2);
     ctx->requires_grad = inputs.at(0)->requires_grad();  // prediction
@@ -52,7 +52,7 @@ class SmoothL1Loss : public OpExprGradFunction<SmoothL1LossInterpState> {
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const SmoothL1LossInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const SmoothL1LossCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
     in_grads->resize(2);
diff --git a/oneflow/core/autograd/gradient_funcs/softmax.cpp b/oneflow/core/autograd/gradient_funcs/softmax.cpp
index 10799b045c9..f8dfeac04eb 100644
--- a/oneflow/core/autograd/gradient_funcs/softmax.cpp
+++ b/oneflow/core/autograd/gradient_funcs/softmax.cpp
@@ -22,16 +22,16 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct SoftmaxInterpState : public OpExprInterpState {
+struct SoftmaxCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class Softmax : public OpExprGradFunction<SoftmaxInterpState> {
+class Softmax : public OpExprGradFunction<SoftmaxCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(SoftmaxInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(SoftmaxCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const SoftmaxInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const SoftmaxCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -46,7 +46,7 @@ Maybe<void> Softmax::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Softmax::Capture(SoftmaxInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Softmax::Capture(SoftmaxCaptureState* ctx, const TensorTuple& inputs,
                              const TensorTuple& outputs, const AttrMap& attrs) const {
   CHECK_EQ_OR_RETURN(inputs.size(), 1);
   ctx->requires_grad = inputs.at(0)->requires_grad();
@@ -57,7 +57,7 @@ Maybe<void> Softmax::Capture(SoftmaxInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Softmax::Apply(const SoftmaxInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Softmax::Apply(const SoftmaxCaptureState* ctx, const TensorTuple& out_grads,
                            TensorTuple* in_grads) const {
   if (!ctx->requires_grad) return Maybe<void>::Ok();
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/sparse_softmax_cross_entropy.cpp b/oneflow/core/autograd/gradient_funcs/sparse_softmax_cross_entropy.cpp
index cb1c93c09f3..3e1dea3df22 100644
--- a/oneflow/core/autograd/gradient_funcs/sparse_softmax_cross_entropy.cpp
+++ b/oneflow/core/autograd/gradient_funcs/sparse_softmax_cross_entropy.cpp
@@ -23,16 +23,16 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct SparseSoftmaxCrossEntropyInterpState : public OpExprInterpState {
+struct SparseSoftmaxCrossEntropyCaptureState : public AutoGradCaptureState {
   int64_t depth;
 };
 
-class SparseSoftmaxCrossEntropy : public OpExprGradFunction<SparseSoftmaxCrossEntropyInterpState> {
+class SparseSoftmaxCrossEntropy : public OpExprGradFunction<SparseSoftmaxCrossEntropyCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(SparseSoftmaxCrossEntropyInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(SparseSoftmaxCrossEntropyCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const SparseSoftmaxCrossEntropyInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const SparseSoftmaxCrossEntropyCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -50,7 +50,7 @@ Maybe<void> SparseSoftmaxCrossEntropy::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> SparseSoftmaxCrossEntropy::Capture(SparseSoftmaxCrossEntropyInterpState* ctx,
+Maybe<void> SparseSoftmaxCrossEntropy::Capture(SparseSoftmaxCrossEntropyCaptureState* ctx,
                                                const TensorTuple& inputs,
                                                const TensorTuple& outputs,
                                                const AttrMap& attrs) const {
@@ -63,7 +63,7 @@ Maybe<void> SparseSoftmaxCrossEntropy::Capture(SparseSoftmaxCrossEntropyInterpSt
   return Maybe<void>::Ok();
 }
 
-Maybe<void> SparseSoftmaxCrossEntropy::Apply(const SparseSoftmaxCrossEntropyInterpState* ctx,
+Maybe<void> SparseSoftmaxCrossEntropy::Apply(const SparseSoftmaxCrossEntropyCaptureState* ctx,
                                              const TensorTuple& out_grads,
                                              TensorTuple* in_grads) const {
   CHECK_EQ_OR_RETURN(out_grads.size(), 2);
diff --git a/oneflow/core/autograd/gradient_funcs/split_like.cpp b/oneflow/core/autograd/gradient_funcs/split_like.cpp
index 7d03cd5f20d..4424263384e 100644
--- a/oneflow/core/autograd/gradient_funcs/split_like.cpp
+++ b/oneflow/core/autograd/gradient_funcs/split_like.cpp
@@ -23,17 +23,17 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct SplitLikeInterpState : public OpExprInterpState {
+struct SplitLikeCaptureState : public AutoGradCaptureState {
   int64_t max_dim_size;
   bool requires_grad;
 };
 
-class SplitLike : public OpExprGradFunction<SplitLikeInterpState> {
+class SplitLike : public OpExprGradFunction<SplitLikeCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(SplitLikeInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(SplitLikeCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const SplitLikeInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const SplitLikeCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -60,7 +60,7 @@ Maybe<void> SplitLike::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> SplitLike::Capture(SplitLikeInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> SplitLike::Capture(SplitLikeCaptureState* ctx, const TensorTuple& inputs,
                                const TensorTuple& outputs, const AttrMap& attrs) const {
   CHECK_EQ_OR_RETURN(inputs.size(), outputs.size() + 1);
   ctx->requires_grad = inputs.at(0)->requires_grad();
@@ -73,7 +73,7 @@ Maybe<void> SplitLike::Capture(SplitLikeInterpState* ctx, const TensorTuple& inp
   return Maybe<void>::Ok();
 }
 
-Maybe<void> SplitLike::Apply(const SplitLikeInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> SplitLike::Apply(const SplitLikeCaptureState* ctx, const TensorTuple& out_grads,
                              TensorTuple* in_grads) const {
   in_grads->resize(1);
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
diff --git a/oneflow/core/autograd/gradient_funcs/squeeze.cpp b/oneflow/core/autograd/gradient_funcs/squeeze.cpp
index a69a600394f..8d1e56ba5a3 100644
--- a/oneflow/core/autograd/gradient_funcs/squeeze.cpp
+++ b/oneflow/core/autograd/gradient_funcs/squeeze.cpp
@@ -22,16 +22,16 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct SqueezeInterpState : public OpExprInterpState {
+struct SqueezeCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class Squeeze : public OpExprGradFunction<SqueezeInterpState> {
+class Squeeze : public OpExprGradFunction<SqueezeCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(SqueezeInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(SqueezeCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const SqueezeInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const SqueezeCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -48,7 +48,7 @@ Maybe<void> Squeeze::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Squeeze::Capture(SqueezeInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Squeeze::Capture(SqueezeCaptureState* ctx, const TensorTuple& inputs,
                              const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -57,7 +57,7 @@ Maybe<void> Squeeze::Capture(SqueezeInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Squeeze::Apply(const SqueezeInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Squeeze::Apply(const SqueezeCaptureState* ctx, const TensorTuple& out_grads,
                            TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/tensor_scalar_binary.cpp b/oneflow/core/autograd/gradient_funcs/tensor_scalar_binary.cpp
index fa3a6264aff..f95225211a4 100644
--- a/oneflow/core/autograd/gradient_funcs/tensor_scalar_binary.cpp
+++ b/oneflow/core/autograd/gradient_funcs/tensor_scalar_binary.cpp
@@ -22,18 +22,18 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct TensorScalarInterpState : public OpExprInterpState {
+struct TensorScalarCaptureState : public AutoGradCaptureState {
   bool x_requires_grad;
   bool scalar_requires_grad;
 };
 
-class TensorScalarAddOrSub : public OpExprGradFunction<TensorScalarInterpState> {
+class TensorScalarAddOrSub : public OpExprGradFunction<TensorScalarCaptureState> {
  public:
   TensorScalarAddOrSub() = default;
   virtual ~TensorScalarAddOrSub() = default;
 
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(TensorScalarInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(TensorScalarCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
 
  protected:
@@ -55,7 +55,7 @@ Maybe<void> TensorScalarAddOrSub::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> TensorScalarAddOrSub::Capture(TensorScalarInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> TensorScalarAddOrSub::Capture(TensorScalarCaptureState* ctx, const TensorTuple& inputs,
                                           const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->x_requires_grad = inputs.at(0)->requires_grad();
   ctx->scalar_requires_grad = inputs.at(1)->requires_grad();
@@ -64,7 +64,7 @@ Maybe<void> TensorScalarAddOrSub::Capture(TensorScalarInterpState* ctx, const Te
 
 class TensorScalarAdd : public TensorScalarAddOrSub {
  public:
-  Maybe<void> Apply(const TensorScalarInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const TensorScalarCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     in_grads->resize(2);
     if (ctx->x_requires_grad) {
@@ -85,7 +85,7 @@ class TensorScalarAdd : public TensorScalarAddOrSub {
 
 class TensorScalarSub : public TensorScalarAddOrSub {
  public:
-  Maybe<void> Apply(const TensorScalarInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const TensorScalarCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     in_grads->resize(2);
     if (ctx->x_requires_grad) {
@@ -108,12 +108,12 @@ class TensorScalarSub : public TensorScalarAddOrSub {
 REGISTER_OP_EXPR_GRAD_FUNCTION("scalar_add_by_tensor", TensorScalarAdd);
 REGISTER_OP_EXPR_GRAD_FUNCTION("scalar_sub_by_tensor", TensorScalarSub);
 
-class TensorScalarMul : public OpExprGradFunction<TensorScalarInterpState> {
+class TensorScalarMul : public OpExprGradFunction<TensorScalarCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(TensorScalarInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(TensorScalarCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const TensorScalarInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const TensorScalarCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -133,7 +133,7 @@ Maybe<void> TensorScalarMul::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> TensorScalarMul::Capture(TensorScalarInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> TensorScalarMul::Capture(TensorScalarCaptureState* ctx, const TensorTuple& inputs,
                                      const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->x_requires_grad = inputs.at(0)->requires_grad();
   ctx->scalar_requires_grad = inputs.at(1)->requires_grad();
@@ -142,8 +142,8 @@ Maybe<void> TensorScalarMul::Capture(TensorScalarInterpState* ctx, const TensorT
   return Maybe<void>::Ok();
 }
 
-Maybe<void> TensorScalarMul::Apply(const TensorScalarInterpState* ctx, const TensorTuple& out_grads,
-                                   TensorTuple* in_grads) const {
+Maybe<void> TensorScalarMul::Apply(const TensorScalarCaptureState* ctx,
+                                   const TensorTuple& out_grads, TensorTuple* in_grads) const {
   in_grads->resize(2);
   if (ctx->x_requires_grad) {
     const auto& scalar = ctx->SavedTensors().at(0);
@@ -165,12 +165,12 @@ Maybe<void> TensorScalarMul::Apply(const TensorScalarInterpState* ctx, const Ten
 
 REGISTER_OP_EXPR_GRAD_FUNCTION("scalar_mul_by_tensor", TensorScalarMul);
 
-class TensorScalarDiv : public OpExprGradFunction<TensorScalarInterpState> {
+class TensorScalarDiv : public OpExprGradFunction<TensorScalarCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(TensorScalarInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(TensorScalarCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const TensorScalarInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const TensorScalarCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -188,7 +188,7 @@ Maybe<void> TensorScalarDiv::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> TensorScalarDiv::Capture(TensorScalarInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> TensorScalarDiv::Capture(TensorScalarCaptureState* ctx, const TensorTuple& inputs,
                                      const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->x_requires_grad = inputs.at(0)->requires_grad();
   ctx->scalar_requires_grad = inputs.at(1)->requires_grad();
@@ -199,8 +199,8 @@ Maybe<void> TensorScalarDiv::Capture(TensorScalarInterpState* ctx, const TensorT
   return Maybe<void>::Ok();
 }
 
-Maybe<void> TensorScalarDiv::Apply(const TensorScalarInterpState* ctx, const TensorTuple& out_grads,
-                                   TensorTuple* in_grads) const {
+Maybe<void> TensorScalarDiv::Apply(const TensorScalarCaptureState* ctx,
+                                   const TensorTuple& out_grads, TensorTuple* in_grads) const {
   in_grads->resize(2);
   if (ctx->x_requires_grad) {
     const auto& scalar = ctx->SavedTensors().at(0);
diff --git a/oneflow/core/autograd/gradient_funcs/transpose.cpp b/oneflow/core/autograd/gradient_funcs/transpose.cpp
index 41e3db11238..570132edbcb 100644
--- a/oneflow/core/autograd/gradient_funcs/transpose.cpp
+++ b/oneflow/core/autograd/gradient_funcs/transpose.cpp
@@ -22,17 +22,17 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct TransposeInterpState : public OpExprInterpState {
+struct TransposeCaptureState : public AutoGradCaptureState {
   std::vector<int32_t> perm;
   bool requires_grad;
 };
 
-class Transpose : public OpExprGradFunction<TransposeInterpState> {
+class Transpose : public OpExprGradFunction<TransposeCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(TransposeInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(TransposeCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const TransposeInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const TransposeCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -50,7 +50,7 @@ Maybe<void> Transpose::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Transpose::Capture(TransposeInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Transpose::Capture(TransposeCaptureState* ctx, const TensorTuple& inputs,
                                const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -60,7 +60,7 @@ Maybe<void> Transpose::Capture(TransposeInterpState* ctx, const TensorTuple& inp
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Transpose::Apply(const TransposeInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Transpose::Apply(const TransposeCaptureState* ctx, const TensorTuple& out_grads,
                              TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/triu.cpp b/oneflow/core/autograd/gradient_funcs/triu.cpp
index 0bcce6ac7da..ed04de8074d 100644
--- a/oneflow/core/autograd/gradient_funcs/triu.cpp
+++ b/oneflow/core/autograd/gradient_funcs/triu.cpp
@@ -20,17 +20,17 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct TriuInterpState : public OpExprInterpState {
+struct TriuCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   int64_t diagonal;
 };
 
-class Triu : public OpExprGradFunction<TriuInterpState> {
+class Triu : public OpExprGradFunction<TriuCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(TriuInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+  Maybe<void> Capture(TriuCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
                       const AttrMap& attrs) const override;
-  Maybe<void> Apply(const TriuInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const TriuCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -44,7 +44,7 @@ Maybe<void> Triu::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Triu::Capture(TriuInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Triu::Capture(TriuCaptureState* ctx, const TensorTuple& inputs,
                           const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -53,7 +53,7 @@ Maybe<void> Triu::Capture(TriuInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Triu::Apply(const TriuInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Triu::Apply(const TriuCaptureState* ctx, const TensorTuple& out_grads,
                         TensorTuple* in_grads) const {
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
   in_grads->resize(1);
diff --git a/oneflow/core/autograd/gradient_funcs/unsqueeze.cpp b/oneflow/core/autograd/gradient_funcs/unsqueeze.cpp
index 3faf92e3d5c..7246b3c1d5e 100644
--- a/oneflow/core/autograd/gradient_funcs/unsqueeze.cpp
+++ b/oneflow/core/autograd/gradient_funcs/unsqueeze.cpp
@@ -22,16 +22,16 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct UnsqueezeInterpState : public OpExprInterpState {
+struct UnsqueezeCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class Unsqueeze : public OpExprGradFunction<UnsqueezeInterpState> {
+class Unsqueeze : public OpExprGradFunction<UnsqueezeCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(UnsqueezeInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(UnsqueezeCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const UnsqueezeInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const UnsqueezeCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -48,7 +48,7 @@ Maybe<void> Unsqueeze::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Unsqueeze::Capture(UnsqueezeInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Unsqueeze::Capture(UnsqueezeCaptureState* ctx, const TensorTuple& inputs,
                                const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -57,7 +57,7 @@ Maybe<void> Unsqueeze::Capture(UnsqueezeInterpState* ctx, const TensorTuple& inp
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Unsqueeze::Apply(const UnsqueezeInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Unsqueeze::Apply(const UnsqueezeCaptureState* ctx, const TensorTuple& out_grads,
                              TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/upsample.cpp b/oneflow/core/autograd/gradient_funcs/upsample.cpp
index 51ab7d58937..c6eeaca4f78 100644
--- a/oneflow/core/autograd/gradient_funcs/upsample.cpp
+++ b/oneflow/core/autograd/gradient_funcs/upsample.cpp
@@ -23,7 +23,7 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct UpsampleInterpState : public OpExprInterpState {
+struct UpsampleCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   float height_scale;
   float width_scale;
@@ -32,12 +32,12 @@ struct UpsampleInterpState : public OpExprInterpState {
   std::string interpolation;
 };
 
-class Upsample : public OpExprGradFunction<UpsampleInterpState> {
+class Upsample : public OpExprGradFunction<UpsampleCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(UpsampleInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(UpsampleCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
-  Maybe<void> Apply(const UpsampleInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const UpsampleCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
@@ -61,7 +61,7 @@ Maybe<void> Upsample::Init(const OpExpr& op) {
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Upsample::Capture(UpsampleInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Upsample::Capture(UpsampleCaptureState* ctx, const TensorTuple& inputs,
                               const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad = inputs.at(0)->requires_grad();
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -75,7 +75,7 @@ Maybe<void> Upsample::Capture(UpsampleInterpState* ctx, const TensorTuple& input
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Upsample::Apply(const UpsampleInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Upsample::Apply(const UpsampleCaptureState* ctx, const TensorTuple& out_grads,
                             TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
@@ -94,18 +94,18 @@ Maybe<void> Upsample::Apply(const UpsampleInterpState* ctx, const TensorTuple& o
 
 REGISTER_OP_EXPR_GRAD_FUNCTION("upsample", Upsample);
 
-struct UpsampleNearest2DInterpState : public OpExprInterpState {
+struct UpsampleNearest2DCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   float height_scale;
   float width_scale;
   std::string data_format;
 };
 
-class UpsampleNearest2D : public OpExprGradFunction<UpsampleNearest2DInterpState> {
+class UpsampleNearest2D : public OpExprGradFunction<UpsampleNearest2DCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(UpsampleNearest2DInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(UpsampleNearest2DCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -119,7 +119,7 @@ class UpsampleNearest2D : public OpExprGradFunction<UpsampleNearest2DInterpState
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const UpsampleNearest2DInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const UpsampleNearest2DCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
@@ -138,7 +138,7 @@ class UpsampleNearest2D : public OpExprGradFunction<UpsampleNearest2DInterpState
 
 REGISTER_OP_EXPR_GRAD_FUNCTION("upsample_nearest_2d", UpsampleNearest2D);
 
-struct UpsampleBilinear2DInterpState : public OpExprInterpState {
+struct UpsampleBilinear2DCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   float height_scale;
   float width_scale;
@@ -146,11 +146,11 @@ struct UpsampleBilinear2DInterpState : public OpExprInterpState {
   std::string data_format;
 };
 
-class UpsampleBilinear2D : public OpExprGradFunction<UpsampleBilinear2DInterpState> {
+class UpsampleBilinear2D : public OpExprGradFunction<UpsampleBilinear2DCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(UpsampleBilinear2DInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(UpsampleBilinear2DCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -165,7 +165,7 @@ class UpsampleBilinear2D : public OpExprGradFunction<UpsampleBilinear2DInterpSta
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const UpsampleBilinear2DInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const UpsampleBilinear2DCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
@@ -185,18 +185,18 @@ class UpsampleBilinear2D : public OpExprGradFunction<UpsampleBilinear2DInterpSta
 
 REGISTER_OP_EXPR_GRAD_FUNCTION("upsample_bilinear_2d", UpsampleBilinear2D);
 
-struct UpsampleLinear1DInterpState : public OpExprInterpState {
+struct UpsampleLinear1DCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   float scale_factor;
   bool align_corners;
   std::string data_format;
 };
 
-class UpsampleLinear1D : public OpExprGradFunction<UpsampleLinear1DInterpState> {
+class UpsampleLinear1D : public OpExprGradFunction<UpsampleLinear1DCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(UpsampleLinear1DInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(UpsampleLinear1DCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -210,7 +210,7 @@ class UpsampleLinear1D : public OpExprGradFunction<UpsampleLinear1DInterpState>
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const UpsampleLinear1DInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const UpsampleLinear1DCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
@@ -229,17 +229,17 @@ class UpsampleLinear1D : public OpExprGradFunction<UpsampleLinear1DInterpState>
 
 REGISTER_OP_EXPR_GRAD_FUNCTION("upsample_linear_1d", UpsampleLinear1D);
 
-struct UpsampleNearest1DInterpState : public OpExprInterpState {
+struct UpsampleNearest1DCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   float scale_factor;
   std::string data_format;
 };
 
-class UpsampleNearest1D : public OpExprGradFunction<UpsampleNearest1DInterpState> {
+class UpsampleNearest1D : public OpExprGradFunction<UpsampleNearest1DCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(UpsampleNearest1DInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(UpsampleNearest1DCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -252,7 +252,7 @@ class UpsampleNearest1D : public OpExprGradFunction<UpsampleNearest1DInterpState
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const UpsampleNearest1DInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const UpsampleNearest1DCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
@@ -271,7 +271,7 @@ class UpsampleNearest1D : public OpExprGradFunction<UpsampleNearest1DInterpState
 
 REGISTER_OP_EXPR_GRAD_FUNCTION("upsample_nearest_1d", UpsampleNearest1D);
 
-struct UpsampleBicubic2DInterpState : public OpExprInterpState {
+struct UpsampleBicubic2DCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   float height_scale;
   float width_scale;
@@ -279,11 +279,11 @@ struct UpsampleBicubic2DInterpState : public OpExprInterpState {
   std::string data_format;
 };
 
-class UpsampleBicubic2D : public OpExprGradFunction<UpsampleBicubic2DInterpState> {
+class UpsampleBicubic2D : public OpExprGradFunction<UpsampleBicubic2DCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(UpsampleBicubic2DInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(UpsampleBicubic2DCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -298,7 +298,7 @@ class UpsampleBicubic2D : public OpExprGradFunction<UpsampleBicubic2DInterpState
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const UpsampleBicubic2DInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const UpsampleBicubic2DCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
@@ -317,7 +317,7 @@ class UpsampleBicubic2D : public OpExprGradFunction<UpsampleBicubic2DInterpState
 
 REGISTER_OP_EXPR_GRAD_FUNCTION("upsample_bicubic_2d", UpsampleBicubic2D);
 
-struct UpsampleNearest3DInterpState : public OpExprInterpState {
+struct UpsampleNearest3DCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   float depth_scale;
   float height_scale;
@@ -325,11 +325,11 @@ struct UpsampleNearest3DInterpState : public OpExprInterpState {
   std::string data_format;
 };
 
-class UpsampleNearest3D : public OpExprGradFunction<UpsampleNearest3DInterpState> {
+class UpsampleNearest3D : public OpExprGradFunction<UpsampleNearest3DCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(UpsampleNearest3DInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(UpsampleNearest3DCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -344,7 +344,7 @@ class UpsampleNearest3D : public OpExprGradFunction<UpsampleNearest3DInterpState
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const UpsampleNearest3DInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const UpsampleNearest3DCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
@@ -364,7 +364,7 @@ class UpsampleNearest3D : public OpExprGradFunction<UpsampleNearest3DInterpState
 
 REGISTER_OP_EXPR_GRAD_FUNCTION("upsample_nearest_3d", UpsampleNearest3D);
 
-struct UpsampleTrilinear3DInterpState : public OpExprInterpState {
+struct UpsampleTrilinear3DCaptureState : public AutoGradCaptureState {
   bool requires_grad;
   float depth_scale;
   float height_scale;
@@ -373,11 +373,11 @@ struct UpsampleTrilinear3DInterpState : public OpExprInterpState {
   std::string data_format;
 };
 
-class UpsampleTrilinear3D : public OpExprGradFunction<UpsampleTrilinear3DInterpState> {
+class UpsampleTrilinear3D : public OpExprGradFunction<UpsampleTrilinear3DCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(UpsampleTrilinear3DInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(UpsampleTrilinear3DCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     CHECK_EQ_OR_RETURN(inputs.size(), 1);
     CHECK_EQ_OR_RETURN(outputs.size(), 1);
@@ -393,7 +393,7 @@ class UpsampleTrilinear3D : public OpExprGradFunction<UpsampleTrilinear3DInterpS
     return Maybe<void>::Ok();
   }
 
-  Maybe<void> Apply(const UpsampleTrilinear3DInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const UpsampleTrilinear3DCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/autograd/gradient_funcs/where.cpp b/oneflow/core/autograd/gradient_funcs/where.cpp
index f31b340d58d..2e2f0f8e0a9 100644
--- a/oneflow/core/autograd/gradient_funcs/where.cpp
+++ b/oneflow/core/autograd/gradient_funcs/where.cpp
@@ -21,27 +21,27 @@ limitations under the License.
 namespace oneflow {
 namespace one {
 
-struct WhereInterpState : public OpExprInterpState {
+struct WhereCaptureState : public AutoGradCaptureState {
   bool requires_grad_x;
   bool requires_grad_y;
 };
 
-struct WhereScalarInterpState : public OpExprInterpState {
+struct WhereScalarCaptureState : public AutoGradCaptureState {
   bool requires_grad;
 };
 
-class Where : public OpExprGradFunction<WhereInterpState> {
+class Where : public OpExprGradFunction<WhereCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override;
-  Maybe<void> Capture(WhereInterpState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+  Maybe<void> Capture(WhereCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
                       const AttrMap& attrs) const override;
-  Maybe<void> Apply(const WhereInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const WhereCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 };
 
 Maybe<void> Where::Init(const OpExpr& op) { return Maybe<void>::Ok(); }
 
-Maybe<void> Where::Capture(WhereInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> Where::Capture(WhereCaptureState* ctx, const TensorTuple& inputs,
                            const TensorTuple& outputs, const AttrMap& attrs) const {
   ctx->requires_grad_x = inputs.at(1)->requires_grad();
   ctx->requires_grad_y = inputs.at(2)->requires_grad();
@@ -53,7 +53,7 @@ Maybe<void> Where::Capture(WhereInterpState* ctx, const TensorTuple& inputs,
   return Maybe<void>::Ok();
 }
 
-Maybe<void> Where::Apply(const WhereInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> Where::Apply(const WhereCaptureState* ctx, const TensorTuple& out_grads,
                          TensorTuple* in_grads) const {
   if ((!ctx->requires_grad_x) && (!ctx->requires_grad_y)) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
@@ -74,10 +74,10 @@ Maybe<void> Where::Apply(const WhereInterpState* ctx, const TensorTuple& out_gra
   return Maybe<void>::Ok();
 }
 
-class WhereScalar : public OpExprGradFunction<WhereScalarInterpState> {
+class WhereScalar : public OpExprGradFunction<WhereScalarCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
-  Maybe<void> Capture(WhereScalarInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> Capture(WhereScalarCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
     ctx->requires_grad = inputs.at(1)->requires_grad();
     if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
@@ -90,7 +90,7 @@ class WhereScalar : public OpExprGradFunction<WhereScalarInterpState> {
 
 class WhereScalarX : public WhereScalar {
  public:
-  Maybe<void> Apply(const WhereScalarInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const WhereScalarCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
@@ -107,7 +107,7 @@ class WhereScalarX : public WhereScalar {
 
 class WhereScalarY : public WhereScalar {
  public:
-  Maybe<void> Apply(const WhereScalarInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> Apply(const WhereScalarCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
     CHECK_EQ_OR_RETURN(out_grads.size(), 1);
diff --git a/oneflow/core/framework/op_expr.h b/oneflow/core/framework/op_expr.h
index 62a982b92f6..8e7d139beda 100644
--- a/oneflow/core/framework/op_expr.h
+++ b/oneflow/core/framework/op_expr.h
@@ -258,11 +258,11 @@ class SelectFirstOpExpr final : public OpExpr {
   mutable std::shared_ptr<OpExprGradFunctionIf> op_grad_func_;
 };
 
-class OpExprInterpState;
+class AutoGradCaptureState;
 // TODO(): Finish the class definition of `FunctionOpExpr`.
 class FunctionOpExpr : public OpExpr {
  public:
-  using FType = std::function<Maybe<void>(const std::shared_ptr<OpExprInterpState>& /*ctx*/,
+  using FType = std::function<Maybe<void>(const std::shared_ptr<AutoGradCaptureState>& /*ctx*/,
                                           const TensorTuple& /*inputs or out_grads*/,
                                           TensorTuple* /*outputs or in_grads*/)>;
 
@@ -287,8 +287,8 @@ class FunctionOpExpr : public OpExpr {
   FType forward() const { return forward_; }
   FType backward() const { return backward_; }
 
-  std::shared_ptr<const OpExprInterpState> state() const { return state_; }
-  std::shared_ptr<OpExprInterpState> mutable_state() { return state_; }
+  std::shared_ptr<const AutoGradCaptureState> state() const { return state_; }
+  std::shared_ptr<AutoGradCaptureState> mutable_state() { return state_; }
 
   Maybe<bool> IsGradDisabled() const override { return false; }
   Maybe<OpExprGradClosure> GetOrCreateOpGradClosure() const override { OF_UNIMPLEMENTED(); }
@@ -296,7 +296,7 @@ class FunctionOpExpr : public OpExpr {
  private:
   FType forward_;
   FType backward_;
-  std::shared_ptr<OpExprInterpState> state_;
+  std::shared_ptr<AutoGradCaptureState> state_;
 };
 
 }  // namespace one
diff --git a/oneflow/core/framework/op_expr_grad_function.h b/oneflow/core/framework/op_expr_grad_function.h
index 70caf033641..b2e381f8409 100644
--- a/oneflow/core/framework/op_expr_grad_function.h
+++ b/oneflow/core/framework/op_expr_grad_function.h
@@ -18,40 +18,57 @@ limitations under the License.
 #define ONEFLOW_CORE_FRAMEWORK_OP_EXPR_GRAD_FUNCTION_H_
 
 #include "oneflow/core/common/auto_registration_factory.h"
-#include "oneflow/core/framework/op_interpreter.h"  // OpExprInterpState
+#include "oneflow/core/framework/op_interpreter.h"
 
 namespace oneflow {
 namespace one {
 
 static constexpr char kGradientOpSuffix[] = ".grad";
 
+class AutoGradCaptureState {
+ public:
+  AutoGradCaptureState() = default;
+  virtual ~AutoGradCaptureState() = default;
+
+  const TensorTuple& SavedTensors() const { return saved_tensors_; }
+
+  size_t SaveTensorForBackward(const std::shared_ptr<Tensor>& tensor) {
+    size_t offset = saved_tensors_.size();
+    saved_tensors_.push_back(tensor);
+    return offset;
+  }
+
+ private:
+  TensorTuple saved_tensors_;
+};
+
 // Stateless container base of the backward op exprs.
 // The backward op exprs should be contained in the derived class.
 class OpExprGradFunctionIf {
  public:
   virtual ~OpExprGradFunctionIf() = default;
 
-  virtual std::shared_ptr<OpExprInterpState> MakeCustomState() const = 0;
+  virtual std::shared_ptr<AutoGradCaptureState> MakeCustomState() const = 0;
 
   virtual Maybe<void> Init(const OpExpr& op) = 0;
 
   // Capture forward inputs and outputs for backward.
-  virtual Maybe<void> CaptureIf(OpExprInterpState* ctx, const TensorTuple& inputs,
+  virtual Maybe<void> CaptureIf(AutoGradCaptureState* ctx, const TensorTuple& inputs,
                                 const TensorTuple& outputs,
                                 const OpExprInterpContext& interp_ctx) const = 0;
 
-  virtual Maybe<void> ApplyIf(const OpExprInterpState* ctx, const TensorTuple& out_grads,
+  virtual Maybe<void> ApplyIf(const AutoGradCaptureState* ctx, const TensorTuple& out_grads,
                               TensorTuple* in_grads) const = 0;
 };
 
 template<typename StateT>
 class OpExprGradFunction : public OpExprGradFunctionIf {
  public:
-  std::shared_ptr<OpExprInterpState> MakeCustomState() const override {
+  std::shared_ptr<AutoGradCaptureState> MakeCustomState() const override {
     return std::make_shared<StateT>();
   }
 
-  Maybe<void> CaptureIf(OpExprInterpState* ctx, const TensorTuple& inputs,
+  Maybe<void> CaptureIf(AutoGradCaptureState* ctx, const TensorTuple& inputs,
                         const TensorTuple& outputs,
                         const OpExprInterpContext& interp_ctx) const override {
     StateT* state = dynamic_cast<StateT*>(ctx);
@@ -71,7 +88,7 @@ class OpExprGradFunction : public OpExprGradFunctionIf {
     return Capture(state, detach_inputs, detach_outputs, interp_ctx);
   }
 
-  Maybe<void> ApplyIf(const OpExprInterpState* ctx, const TensorTuple& out_grads,
+  Maybe<void> ApplyIf(const AutoGradCaptureState* ctx, const TensorTuple& out_grads,
                       TensorTuple* in_grads) const override {
     const StateT* state = dynamic_cast<const StateT*>(ctx);
     CHECK_NOTNULL_OR_RETURN(state);
@@ -104,7 +121,7 @@ class OpExprGradClosure {
   explicit OpExprGradClosure(const std::shared_ptr<OpExprGradFunctionIf>& impl)
       : impl_(impl), state_(impl->MakeCustomState()) {}
   explicit OpExprGradClosure(const std::shared_ptr<OpExprGradFunctionIf>& impl,
-                             const std::shared_ptr<OpExprInterpState>& state)
+                             const std::shared_ptr<AutoGradCaptureState>& state)
       : impl_(impl), state_(state) {}
 
   virtual ~OpExprGradClosure() = default;
@@ -120,7 +137,7 @@ class OpExprGradClosure {
 
  private:
   std::shared_ptr<OpExprGradFunctionIf> impl_;
-  std::shared_ptr<OpExprInterpState> state_;
+  std::shared_ptr<AutoGradCaptureState> state_;
 };
 
 #define REGISTER_OP_EXPR_GRAD_FUNCTION(op_type, op_grad) \
diff --git a/oneflow/core/framework/op_interpreter.h b/oneflow/core/framework/op_interpreter.h
index b44129a2b6b..71ea6b844d4 100644
--- a/oneflow/core/framework/op_interpreter.h
+++ b/oneflow/core/framework/op_interpreter.h
@@ -33,23 +33,6 @@ class NdSbp;
 
 namespace one {
 
-class OpExprInterpState {
- public:
-  OpExprInterpState() = default;
-  virtual ~OpExprInterpState() = default;
-
-  const TensorTuple& SavedTensors() const { return saved_tensors_; }
-
-  size_t SaveTensorForBackward(const std::shared_ptr<Tensor>& tensor) {
-    size_t offset = saved_tensors_.size();
-    saved_tensors_.push_back(tensor);
-    return offset;
-  }
-
- private:
-  TensorTuple saved_tensors_;
-};
-
 struct OpExprInterpContext {
   OpExprInterpContext(const AttrMap& attrs_arg) : attrs(attrs_arg) {}
   OpExprInterpContext(const AttrMap& attrs_arg, Symbol<Device> device_arg)