计算图怎么记下来

pub struct Tensor_ {
    id: TensorId,
    // ...
    op: BackpropOp,
    is_variable: bool,
    // ...
}

/// `BackpropOp` is a wrapper around `Option<Op>`. The main goal is to ensure that dependencies are
/// properly checked when creating a new value
#[derive(Clone)]
pub struct BackpropOp(Option<Op>);

#[derive(Clone)]
pub enum Op {
    Binary(Tensor, Tensor, BinaryOp),
    Unary(Tensor, UnaryOp),
    Cmp(Tensor, CmpOp),
    // The third argument is the reduced shape with `keepdim=true`.
    Reduce(Tensor, ReduceOp, Vec<usize>),
    Matmul(Tensor, Tensor),
    // ...
    Broadcast(Tensor),
    Narrow(Tensor, usize, usize, usize),
    Reshape(Tensor),
    Transpose(Tensor, usize, usize),
    // ...
    CustomOp1(
        Tensor,
        std::sync::Arc<Box<dyn crate::CustomOp1 + Send + Sync>>,
    ),
    // ...
}

            let storage = self.storage().binary_impl::<crate::op::$op_name>(
                &*rhs.storage(),
                self.layout(),
                rhs.layout(),
            )?;
            let op = BackpropOp::new2(self, rhs, |t1, t2| Op::Binary(t1, t2, BinaryOp::$op_name));
            Ok(from_storage(storage, shape.clone(), op, false))

    pub(crate) fn new2(arg1: &Tensor, arg2: &Tensor, f: impl Fn(Tensor, Tensor) -> Op) -> Self {
        let op = if arg1.track_op() || arg2.track_op() {
            Some(f(arg1.clone(), arg2.clone()))
        } else {
            None
        };
        Self(op)
    }

    pub fn track_op(&self) -> bool {
        self.is_variable || self.op.is_some()
    }

let x = Tensor::new(&[[1f32, 2.], [3., 4.]], &dev)?;

// w 是普通张量
let w = Tensor::new(&[[0.5f32, -1.], [2., 0.]], &dev)?;
let y = x.matmul(&w)?;
let loss = y.sqr()?.sum_all()?;

// w 是 Var
let w = Var::new(&[[0.5f32, -1.], [2., 0.]], &dev)?;
let y = x.matmul(w.as_tensor())?;
let loss = y.sqr()?.sum_all()?;

-- w is a plain Tensor
w                      id=TensorId(2) is_variable=false track_op=false sorted_nodes=0
y = x.matmul(w)        id=TensorId(3) is_variable=false track_op=false sorted_nodes=0
loss                   id=TensorId(6) is_variable=false track_op=false sorted_nodes=0
-- w is a Var
w                      id=TensorId(7) is_variable=true  track_op=true  sorted_nodes=1
x                      id=TensorId(1) is_variable=false track_op=false sorted_nodes=0
y = x.matmul(w)        id=TensorId(8) is_variable=false track_op=true  sorted_nodes=2

// Variables are wrappers around tensors that can be modified, they are typically used for holding
// weights and being modified by gradient descent.
// We do not expose a public way to create variables as this would break the invariant that the
// tensor within a variable is actually with `is_variable` set to `true`.
use crate::{DType, Device, Error, Result, Shape, Tensor};

/// A variable is a wrapper around a tensor, however variables can have their content modified
/// whereas tensors are immutable.
#[derive(Clone, Debug)]
pub struct Var(Tensor);

    pub fn set(&self, src: &Tensor) -> Result<()> {
        if self.same_storage(src) {
            let msg = "cannot set a variable to a tensor that is derived from its value";
            Err(Error::CannotSetVar { msg }.bt())?
        }
        let (mut dst, layout) = self.storage_mut_and_layout();
        if !layout.is_contiguous() {
            let msg = "cannot set a non-contiguous variable";
            Err(Error::CannotSetVar { msg }.bt())?
        }
        // ...
        src.copy_strided_src(&mut dst, layout.start_offset(), src_l)?;
        Ok(())
    }

set(w.t()) -> cannot set variable cannot set a variable to a tensor that is derived from its value
set(w.as_detached_tensor()) -> cannot set variable cannot set a variable to a tensor that is derived from its value
after set(w * 2): [[2.0, 4.0], [6.0, 8.0]]

    pub fn detach(&self) -> Tensor {
        if self.op.is_none() && !self.is_variable {
            self.clone()
        } else {
            let tensor_ = Tensor_ {
                id: TensorId::new(),
                storage: self.storage.clone(),
                layout: self.layout.clone(),
                op: BackpropOp::none(),
                is_variable: false,
                dtype: self.dtype,
                device: self.device.clone(),
            };
            Tensor(Arc::new(tensor_))
        }
    }

-- detach / argmax / views
y.detach()             id=TensorId(12) is_variable=false track_op=false sorted_nodes=0
y.detach().exp()       id=TensorId(13) is_variable=false track_op=false sorted_nodes=0
y.argmax(1)            id=TensorId(15) is_variable=false track_op=false sorted_nodes=0
w.t()                  id=TensorId(16) is_variable=false track_op=true  sorted_nodes=2
y.to_dtype(U32)        id=TensorId(17) is_variable=false track_op=true  sorted_nodes=0
w.as_detached_tensor() id=TensorId(18) is_variable=false track_op=false sorted_nodes=0

let big = Var::zeros((1024, 1024), DType::F32, &dev)?;
let h = big.as_tensor().exp()?; // 4 MiB intermediate
let out = h.sum_all()?;
let h_id = h.id();
drop(h);
let still_there = out.sorted_nodes().iter().any(|n| n.id() == h_id);

dropped h, still reachable from out: true

loss                   id=TensorId(11) is_variable=false track_op=true  sorted_nodes=5
  sorted[0] id=TensorId(11) dims=[] is_variable=false
  sorted[1] id=TensorId(10) dims=[1, 1] is_variable=false
  sorted[2] id=TensorId(9) dims=[2, 2] is_variable=false
  sorted[3] id=TensorId(8) dims=[2, 2] is_variable=false
  sorted[4] id=TensorId(7) dims=[2, 2] is_variable=true