一个算子的反向

pub trait Backward<B, const N: usize>: Send + core::fmt::Debug
where
    Self: Sized + 'static,
    B: Backend,
{
    /// Associated type to compute the backward pass.
    type State: Clone + Send + core::fmt::Debug + 'static;

    /// The backward pass.
    fn backward(
        self,
        ops: Ops<Self::State, N>,
        grads: &mut Gradients,
        checkpointer: &mut Checkpointer,
    );

    /// Prepare the backward ops, retaining input references through registration.
    ///
    /// Obtain guards from the input tensors before consuming their primitives.
    fn prepare<C: CheckpointStrategy>(
        self,
        nodes: [NodeGuard; N],
    ) -> OpsPrep<Self, B, Self::State, C, N> {
        let requirement =
            Requirement::from_nodes(&nodes.each_ref().map(|guard| guard.node_ref().clone()));
        OpsPrep::new(
            nodes,
            requirement,
            self,
            ComputingProperty::Ambiguous, // If not specified we start with ambiguous
            CheckpointerBuilder::default(),
        )
    }
}

    fn float_mul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
        #[derive(Debug)]
        struct Mul;

        retro_binary!(RetroMul, B::float_mul);

        impl<B: Backend> Backward<B, 2> for Mul {
            type State = (Option<NodeId>, Option<NodeId>, BinaryOpsBroadcast);

            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
                let (lhs, rhs, broadcast) = ops.state;
                let lhs = lhs.map(|lhs| checkpointer.retrieve_node_output(lhs));
                let rhs = rhs.map(|rhs| checkpointer.retrieve_node_output(rhs));

                binary::<B, _, _>(
                    ops.parents,
                    ops.node,
                    grads,
                    |grad| {
                        let grad = B::float_mul(grad, rhs.unwrap());
                        broadcast.backward_lhs::<B>(grad)
                    },
                    |grad| {
                        let grad = B::float_mul(grad, lhs.unwrap());
                        broadcast.backward_rhs::<B>(grad)
                    },
                );
            }
        }

        let lhs_tracked = lhs.is_tracked();
        let rhs_tracked = rhs.is_tracked();
        let broadcast = BinaryOpsBroadcast::new::<B>(&lhs.primitive, &rhs.primitive);

        match Mul
            .prepare::<C>([lhs.node(), rhs.node()])
            .memory_bound()
            .retro_forward(RetroMul::<B>::new(lhs.node.id, rhs.node.id))
            .parents([&lhs, &rhs])
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let lhs_state = rhs_tracked.then(|| prep.checkpoint(&lhs));
                let rhs_state = lhs_tracked.then(|| prep.checkpoint(&rhs));

                prep.finish(
                    (lhs_state, rhs_state, broadcast),
                    B::float_mul(lhs.primitive, rhs.primitive),
                )
            }
            OpsKind::UnTracked(prep) => prep.finish(B::float_mul(lhs.primitive, rhs.primitive)),
        }
    }

pub(crate) fn register_step<B: Backend, S: Step + 'static>(
    input_guards: impl IntoIterator<Item = NodeGuard>,
    output: AutodiffTensor<B>,
    step: S,
    checkpointer_builder: CheckpointerBuilder,
) -> AutodiffTensor<B> {
    let output = output.register_step(step, checkpointer_builder);
    // Registration publishes the dependency before its input references are released.
    drop(input_guards);
    output
}

        // We register the ops in the graph even if untracked, otherwise memory bound operations
        // that have an untracked parent would not be able to retrieve it

pub fn binary<B, FLhs, FRhs>(
    parents: [Option<NodeRef>; 2],
    node: NodeRef,
    grads: &mut Gradients,
    func_lhs: FLhs,
    func_rhs: FRhs,
) where
    B: Backend,
    FLhs: FnOnce(B::FloatTensorPrimitive) -> B::FloatTensorPrimitive,
    FRhs: FnOnce(B::FloatTensorPrimitive) -> B::FloatTensorPrimitive,
{
    let [grad_4lhs, grad_4rhs] = duplicate(&parents, Some(grads.consume::<B>(&node)));
    let [node_lhs, node_rhs] = parents;

    if let Some(node) = node_lhs {
        let grad = func_lhs(grad_4lhs.unwrap());
        grads.register::<B>(node.id, grad)
    }

    if let Some(node) = node_rhs {
        let grad = func_rhs(grad_4rhs.unwrap());
        grads.register::<B>(node.id, grad)
    }
}

        impl<B: Backend> Backward<B, 2> for Matmul {
            type State = (Option<NodeId>, Option<NodeId>, BinaryOpsBroadcast);

            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
                let (lhs, rhs, broadcast) = ops.state;
                let lhs = lhs.map(|lhs| checkpointer.retrieve_node_output(lhs));
                let rhs = rhs.map(|rhs| checkpointer.retrieve_node_output(rhs));

                binary::<B, _, _>(
                    ops.parents,
                    ops.node,
                    grads,
                    |grad| {
                        let rhs = B::float_transpose(rhs.unwrap());
                        let grad = B::float_matmul(grad, rhs);

                        broadcast.backward_lhs::<B>(grad)
                    },
                    |grad| {
                        let lhs = B::float_transpose(lhs.unwrap());
                        let grad = B::float_matmul(lhs, grad);

                        broadcast.backward_rhs::<B>(grad)
                    },
                );
            }
        }

        // ...
        match Matmul
            .prepare::<C>([lhs.node(), rhs.node()])
            .compute_bound()
            .stateful()
        {

pub fn broadcast_shape<B: Backend>(grad: FloatTensor<B>, shape: &Shape) -> FloatTensor<B> {
    let shape_grad = grad.shape();
    let ndims = shape_grad.num_dims();

    let mut broadcast_dims = Vec::new();
    for i in 0..ndims {
        if shape_grad[i] != shape[i] {
            if shape[i] != 1 {
                panic!(
                    "Invalid broadcast shapes: Next grad shape {:?}, Previous grad shape {:?}. {}",
                    shape, shape_grad, "Expected the shape of the next grad to be 1."
                );
            }
            broadcast_dims.push(i);
        }
    }

    // One reduction over every broadcast dimension rather than one per
    // dimension: a parameter broadcast over batch and space would otherwise
    // pay a launch and an intermediate for each, the first of them nearly the
    // size of the gradient itself.
    match broadcast_dims.as_slice() {
        [] => grad,
        [dim] => B::float_sum_dim(grad, *dim),
        dims => B::float_sum_dims(grad, dims),
    }
}

fn check(name: &str, f: F, inputs: &[(Vec<f32>, [usize; 3])]) {
  let ad = Device::flex().autodiff();
  let plain = Device::flex();

  let leaves: Vec<Tensor<3>> = inputs
    .iter()
    .map(|(v, s)| make(v, *s, &ad).require_grad())
    .collect();
  let grads = f(&leaves).backward();

  for (i, (v, s)) in inputs.iter().enumerate() {
    let g = leaves[i].grad(&grads).unwrap();
    // ... 对第 i 个输入的每个元素 k：xs[i][k] ± eps，在 plain 设备上求 f，取中心差分
  }
}

fn main() {
  check(
    "mul broadcast [1,1,3]",
    |t| {
      let w = make(&[0.3, -1.2, 0.7, 2.0, -0.5, 1.1], [1, 2, 3], &t[0].device());
      (t[0].clone() * t[1].clone() * w).sum()
    },
    &[(seq(6, 0.5, -1.0), [1, 2, 3]), (vec![2.0, -1.0, 0.5], [1, 1, 3])],
  );
  // ... "mul"、"matmul"、"matmul batch-broadcast"（[2,2,3] 乘 [1,3,2]）同理
}

mul                    input0 shape=[1, 2, 3] grad shape=[1, 2, 3] max|analytic-numeric|=1.46e-6
                       grad=[0.3, -0.90000004, 0.35, 0.5, -0.0, -0.275]
mul                    input1 shape=[1, 2, 3] grad shape=[1, 2, 3] max|analytic-numeric|=1.79e-6
                       grad=[-0.3, 0.6, 0.0, 1.0, -0.5, 1.6500001]
mul broadcast [1,1,3]  input0 shape=[1, 2, 3] grad shape=[1, 2, 3] max|analytic-numeric|=1.14e-5
                       grad=[0.6, 1.2, 0.35, 4.0, 0.5, 0.55]
mul broadcast [1,1,3]  input1 shape=[1, 1, 3] grad shape=[1, 1, 3] max|analytic-numeric|=7.27e-6
                       grad=[0.7, 0.100000024, 1.6500001]
matmul                 input0 shape=[1, 2, 3] grad shape=[1, 2, 3] max|analytic-numeric|=4.53e-6
                       grad=[-0.6, -0.15, 0.3, 2.2, 0.85, -0.5]
matmul                 input1 shape=[1, 3, 2] grad shape=[1, 3, 2] max|analytic-numeric|=4.53e-6
                       grad=[0.049999982, 2.2, 0.54999995, 2.6, 1.05, 3.0]
matmul batch-broadcast input0 shape=[2, 2, 3] grad shape=[2, 2, 3] max|analytic-numeric|=4.35e-5
                       grad=[-0.6, -0.15, 0.3, 2.2, 0.85, -0.5, 1.375, 0.625, -0.125, 0.625, 0.125, -0.375]
matmul batch-broadcast input1 shape=[1, 3, 2] grad shape=[1, 3, 2] max|analytic-numeric|=3.15e-5
                       grad=[0.29999998, 8.45, 1.05, 9.85, 1.8, 11.25]