1use fandhe_ai_tensor_core::{ScalarBinaryOp, ScalarUnaryOp, Tensor, batch_norm_layout};
120
121use crate::backward::Gradients;
122use crate::error::AutodiffError;
123use crate::tape::{NodeId, Op, Tape, TapeId, TapeNode, materialize_fallible};
124use crate::var::Var;
125
126mod conv;
127mod index_cum;
128mod linalg_fft;
129mod pad_mask;
130mod pool;
131
132pub(crate) fn scalar_unary_replayable(op: ScalarUnaryOp) -> bool {
146 matches!(
147 op,
148 ScalarUnaryOp::Neg
149 | ScalarUnaryOp::Abs
150 | ScalarUnaryOp::Sqrt
151 | ScalarUnaryOp::Log
152 | ScalarUnaryOp::Log2
153 | ScalarUnaryOp::Log10
154 | ScalarUnaryOp::Sin
155 | ScalarUnaryOp::Cos
156 | ScalarUnaryOp::Tan
157 | ScalarUnaryOp::Relu
158 | ScalarUnaryOp::Exp
159 | ScalarUnaryOp::Tanh
160 | ScalarUnaryOp::Sigmoid
161 | ScalarUnaryOp::Silu
162 | ScalarUnaryOp::Hardswish
163 | ScalarUnaryOp::LeakyRelu { .. }
164 | ScalarUnaryOp::Elu { .. }
165 | ScalarUnaryOp::Softplus { .. }
166 | ScalarUnaryOp::Clamp { .. }
167 | ScalarUnaryOp::PowScalar { .. }
168 | ScalarUnaryOp::Gelu
169 | ScalarUnaryOp::GeluTanh
170 | ScalarUnaryOp::Erf
171 )
172}
173
174pub(crate) fn scalar_binary_replayable(op: ScalarBinaryOp) -> bool {
181 matches!(
182 op,
183 ScalarBinaryOp::Add
184 | ScalarBinaryOp::Sub
185 | ScalarBinaryOp::Mul
186 | ScalarBinaryOp::Div
187 | ScalarBinaryOp::Pow
188 | ScalarBinaryOp::Maximum
189 | ScalarBinaryOp::Minimum
190 | ScalarBinaryOp::Gt
191 | ScalarBinaryOp::Ge
192 | ScalarBinaryOp::Lt
193 | ScalarBinaryOp::Le
194 | ScalarBinaryOp::Eq
195 | ScalarBinaryOp::Ne
196 )
197}
198
199#[derive(Debug)]
205pub struct CreateGraphResult<'c> {
206 first_order: Gradients,
207 parent_id: TapeId,
208 parent_epoch: u64,
209 mirror: Vec<Option<Var<'c>>>,
212 cgrads: Vec<Option<Var<'c>>>,
215}
216
217impl<'c> CreateGraphResult<'c> {
218 pub fn first_order(&self) -> &Gradients {
224 &self.first_order
225 }
226
227 pub fn grad(&self, parent_var: &Var<'_>) -> Result<Option<Var<'c>>, AutodiffError> {
241 self.check(parent_var)?;
242 Ok(self.cgrads.get(parent_var.node_id().0).copied().flatten())
243 }
244
245 pub fn child_var(&self, parent_var: &Var<'_>) -> Result<Option<Var<'c>>, AutodiffError> {
252 self.check(parent_var)?;
253 Ok(self.mirror.get(parent_var.node_id().0).copied().flatten())
254 }
255
256 pub(crate) fn mirror(&self) -> &[Option<Var<'c>>] {
260 &self.mirror
261 }
262
263 fn check(&self, parent_var: &Var<'_>) -> Result<(), AutodiffError> {
264 if parent_var.tape_id() != self.parent_id || parent_var.tape_epoch() != self.parent_epoch {
265 return Err(AutodiffError::TapeMismatch);
266 }
267 if !parent_var.requires_grad() {
268 return Err(AutodiffError::GradientTrackingDisabled);
269 }
270 Ok(())
271 }
272}
273
274impl Tape {
275 pub fn backward_create_graph<'c>(
311 &self,
312 loss: &Var<'_>,
313 child: &'c Tape,
314 ) -> Result<CreateGraphResult<'c>, AutodiffError> {
315 if loss.tape_id() != self.id {
316 return Err(AutodiffError::TapeMismatch);
317 }
318 if child.id == self.id {
319 return Err(AutodiffError::Backward(
320 "create_graph: 子テープに親と同一の Tape は指定できない".into(),
321 ));
322 }
323 if child.device() != self.device() {
324 return Err(AutodiffError::DeviceMismatch {
325 requested: child.device(),
326 actual: self.device(),
327 });
328 }
329 if !child.is_empty() {
330 return Err(AutodiffError::Backward(
331 "create_graph: 子テープは空である必要がある(既存ノードを持つ Tape は渡せない)"
332 .into(),
333 ));
334 }
335 if self.has_registered_checkpoints() {
336 return Err(AutodiffError::Backward(
337 "create_graph: checkpoint 区間が登録済みの親テープでは create_graph は未対応(\
338 docs/autodiff-higher-order-grad-decision.md §8 参照)"
339 .into(),
340 ));
341 }
342
343 let ancestors = collect_ancestors(self, loss.node_id());
356 validate_ancestors(self, &ancestors)?;
357
358 let first_order = self.backward(loss)?;
359
360 let n = self.nodes.borrow().len();
361 let mut mirror: Vec<Option<Var<'c>>> = vec![None; n];
362 let point = child.rollback_point();
367 let built = build_mirror(self, child, &ancestors, &mut mirror)
368 .and_then(|()| build_cgrads(self, child, &ancestors, &mirror, loss.node_id()));
369 let cgrads = match built {
370 Ok(c) => c,
371 Err(e) => {
372 mirror.clear();
373 child.rollback_to(point);
374 return Err(e);
375 }
376 };
377
378 Ok(CreateGraphResult {
379 first_order,
380 parent_id: self.id,
381 parent_epoch: self.epoch(),
382 mirror,
383 cgrads,
384 })
385 }
386}
387
388fn collect_ancestors(tape: &Tape, root: NodeId) -> Vec<NodeId> {
401 let nodes = tape.nodes.borrow();
402 let n = nodes.len();
403 let mut visited = vec![false; n];
404 let mut stack = vec![root];
405 let mut ids: Vec<NodeId> = Vec::new();
406 while let Some(id) = stack.pop() {
407 if visited[id.0] {
408 continue;
409 }
410 visited[id.0] = true;
411 ids.push(id);
412 let node = &nodes[id.0];
413 if node.requires_grad {
414 node.op.for_each_input(|input_id| stack.push(input_id));
415 }
416 }
417 ids.sort_unstable_by_key(|id| id.0);
418 ids
419}
420
421fn validate_ancestors(parent: &Tape, ancestors: &[NodeId]) -> Result<(), AutodiffError> {
445 let parent_nodes = parent.nodes.borrow();
446 for &id in ancestors {
447 let node = &parent_nodes[id.0];
448 if matches!(
449 node.op,
450 Op::ResidentLeaf { .. } | Op::LinearResident { .. } | Op::LinearAct { .. }
451 ) {
452 return Err(AutodiffError::Backward(format!(
453 "create_graph: resident/fused 経路の Op(NodeId({}))は非対応(\
454 docs/autodiff-higher-order-grad-decision.md §8 参照)",
455 id.0
456 )));
457 }
458 if !node.requires_grad || matches!(node.op, Op::Leaf) {
459 continue;
460 }
461 if node.low_precision {
471 return Err(AutodiffError::Backward(format!(
472 "create_graph: 低精度 forward のノード(NodeId({}))は非対応(\
473 テープへ精度情報を保持しないため子テープへ replay すると静かに \
474 f32 精度へフォールバックする。イシュー #2071・#2628)",
475 id.0
476 )));
477 }
478 if !node.op.supports_create_graph() {
479 return Err(AutodiffError::Backward(format!(
480 "create_graph: 未対応の Op(NodeId({}))へ到達した(\
481 docs/autodiff-higher-order-grad-decision.md §8 の対象 Op のみ再生可能)",
482 id.0
483 )));
484 }
485 if let Op::Rfft { norm, .. }
488 | Op::Irfft { norm, .. }
489 | Op::Fft { norm, .. }
490 | Op::Ifft { norm, .. } = &node.op
491 {
492 linalg_fft::check_fft_norm(*norm)?;
493 }
494 if let Op::BatchNorm { input, .. } = &node.op
498 && parent_nodes[input.0].shape.contains(&0)
501 {
502 return Err(AutodiffError::Backward(format!(
503 "create_graph: 空テンソルの BatchNorm(NodeId({}))は非対応(\
504 docs/autodiff-higher-order-grad-decision.md §28 参照)",
505 id.0
506 )));
507 }
508 if let Op::Conv2d { compute_dtype, .. } = &node.op
513 && *compute_dtype != fandhe_ai_tensor_core::ScalarDType::F32
514 {
515 return Err(AutodiffError::Backward(format!(
516 "create_graph: 低精度 conv2d({compute_dtype:?}。NodeId({}))は非対応(\
517 子テープへ精度情報を持たず静かに f32 再生される。イシュー #2071・#3080)",
518 id.0
519 )));
520 }
521 if let Op::Scatter { reduce, .. } = &node.op
524 && !matches!(
525 reduce,
526 fandhe_ai_tensor_core::ScatterReduce::Add
527 | fandhe_ai_tensor_core::ScatterReduce::Overwrite
528 )
529 {
530 return Err(AutodiffError::Backward(format!(
531 "create_graph: 未知の ScatterReduce variant {reduce:?}(NodeId({}))は非対応",
532 id.0
533 )));
534 }
535 if let Op::Embedding {
538 weight,
539 padding_idx: Some(p),
540 ..
541 } = &node.op
542 {
543 let w_shape = &parent_nodes[weight.0].shape;
544 if w_shape.len() != 2 || *p >= w_shape[0] {
545 return Err(AutodiffError::Backward(format!(
546 "create_graph: Embedding の padding_idx={p} が weight shape {w_shape:?} \
547 の範囲外(NodeId({}))",
548 id.0
549 )));
550 }
551 }
552 if let Op::MatMul(a, b) = node.op {
553 let a_rank = parent_nodes[a.0].shape.len();
554 let b_rank = parent_nodes[b.0].shape.len();
555 if a_rank != 2 || b_rank != 2 {
556 return Err(AutodiffError::Backward(format!(
557 "create_graph: rank≥3 の MatMul(NodeId({}))は非対応(\
558 子テープの matmul VJP は rank 2 × rank 2 限定。\
559 docs/autodiff-higher-order-grad-decision.md §14 参照)",
560 id.0
561 )));
562 }
563 }
564 crate::create_graph_norm::validate_norm_s1(&parent_nodes, parent.ops(), &node.op)?;
567 }
568 Ok(())
569}
570
571fn build_mirror<'c>(
589 parent: &Tape,
590 child: &'c Tape,
591 ancestors: &[NodeId],
592 mirror: &mut [Option<Var<'c>>],
593) -> Result<(), AutodiffError> {
594 let parent_nodes = parent.nodes.borrow();
595 let parent_ops = parent.ops();
596
597 for &id in ancestors {
599 let node = &parent_nodes[id.0];
600 if matches!(
601 node.op,
602 Op::ResidentLeaf { .. } | Op::LinearResident { .. } | Op::LinearAct { .. }
603 ) {
604 return Err(AutodiffError::Backward(format!(
605 "create_graph: resident/fused 経路の Op(NodeId({}))は非対応(\
606 docs/autodiff-higher-order-grad-decision.md §8 参照)",
607 id.0
608 )));
609 }
610 if !node.requires_grad {
611 let value = materialize_fallible(&parent_nodes, parent_ops, id)?.clone();
612 mirror[id.0] = Some(child.var_no_grad(&value));
613 continue;
614 }
615 if matches!(node.op, Op::Leaf) {
616 let value = materialize_fallible(&parent_nodes, parent_ops, id)?.clone();
617 mirror[id.0] = Some(child.var(&value));
618 }
619 }
620
621 for &id in ancestors {
623 if mirror[id.0].is_some() {
624 continue;
625 }
626 let node = &parent_nodes[id.0];
627 if !node.op.supports_create_graph() {
628 return Err(AutodiffError::Backward(format!(
629 "create_graph: 未対応の Op(NodeId({}))へ到達した(\
630 docs/autodiff-higher-order-grad-decision.md §8 の対象 Op のみ再生可能)",
631 id.0
632 )));
633 }
634 let replayed = replay_op(node.op.clone(), &node.shape, node.fp32_strict, mirror)?;
635 mirror[id.0] = Some(replayed);
636 }
637
638 Ok(())
639}
640
641fn get_mirror<'c>(mirror: &[Option<Var<'c>>], id: NodeId) -> Result<Var<'c>, AutodiffError> {
648 mirror[id.0].ok_or_else(|| {
649 AutodiffError::Backward(format!(
650 "create_graph: 祖先ノード(NodeId({}))の写しが未構築(内部不変条件違反)",
651 id.0
652 ))
653 })
654}
655
656fn replay_op<'c>(
677 op: Op,
678 shape: &[usize],
679 fp32_strict: bool,
680 mirror: &[Option<Var<'c>>],
681) -> Result<Var<'c>, AutodiffError> {
682 match op {
683 Op::Add(a, b) => get_mirror(mirror, a)?.add(&get_mirror(mirror, b)?),
684 Op::Mul(a, b) => get_mirror(mirror, a)?.mul(&get_mirror(mirror, b)?),
685 Op::MatMul(a, b) => {
693 let lhs = get_mirror(mirror, a)?;
694 let rhs = get_mirror(mirror, b)?;
695 if fp32_strict {
696 lhs.matmul_fp32_strict(&rhs)
697 } else {
698 lhs.matmul(&rhs)
699 }
700 }
701 Op::Relu(a) => Ok(get_mirror(mirror, a)?.relu()),
702 Op::Exp(a) => Ok(get_mirror(mirror, a)?.exp()),
703 Op::Tanh(a) => Ok(get_mirror(mirror, a)?.tanh()),
704 Op::Sigmoid(a) => Ok(get_mirror(mirror, a)?.sigmoid()),
705 Op::Sum { input, dim } => get_mirror(mirror, input)?.sum(dim),
706 Op::Mean { input, dim } => get_mirror(mirror, input)?.mean(dim),
707 Op::Reshape { input } => get_mirror(mirror, input)?.contiguous()?.reshape(shape),
708 Op::BroadcastTo { input } => get_mirror(mirror, input)?.broadcast_to(shape),
709 Op::ScalarUnary { op: sop, input } => get_mirror(mirror, input)?.scalar_unary(sop),
715 Op::ScalarBinary { op: sop, a, b } => {
716 get_mirror(mirror, a)?.scalar_binary(&get_mirror(mirror, b)?, sop)
717 }
718 Op::Transpose { input, dim0, dim1 } => get_mirror(mirror, input)?.transpose(dim0, dim1),
719 Op::Permute { input, perm } => get_mirror(mirror, input)?.permute(&perm),
720 Op::Narrow {
721 input,
722 dim,
723 start,
724 len,
725 } => get_mirror(mirror, input)?.narrow(dim, start, len),
726 Op::Concat { inputs, dim } => {
727 let mirrored: Result<Vec<Var<'c>>, AutodiffError> =
728 inputs.iter().map(|&id| get_mirror(mirror, id)).collect();
729 Var::cat(&mirrored?, dim)
730 }
731 Op::Contiguous { input } => get_mirror(mirror, input)?.contiguous(),
732 Op::Softmax { input, dim } => get_mirror(mirror, input)?.softmax(dim),
735 Op::LogSoftmax { input, dim } => get_mirror(mirror, input)?.log_softmax(dim),
736 Op::ConvWeightGradHost {
739 structure,
740 g,
741 x,
742 params,
743 transposed,
744 } => {
745 let structure_m = get_mirror(mirror, structure)?;
746 conv::host_weight_grad_node(
747 structure_m.tape(),
748 &structure_m,
749 &get_mirror(mirror, g)?,
750 &get_mirror(mirror, x)?,
751 ¶ms,
752 transposed,
753 )
754 }
755 Op::BatchNormGradHost {
758 structure,
759 g,
760 x,
761 weight,
762 has_bias,
763 eps,
764 fixed_stats,
765 kind,
766 } => {
767 let structure_m = get_mirror(mirror, structure)?;
768 let w_m = weight.map(|w| get_mirror(mirror, w)).transpose()?;
769 crate::create_graph_norm::host_grad_node(
770 structure_m.tape(),
771 &structure_m,
772 &get_mirror(mirror, g)?,
773 &get_mirror(mirror, x)?,
774 w_m.as_ref(),
775 has_bias,
776 eps,
777 fixed_stats,
778 kind,
779 )
780 }
781 Op::Conv2d {
784 input,
785 weight,
786 bias,
787 params,
788 compute_dtype,
789 } => {
790 let bias_m = bias.map(|b| get_mirror(mirror, b)).transpose()?;
791 conv::replay_conv2d(
792 &get_mirror(mirror, input)?,
793 &get_mirror(mirror, weight)?,
794 bias_m.as_ref(),
795 ¶ms,
796 compute_dtype,
797 )
798 }
799 Op::ConvTranspose2d {
800 input,
801 weight,
802 bias,
803 params,
804 } => {
805 let bias_m = bias.map(|b| get_mirror(mirror, b)).transpose()?;
806 conv::replay_conv_transpose2d(
807 shape,
808 &get_mirror(mirror, input)?,
809 &get_mirror(mirror, weight)?,
810 bias_m.as_ref(),
811 ¶ms,
812 )
813 }
814 Op::Unfold { input, params } => conv::replay_unfold(&get_mirror(mirror, input)?, ¶ms),
815 Op::Fold { input, params } => {
816 conv::replay_fold(shape, &get_mirror(mirror, input)?, ¶ms)
817 }
818 Op::MaxPool2d { input, index } | Op::MaxPool3d { input, index } => {
820 pool::replay_max_pool(&get_mirror(mirror, input)?, &index, shape)
821 }
822 Op::AvgPool2d {
823 input,
824 params,
825 count_include_pad,
826 } => pool::replay_avg_pool2d(&get_mirror(mirror, input)?, ¶ms, count_include_pad),
827 Op::AvgPool3d {
828 input,
829 params,
830 count_include_pad,
831 } => pool::replay_avg_pool3d(&get_mirror(mirror, input)?, ¶ms, count_include_pad),
832 Op::AdaptiveAvgPool2d { input } => {
833 pool::replay_adaptive_avg_pool2d(&get_mirror(mirror, input)?, shape)
834 }
835 Op::Gather { input, dim, index }
838 | Op::Sort { input, dim, index }
839 | Op::Topk { input, dim, index }
840 | Op::Cummax { input, dim, index }
841 | Op::Cummin { input, dim, index } => get_mirror(mirror, input)?.gather(dim, &index),
842 Op::Embedding { weight, index, .. } => get_mirror(mirror, weight)?.gather(0, &index),
843 Op::Scatter {
844 input,
845 dim,
846 index,
847 src,
848 reduce,
849 } => index_cum::replay_scatter(
850 &get_mirror(mirror, input)?,
851 &get_mirror(mirror, src)?,
852 dim,
853 &index,
854 reduce,
855 ),
856 Op::Cumsum { input, dim } => get_mirror(mirror, input)?.cumsum(dim),
857 Op::Pad { input, pads, value } => get_mirror(mirror, input)?.pad(&pads, value),
861 Op::MaskedFill { input, mask, value } => {
862 get_mirror(mirror, input)?.masked_fill(&f32_mask_to_bool(&mask)?, value)
863 }
864 Op::PadMode { input, pads, mode } => {
865 crate::pad_ops::pad_with_mode(&get_mirror(mirror, input)?, &pads, mode)
866 }
867 Op::Where { cond, a, b } => {
873 let cond_bool = f32_mask_to_bool(&cond)?;
874 Var::where_cond(&cond_bool, &get_mirror(mirror, a)?, &get_mirror(mirror, b)?)
875 }
876 Op::BatchNorm {
882 input,
883 weight,
884 bias,
885 eps,
886 fixed_stats,
887 } => {
888 let x = get_mirror(mirror, input)?;
889 let w = weight.map(|id| get_mirror(mirror, id)).transpose()?;
890 let b = bias.map(|id| get_mirror(mirror, id)).transpose()?;
891 match fixed_stats {
892 None => x.batch_norm(w.as_ref(), b.as_ref(), eps),
893 Some((mean, var)) => x.batch_norm_infer(w.as_ref(), b.as_ref(), &mean, &var, eps),
894 }
895 }
896 Op::RmsNorm { .. } | Op::LayerNorm { .. } => {
897 crate::create_graph_norm::replay_norm(&op, mirror)
898 }
899 Op::MseLoss {
902 pred,
903 target,
904 reduction,
905 } => get_mirror(mirror, pred)?.mse_loss_with(&get_mirror(mirror, target)?, reduction),
906 Op::L1Loss {
907 pred,
908 target,
909 reduction,
910 } => get_mirror(mirror, pred)?.l1_loss(&get_mirror(mirror, target)?, reduction),
911 Op::HuberLoss {
912 pred,
913 target,
914 kind,
915 delta,
916 reduction,
917 } => get_mirror(mirror, pred)?.huber_loss_impl(
918 &get_mirror(mirror, target)?,
919 kind,
920 delta,
921 reduction,
922 ),
923 Op::CrossEntropyLoss {
924 logits,
925 targets,
926 class_dim,
927 reduction,
928 } => get_mirror(mirror, logits)?.cross_entropy_loss(&targets, class_dim, reduction),
929 Op::NllLoss {
930 input,
931 targets,
932 class_dim,
933 reduction,
934 } => get_mirror(mirror, input)?.nll_loss(&targets, class_dim, reduction),
935 Op::Rfft {
940 input,
941 n,
942 dim,
943 norm,
944 } => get_mirror(mirror, input)?.rfft(Some(n), Some(dim), norm),
945 Op::Irfft {
946 input,
947 n,
948 dim,
949 norm,
950 } => get_mirror(mirror, input)?.irfft(Some(n), Some(dim), norm),
951 Op::Fft {
952 input,
953 n,
954 dim,
955 norm,
956 } => get_mirror(mirror, input)?.fft(Some(n), Some(dim), norm),
957 Op::Ifft {
958 input,
959 n,
960 dim,
961 norm,
962 } => get_mirror(mirror, input)?.ifft(Some(n), Some(dim), norm),
963 Op::Inv { input } => get_mirror(mirror, input)?.inv(),
964 Op::Solve { a, b } => get_mirror(mirror, a)?.solve(&get_mirror(mirror, b)?),
965 Op::Det { input } => get_mirror(mirror, input)?.det(),
966 Op::SlogdetLogAbsDet { input } => Ok(get_mirror(mirror, input)?.slogdet()?.logabsdet),
967 _ => Err(AutodiffError::Backward(
968 "create_graph: replay_op: supports_create_graph() が true の未対応 Op(内部契約違反)"
969 .into(),
970 )),
971 }
972}
973
974fn child_value(child: &Tape, v: &Var<'_>) -> Result<Tensor<f32>, AutodiffError> {
977 let nodes = child.nodes.borrow();
978 Ok(materialize_fallible(&nodes, child.ops(), v.node_id())?.clone())
979}
980
981fn push_v3<'c>(
987 child: &'c Tape,
988 composed: &Var<'c>,
989 host_value: Tensor<f32>,
990) -> Result<Var<'c>, AutodiffError> {
991 if host_value.shape() != composed.shape().as_slice() {
992 return Err(AutodiffError::Backward(
993 "create_graph: V3 のホスト値と構造グラフの shape が不一致(内部契約違反)".into(),
994 ));
995 }
996 let nid = child.push_eager(
997 Op::Contiguous {
998 input: composed.node_id(),
999 },
1000 host_value,
1001 );
1002 Ok(Var::from_raw(child, nid))
1003}
1004
1005fn reduce_dim_bcast<'c>(
1007 v: &Var<'c>,
1008 dim: usize,
1009 shape: &[usize],
1010) -> Result<Var<'c>, AutodiffError> {
1011 v.sum(Some(dim))?
1012 .contiguous()?
1013 .unsqueeze(dim)?
1014 .broadcast_to(shape)
1015}
1016
1017fn softmax_cgrad_v3<'c>(
1029 parent: &Tape,
1030 child: &'c Tape,
1031 parent_nodes: &[TapeNode],
1032 id: NodeId,
1033 x_m: Var<'c>,
1034 g: &Var<'c>,
1035 dim: usize,
1036) -> Result<Var<'c>, AutodiffError> {
1037 let shape = parent_nodes[id.0].shape.clone();
1038 if shape.contains(&0) {
1039 return Ok(child.var_no_grad(&Tensor::zeros(&shape)?));
1040 }
1041 let y_a = x_m.softmax(dim)?;
1042 let y_b = x_m.softmax(dim)?;
1043 let y_c = x_m.softmax(dim)?;
1044 let out_value = materialize_fallible(parent_nodes, parent.ops(), id)?;
1045 let g_value = child_value(child, g)?;
1046 let t1 = y_a.mul(g)?;
1047 let dot = reduce_dim_bcast(&g.mul(&y_c)?, dim, &shape)?;
1048 let composed = t1.sub(&y_b.mul(&dot)?)?;
1049 let host = crate::grad::softmax_vjp_along(out_value, &g_value, dim);
1050 push_v3(child, &composed, host)
1051}
1052
1053const LANE_SCALE_FLOOR: f32 = 9.536_743e-7;
1064
1065fn lane_scale(g: &Tensor<f32>, dim: usize) -> Result<Tensor<f32>, AutodiffError> {
1066 let shape = g.shape().to_vec();
1067 let outer: usize = shape[..dim].iter().product();
1068 let axis_len = shape[dim];
1069 let inner: usize = shape[dim + 1..].iter().product();
1070 let data = crate::eval::dense_vec(g);
1071 let mut s = Vec::with_capacity(outer * inner);
1072 for o in 0..outer {
1073 for i in 0..inner {
1074 let mut m = 0f32;
1075 for a in 0..axis_len {
1076 let v = data[(o * axis_len + a) * inner + i].abs();
1077 if v.is_nan() {
1079 m = f32::NAN;
1080 break;
1081 }
1082 if v > m {
1083 m = v;
1084 }
1085 }
1086 s.push(if m == 0.0 || !m.is_finite() {
1087 1.0
1088 } else {
1089 m.max(LANE_SCALE_FLOOR)
1090 });
1091 }
1092 }
1093 let mut sshape = shape;
1094 sshape[dim] = 1;
1095 Ok(Tensor::new(s, &sshape)?)
1096}
1097
1098const SPLIT_LIMIT: f64 = 2.0e38;
1103const SPLIT_MAX: usize = 64;
1104
1105fn log_softmax_split_count(
1106 y: &Tensor<f32>,
1107 g: &Tensor<f32>,
1108 s_lane: &Tensor<f32>,
1109 dim: usize,
1110) -> usize {
1111 let shape = g.shape().to_vec();
1112 let outer: usize = shape[..dim].iter().product();
1113 let n = shape[dim];
1114 let inner: usize = shape[dim + 1..].iter().product();
1115 let yd = crate::eval::dense_vec(y);
1116 let gd = crate::eval::dense_vec(g);
1117 let sd = crate::eval::dense_vec(s_lane);
1118 let mut worst = 0f64;
1119 for o in 0..outer {
1120 for i in 0..inner {
1121 let s = sd[o * inner + i] as f64;
1122 let mut sum = 0f64;
1123 let mut pmax = 0f64;
1124 for a in 0..n {
1125 let idx = (o * n + a) * inner + i;
1126 sum += gd[idx] as f64 / s;
1127 pmax = pmax.max((yd[idx] as f64).exp());
1128 }
1129 let mean = (sum / n as f64).abs();
1130 let m = n as f64 * pmax * s * mean;
1131 if m.is_finite() {
1133 worst = worst.max(m);
1134 }
1135 }
1136 }
1137 ((worst / SPLIT_LIMIT).ceil() as usize).clamp(1, SPLIT_MAX)
1138}
1139
1140fn log_softmax_cgrad_v3<'c>(
1144 parent: &Tape,
1145 child: &'c Tape,
1146 parent_nodes: &[TapeNode],
1147 x_m: Var<'c>,
1148 id: NodeId,
1149 g: &Var<'c>,
1150 dim: usize,
1151) -> Result<Var<'c>, AutodiffError> {
1152 let shape = parent_nodes[id.0].shape.clone();
1153 if shape.contains(&0) {
1154 return Ok(child.var_no_grad(&Tensor::zeros(&shape)?));
1155 }
1156 let out_value = materialize_fallible(parent_nodes, parent.ops(), id)?;
1157 let g_value = child_value(child, g)?;
1158 let s_lane = lane_scale(&g_value, dim)?;
1159 let pieces = log_softmax_split_count(out_value, &g_value, &s_lane, dim);
1160 let n_lane = shape[dim] as f32;
1170 let ln_n = child.var_no_grad(&Tensor::full(&shape, n_lane.ln())?);
1171 let inv_nk = child.var_no_grad(&Tensor::full(&shape, 1.0 / (n_lane * pieces as f32))?);
1172 let mean_g_k = reduce_dim_bcast(&g.mul(&inv_nk)?, dim, &shape)?;
1177 let mut sub: Option<Var<'c>> = None;
1178 for _ in 0..pieces {
1179 let y_k = x_m.log_softmax(dim)?;
1180 let term = y_k.add(&ln_n)?.exp().mul(&mean_g_k)?;
1181 sub = Some(match sub {
1182 Some(acc) => acc.add(&term)?,
1183 None => term,
1184 });
1185 }
1186 let sub = sub.ok_or_else(|| AutodiffError::Backward("LogSoftmax 分割数が 0".into()))?;
1187 let composed = g.sub(&sub)?;
1188 let host = crate::grad::log_softmax_vjp_dispatch(child.ops(), out_value, &g_value, dim)?;
1189 push_v3(child, &composed, host)
1190}
1191
1192fn build_cgrads<'c>(
1198 parent: &Tape,
1199 child: &'c Tape,
1200 ancestors: &[NodeId],
1201 mirror: &[Option<Var<'c>>],
1202 loss_id: NodeId,
1203) -> Result<Vec<Option<Var<'c>>>, AutodiffError> {
1204 let parent_nodes = parent.nodes.borrow();
1205 let parent_ops = parent.ops();
1206 let n = parent_nodes.len();
1207 let mut cgrads: Vec<Option<Var<'c>>> = vec![None; n];
1208
1209 let loss_shape = parent_nodes[loss_id.0].shape.clone();
1212 let seed_tensor = Tensor::full(&loss_shape, 1.0f32)?;
1213 cgrads[loss_id.0] = Some(child.var_no_grad(&seed_tensor));
1214
1215 for &id in ancestors.iter().rev() {
1216 let Some(g) = cgrads[id.0] else {
1217 continue;
1218 };
1219 let node = &parent_nodes[id.0];
1220 match node.op.clone() {
1221 Op::Leaf => {}
1222 Op::Add(a, b) => {
1223 let a_shape = parent_nodes[a.0].shape.clone();
1224 let b_shape = parent_nodes[b.0].shape.clone();
1225 let da = reduce_bias_grad_var(child, &g, &a_shape)?;
1235 let db = reduce_bias_grad_var(child, &g, &b_shape)?;
1236 accumulate(&parent_nodes, &mut cgrads, a, da)?;
1237 accumulate(&parent_nodes, &mut cgrads, b, db)?;
1238 }
1239 Op::Mul(a, b) => {
1240 let a_m = get_mirror(mirror, a)?;
1241 let b_m = get_mirror(mirror, b)?;
1242 let ga = g.mul(&b_m)?;
1243 let gb = g.mul(&a_m)?;
1244 let da = reduce_to(child, &ga, &a_m.shape())?;
1245 let db = reduce_to(child, &gb, &b_m.shape())?;
1246 accumulate(&parent_nodes, &mut cgrads, a, da)?;
1247 accumulate(&parent_nodes, &mut cgrads, b, db)?;
1248 }
1249 Op::MatMul(a, b) => {
1250 let a_m = get_mirror(mirror, a)?;
1267 let b_m = get_mirror(mirror, b)?;
1268 let da = g.matmul_fp32_strict(&b_m.transpose(0, 1)?)?;
1269 let db = a_m.transpose(0, 1)?.matmul_fp32_strict(&g)?;
1270 accumulate(&parent_nodes, &mut cgrads, a, da)?;
1271 accumulate(&parent_nodes, &mut cgrads, b, db)?;
1272 }
1273 Op::Relu(a) => {
1274 let a_val = materialize_fallible(&parent_nodes, parent_ops, a)?;
1279 let mask = positive_mask(a_val)?;
1280 let zeros = child.var_no_grad(&Tensor::zeros(&g.shape())?);
1281 let da = Var::where_cond(&mask, &g, &zeros)?;
1282 accumulate(&parent_nodes, &mut cgrads, a, da)?;
1283 }
1284 Op::Exp(a) => {
1285 let out_c = get_mirror(mirror, id)?;
1289 let da = g.mul(&out_c)?;
1290 accumulate(&parent_nodes, &mut cgrads, a, da)?;
1291 }
1292 Op::Tanh(a) => {
1293 let out_c = get_mirror(mirror, id)?;
1295 let one = child.var_no_grad(&Tensor::full(&out_c.shape(), 1.0f32)?);
1296 let sq = out_c.mul(&out_c)?;
1297 let factor = one.sub(&sq)?;
1298 let da = g.mul(&factor)?;
1299 accumulate(&parent_nodes, &mut cgrads, a, da)?;
1300 }
1301 Op::Sigmoid(a) => {
1302 let out_c = get_mirror(mirror, id)?;
1304 let one = child.var_no_grad(&Tensor::full(&out_c.shape(), 1.0f32)?);
1305 let comp = one.sub(&out_c)?;
1306 let factor = out_c.mul(&comp)?;
1307 let da = g.mul(&factor)?;
1308 accumulate(&parent_nodes, &mut cgrads, a, da)?;
1309 }
1310 Op::Sum { input, dim } => {
1311 let input_shape = parent_nodes[input.0].shape.clone();
1312 let da = match dim {
1313 None => g.broadcast_to(&input_shape)?,
1314 Some(axis) => g
1315 .contiguous()?
1316 .unsqueeze(axis)?
1317 .broadcast_to(&input_shape)?,
1318 };
1319 accumulate(&parent_nodes, &mut cgrads, input, da)?;
1320 }
1321 Op::Mean { input, dim } => {
1322 let input_shape = parent_nodes[input.0].shape.clone();
1323 let count: usize = match dim {
1324 None => input_shape.iter().product(),
1325 Some(axis) => input_shape[axis],
1326 };
1327 let scale = child.var_no_grad(&Tensor::scalar(1.0f32 / count as f32));
1328 let g_scaled = g.mul(&scale)?;
1329 let da = match dim {
1330 None => g_scaled.broadcast_to(&input_shape)?,
1331 Some(axis) => g_scaled
1332 .contiguous()?
1333 .unsqueeze(axis)?
1334 .broadcast_to(&input_shape)?,
1335 };
1336 accumulate(&parent_nodes, &mut cgrads, input, da)?;
1337 }
1338 Op::Reshape { input } => {
1339 let input_shape = parent_nodes[input.0].shape.clone();
1340 let da = g.contiguous()?.reshape(&input_shape)?;
1341 accumulate(&parent_nodes, &mut cgrads, input, da)?;
1342 }
1343 Op::BroadcastTo { input } => {
1344 let input_shape = parent_nodes[input.0].shape.clone();
1345 let da = reduce_to(child, &g, &input_shape)?;
1346 accumulate(&parent_nodes, &mut cgrads, input, da)?;
1347 }
1348 Op::Transpose { input, dim0, dim1 } => {
1349 let da = g.transpose(dim0, dim1)?;
1353 accumulate(&parent_nodes, &mut cgrads, input, da)?;
1354 }
1355 Op::Permute { input, perm } => {
1356 let inv = crate::grad::inverse_permutation(&perm);
1360 let da = g.permute(&inv)?;
1361 accumulate(&parent_nodes, &mut cgrads, input, da)?;
1362 }
1363 Op::Narrow {
1364 input,
1365 dim,
1366 start,
1367 len,
1368 } => {
1369 let input_shape = parent_nodes[input.0].shape.clone();
1373 let before_len = start;
1374 let after_len = input_shape[dim] - start - len;
1375 let mut pieces: Vec<Var<'c>> = Vec::with_capacity(3);
1376 if before_len > 0 {
1377 let mut before_shape = input_shape.clone();
1378 before_shape[dim] = before_len;
1379 pieces.push(child.var_no_grad(&Tensor::zeros(&before_shape)?));
1380 }
1381 pieces.push(g);
1382 if after_len > 0 {
1383 let mut after_shape = input_shape.clone();
1384 after_shape[dim] = after_len;
1385 pieces.push(child.var_no_grad(&Tensor::zeros(&after_shape)?));
1386 }
1387 let da = if pieces.len() == 1 {
1388 g
1389 } else {
1390 Var::cat(&pieces, dim)?
1391 };
1392 accumulate(&parent_nodes, &mut cgrads, input, da)?;
1393 }
1394 Op::Concat { inputs, dim } => {
1395 let mut off = 0usize;
1399 for input in inputs {
1400 let len = parent_nodes[input.0].shape[dim];
1401 let da = g.narrow(dim, off, len)?;
1402 accumulate(&parent_nodes, &mut cgrads, input, da)?;
1403 off += len;
1404 }
1405 }
1406 Op::Softmax { input, dim } => {
1407 let da = softmax_cgrad_v3(
1408 parent,
1409 child,
1410 &parent_nodes,
1411 id,
1412 get_mirror(mirror, input)?,
1413 &g,
1414 dim,
1415 )?;
1416 accumulate(&parent_nodes, &mut cgrads, input, da)?;
1417 }
1418 Op::LogSoftmax { input, dim } => {
1419 let da = log_softmax_cgrad_v3(
1420 parent,
1421 child,
1422 &parent_nodes,
1423 get_mirror(mirror, input)?,
1424 id,
1425 &g,
1426 dim,
1427 )?;
1428 accumulate(&parent_nodes, &mut cgrads, input, da)?;
1429 }
1430 Op::MseLoss {
1431 pred,
1432 target,
1433 reduction,
1434 } => {
1435 let pred_m = get_mirror(mirror, pred)?;
1436 let target_m = get_mirror(mirror, target)?;
1437 let shape = parent_nodes[pred.0].shape.clone();
1438 let dpred = if shape.contains(&0) {
1439 child.var_no_grad(&Tensor::zeros(&shape)?)
1440 } else {
1441 let n: usize = shape.iter().product();
1442 let s_b = loss_scale_bcast(child, &g, reduction, true, n, &shape)?;
1443 pred_m.sub(&target_m)?.mul(&s_b)?
1444 };
1445 let dtarget = dpred.neg()?;
1446 accumulate(&parent_nodes, &mut cgrads, pred, dpred)?;
1447 accumulate(&parent_nodes, &mut cgrads, target, dtarget)?;
1448 }
1449 Op::L1Loss {
1450 pred,
1451 target,
1452 reduction,
1453 } => {
1454 let pred_m = get_mirror(mirror, pred)?;
1455 let target_m = get_mirror(mirror, target)?;
1456 let shape = parent_nodes[pred.0].shape.clone();
1457 let dpred = if shape.contains(&0) {
1458 child.var_no_grad(&Tensor::zeros(&shape)?)
1459 } else {
1460 let n: usize = shape.iter().product();
1461 let g_b = loss_scale_bcast(child, &g, reduction, false, n, &shape)?;
1462 let p_val = materialize_fallible(&parent_nodes, parent_ops, pred)?;
1463 let t_val = materialize_fallible(&parent_nodes, parent_ops, target)?;
1464 let pos = pairwise_mask(p_val, t_val, |p, t| p - t > 0.0)?;
1465 let neg = pairwise_mask(p_val, t_val, |p, t| p - t < 0.0)?;
1466 let nan = pairwise_mask(p_val, t_val, |p, t| (p - t).is_nan())?;
1467 let d_m = pred_m.sub(&target_m)?;
1468 let zeros = child.var_no_grad(&Tensor::zeros(&shape)?);
1469 let nan_branch = Var::where_cond(&nan, &d_m, &zeros)?;
1472 let neg_branch = Var::where_cond(&neg, &g_b.neg()?, &nan_branch)?;
1473 Var::where_cond(&pos, &g_b, &neg_branch)?
1474 };
1475 let dtarget = dpred.neg()?;
1476 accumulate(&parent_nodes, &mut cgrads, pred, dpred)?;
1477 accumulate(&parent_nodes, &mut cgrads, target, dtarget)?;
1478 }
1479 Op::HuberLoss {
1480 pred,
1481 target,
1482 kind,
1483 delta,
1484 reduction,
1485 } => {
1486 let pred_m = get_mirror(mirror, pred)?;
1487 let target_m = get_mirror(mirror, target)?;
1488 let shape = parent_nodes[pred.0].shape.clone();
1489 let smooth_l1 = match kind {
1490 fandhe_ai_tensor_core::HuberKind::Huber => false,
1491 fandhe_ai_tensor_core::HuberKind::SmoothL1 => true,
1492 _ => {
1493 return Err(AutodiffError::Backward(
1494 "create_graph: HuberLoss: 未知の HuberKind は fail-closed".into(),
1495 ));
1496 }
1497 };
1498 let dpred = if shape.contains(&0) {
1499 child.var_no_grad(&Tensor::zeros(&shape)?)
1500 } else {
1501 let n: usize = shape.iter().product();
1502 let g_b = loss_scale_bcast(child, &g, reduction, false, n, &shape)?;
1503 let p_val = materialize_fallible(&parent_nodes, parent_ops, pred)?;
1504 let t_val = materialize_fallible(&parent_nodes, parent_ops, target)?;
1505 let quad = pairwise_mask(p_val, t_val, |p, t| {
1508 let a = (p - t).abs();
1509 a < delta || a.is_nan()
1510 })?;
1511 let lin_data: Vec<f32> = {
1512 let pc = p_val.contiguous();
1513 let tc = t_val.contiguous();
1514 let ps = pc.as_slice().unwrap_or(&[]);
1515 let ts = tc.as_slice().unwrap_or(&[]);
1516 ps.iter()
1517 .zip(ts.iter())
1518 .map(|(&p, &t)| {
1519 let d = p - t;
1520 if smooth_l1 {
1521 1.0f32.copysign(d)
1522 } else {
1523 delta.copysign(d)
1524 }
1525 })
1526 .collect()
1527 };
1528 let lin_const = child.var_no_grad(&Tensor::new(lin_data, &shape)?);
1529 let d_m = pred_m.sub(&target_m)?;
1530 let zeros_q = child.var_no_grad(&Tensor::zeros(&shape)?);
1536 let d_safe = Var::where_cond(&quad, &d_m, &zeros_q)?;
1537 let quad_grad = if smooth_l1 {
1538 let delta_c = child.var_no_grad(&Tensor::scalar(delta));
1539 g_b.mul(&d_safe.div(&delta_c)?)?
1540 } else {
1541 g_b.mul(&d_safe)?
1542 };
1543 let lin_grad = g_b.mul(&lin_const)?;
1544 Var::where_cond(&quad, &quad_grad, &lin_grad)?
1545 };
1546 let dtarget = dpred.neg()?;
1547 accumulate(&parent_nodes, &mut cgrads, pred, dpred)?;
1548 accumulate(&parent_nodes, &mut cgrads, target, dtarget)?;
1549 }
1550 Op::CrossEntropyLoss {
1551 logits,
1552 targets,
1553 class_dim,
1554 reduction,
1555 } => {
1556 let logits_m = get_mirror(mirror, logits)?;
1557 let shape = parent_nodes[logits.0].shape.clone();
1558 let dlogits = if shape.contains(&0) {
1559 child.var_no_grad(&Tensor::zeros(&shape)?)
1560 } else {
1561 let n = class_sample_count(&shape, class_dim);
1562 let s_b = loss_scale_bcast(child, &g, reduction, false, n, &shape)?;
1563 let onehot = child.var_no_grad(&onehot_f32(&targets, &shape, class_dim)?);
1564 let y = logits_m.softmax(class_dim)?;
1565 y.sub(&onehot)?.mul(&s_b)?
1566 };
1567 accumulate(&parent_nodes, &mut cgrads, logits, dlogits)?;
1568 }
1569 Op::NllLoss {
1570 input,
1571 targets,
1572 class_dim,
1573 reduction,
1574 } => {
1575 let shape = parent_nodes[input.0].shape.clone();
1576 let dinput = if shape.contains(&0) {
1577 child.var_no_grad(&Tensor::zeros(&shape)?)
1578 } else {
1579 let n = class_sample_count(&shape, class_dim);
1580 let s_b = loss_scale_bcast(child, &g, reduction, false, n, &shape)?;
1581 let mask = onehot_mask(&targets, &shape, class_dim)?;
1582 let zeros = child.var_no_grad(&Tensor::zeros(&shape)?);
1583 Var::where_cond(&mask, &s_b.neg()?, &zeros)?
1584 };
1585 accumulate(&parent_nodes, &mut cgrads, input, dinput)?;
1586 }
1587 Op::Contiguous { input } => {
1588 accumulate(&parent_nodes, &mut cgrads, input, g)?;
1590 }
1591 Op::ConvWeightGradHost { structure, .. } | Op::BatchNormGradHost { structure, .. } => {
1592 accumulate(&parent_nodes, &mut cgrads, structure, g)?;
1594 }
1595 Op::Where { cond, a, b } => {
1596 let a_shape = parent_nodes[a.0].shape.clone();
1601 let b_shape = parent_nodes[b.0].shape.clone();
1602 let mask = f32_mask_to_bool(&cond)?;
1603 let zeros = child.var_no_grad(&Tensor::zeros(&g.shape())?);
1604 let ga = Var::where_cond(&mask, &g, &zeros)?;
1605 let gb = Var::where_cond(&mask, &zeros, &g)?;
1606 let da = reduce_to(child, &ga, &a_shape)?;
1607 let db = reduce_to(child, &gb, &b_shape)?;
1608 accumulate(&parent_nodes, &mut cgrads, a, da)?;
1609 accumulate(&parent_nodes, &mut cgrads, b, db)?;
1610 }
1611 Op::BatchNorm {
1612 input,
1613 weight,
1614 bias,
1615 eps,
1616 fixed_stats,
1617 } => {
1618 let x_host = materialize_fallible(&parent_nodes, parent_ops, input)?.clone();
1621 let (n_, c_, spatial_) = batch_norm_layout(x_host.shape())?;
1622 let inputs = crate::create_graph_norm::BatchNormCgradInputs {
1623 x_mirror: get_mirror(mirror, input)?,
1624 w_mirror: weight.map(|w| get_mirror(mirror, w)).transpose()?,
1625 has_bias: bias.is_some(),
1626 eps,
1627 fixed_stats: fixed_stats.as_ref(),
1628 x_host: &x_host,
1629 need_dx: parent_nodes[input.0].requires_grad,
1630 need_dw: weight.is_some_and(|w| parent_nodes[w.0].requires_grad),
1631 need_db: bias.is_some_and(|b| parent_nodes[b.0].requires_grad),
1632 n: n_,
1633 c: c_,
1634 spatial: spatial_,
1635 };
1636 let out = crate::create_graph_norm::batch_norm_cgrads(child, &inputs, &g)?;
1637 if let Some(dx) = out.dx {
1638 accumulate(&parent_nodes, &mut cgrads, input, dx)?;
1639 }
1640 if let (Some(w), Some(dw)) = (weight, out.dw) {
1641 accumulate(&parent_nodes, &mut cgrads, w, dw)?;
1642 }
1643 if let (Some(b), Some(db)) = (bias, out.db) {
1644 accumulate(&parent_nodes, &mut cgrads, b, db)?;
1645 }
1646 }
1647 Op::ScalarUnary { op: sop, input } => {
1648 let x_m = get_mirror(mirror, input)?;
1653 let da = match sop {
1654 ScalarUnaryOp::Neg => g.neg()?,
1655 ScalarUnaryOp::Abs => {
1656 let x_val = materialize_fallible(&parent_nodes, parent_ops, input)?;
1657 let pos = mask_from_pred(x_val, |v| v > 0.0)?;
1658 let neg = mask_from_pred(x_val, |v| v < 0.0)?;
1659 let zeros = child.var_no_grad(&Tensor::zeros(&g.shape())?);
1660 let neg_branch = Var::where_cond(&neg, &g.neg()?, &zeros)?;
1661 Var::where_cond(&pos, &g, &neg_branch)?
1662 }
1663 ScalarUnaryOp::Sqrt => {
1664 let y_c = get_mirror(mirror, id)?;
1665 let half = child.var_no_grad(&Tensor::scalar(0.5f32));
1666 let factor = half.div(&y_c)?;
1667 g.mul(&factor)?
1668 }
1669 ScalarUnaryOp::Log => {
1670 let one = child.var_no_grad(&Tensor::scalar(1.0f32));
1671 let factor = one.div(&x_m)?;
1672 g.mul(&factor)?
1673 }
1674 ScalarUnaryOp::Log2 => {
1675 let one = child.var_no_grad(&Tensor::scalar(1.0f32));
1676 let ln2 = child.var_no_grad(&Tensor::scalar(std::f32::consts::LN_2));
1677 let factor = one.div(&x_m.mul(&ln2)?)?;
1678 g.mul(&factor)?
1679 }
1680 ScalarUnaryOp::Log10 => {
1681 let one = child.var_no_grad(&Tensor::scalar(1.0f32));
1682 let ln10 = child.var_no_grad(&Tensor::scalar(std::f32::consts::LN_10));
1683 let factor = one.div(&x_m.mul(&ln10)?)?;
1684 g.mul(&factor)?
1685 }
1686 ScalarUnaryOp::Sin => {
1687 let factor = x_m.cos()?;
1688 g.mul(&factor)?
1689 }
1690 ScalarUnaryOp::Cos => {
1691 let factor = x_m.sin()?.neg()?;
1692 g.mul(&factor)?
1693 }
1694 ScalarUnaryOp::Tan => {
1695 let one = child.var_no_grad(&Tensor::scalar(1.0f32));
1696 let c = x_m.cos()?;
1697 let factor = one.div(&c.mul(&c)?)?;
1698 g.mul(&factor)?
1699 }
1700 ScalarUnaryOp::Relu => {
1701 let x_val = materialize_fallible(&parent_nodes, parent_ops, input)?;
1702 let mask = positive_mask(x_val)?;
1703 let zeros = child.var_no_grad(&Tensor::zeros(&g.shape())?);
1704 Var::where_cond(&mask, &g, &zeros)?
1705 }
1706 ScalarUnaryOp::Exp => {
1707 let out_c = get_mirror(mirror, id)?;
1708 g.mul(&out_c)?
1709 }
1710 ScalarUnaryOp::Tanh => {
1711 let out_c = get_mirror(mirror, id)?;
1712 let onec = child.var_no_grad(&Tensor::full(&out_c.shape(), 1.0f32)?);
1713 let sq = out_c.mul(&out_c)?;
1714 let factor = onec.sub(&sq)?;
1715 g.mul(&factor)?
1716 }
1717 ScalarUnaryOp::Sigmoid => {
1718 let out_c = get_mirror(mirror, id)?;
1719 let onec = child.var_no_grad(&Tensor::full(&out_c.shape(), 1.0f32)?);
1720 let comp = onec.sub(&out_c)?;
1721 let factor = out_c.mul(&comp)?;
1722 g.mul(&factor)?
1723 }
1724 ScalarUnaryOp::Silu => {
1725 let s = x_m.sigmoid();
1726 let onec = child.var_no_grad(&Tensor::full(&x_m.shape(), 1.0f32)?);
1727 let one_minus_s = onec.sub(&s)?;
1728 let term = x_m.mul(&one_minus_s)?;
1729 let inner = onec.add(&term)?;
1730 let factor = s.mul(&inner)?;
1731 g.mul(&factor)?
1732 }
1733 ScalarUnaryOp::Hardswish => {
1734 let x_val = materialize_fallible(&parent_nodes, parent_ops, input)?;
1735 let lo = mask_from_pred(x_val, |v| v <= -3.0)?;
1736 let hi = mask_from_pred(x_val, |v| v >= 3.0)?;
1737 let two = child.var_no_grad(&Tensor::scalar(2.0f32));
1738 let three = child.var_no_grad(&Tensor::scalar(3.0f32));
1739 let six = child.var_no_grad(&Tensor::scalar(6.0f32));
1740 let mid = x_m.mul(&two)?.add(&three)?.div(&six)?;
1741 let onec = child.var_no_grad(&Tensor::full(&x_m.shape(), 1.0f32)?);
1742 let hi_branch = Var::where_cond(&hi, &onec, &mid)?;
1743 let zeros = child.var_no_grad(&Tensor::zeros(&x_m.shape())?);
1744 let factor = Var::where_cond(&lo, &zeros, &hi_branch)?;
1745 g.mul(&factor)?
1746 }
1747 ScalarUnaryOp::LeakyRelu { negative_slope } => {
1748 let x_val = materialize_fallible(&parent_nodes, parent_ops, input)?;
1749 let mask = mask_from_pred(x_val, |v| v >= 0.0)?;
1750 let slope = child.var_no_grad(&Tensor::scalar(negative_slope));
1751 let else_branch = g.mul(&slope)?;
1752 Var::where_cond(&mask, &g, &else_branch)?
1753 }
1754 ScalarUnaryOp::Elu { alpha } => {
1755 let x_val = materialize_fallible(&parent_nodes, parent_ops, input)?;
1756 let mask = mask_from_pred(x_val, |v| v > 0.0)?;
1757 let alpha_c = child.var_no_grad(&Tensor::scalar(alpha));
1758 let factor = x_m.exp().mul(&alpha_c)?;
1759 let else_branch = g.mul(&factor)?;
1760 Var::where_cond(&mask, &g, &else_branch)?
1761 }
1762 ScalarUnaryOp::Softplus { beta, threshold } => {
1763 let x_val = materialize_fallible(&parent_nodes, parent_ops, input)?;
1764 let mask = mask_from_pred(x_val, |v| v * beta > threshold)?;
1765 let beta_c = child.var_no_grad(&Tensor::scalar(beta));
1766 let factor = x_m.mul(&beta_c)?.sigmoid();
1767 let else_branch = g.mul(&factor)?;
1768 Var::where_cond(&mask, &g, &else_branch)?
1769 }
1770 ScalarUnaryOp::Clamp { min, max } => {
1771 let x_val = materialize_fallible(&parent_nodes, parent_ops, input)?;
1772 let mask = mask_from_pred(x_val, |v| {
1773 !(v.is_nan() || min > max || v < min || v > max)
1774 })?;
1775 let zeros = child.var_no_grad(&Tensor::zeros(&g.shape())?);
1776 Var::where_cond(&mask, &g, &zeros)?
1777 }
1778 ScalarUnaryOp::PowScalar { exponent } => {
1779 if exponent == 0.0 {
1780 child.var_no_grad(&Tensor::zeros(&g.shape())?)
1781 } else {
1782 let exp_c = child.var_no_grad(&Tensor::scalar(exponent));
1783 let factor = x_m
1784 .scalar_unary(ScalarUnaryOp::PowScalar {
1785 exponent: exponent - 1.0,
1786 })?
1787 .mul(&exp_c)?;
1788 g.mul(&factor)?
1789 }
1790 }
1791 ScalarUnaryOp::Erf => {
1792 let c = child.var_no_grad(&Tensor::scalar(
1794 (2.0f64 * 0.564_189_583_547_756_3f64) as f32,
1795 ));
1796 let factor = x_m.mul(&x_m)?.neg()?.exp().mul(&c)?;
1797 g.mul(&factor)?
1798 }
1799 ScalarUnaryOp::Gelu => {
1800 let inv_sqrt2 =
1803 child.var_no_grad(&Tensor::scalar(std::f32::consts::FRAC_1_SQRT_2));
1804 let half = child.var_no_grad(&Tensor::scalar(0.5f32));
1805 let one = child.var_no_grad(&Tensor::scalar(1.0f32));
1806 let neg_half = child.var_no_grad(&Tensor::scalar(-0.5f32));
1807 let inv_sqrt_2pi = child.var_no_grad(&Tensor::scalar(0.398_942_3_f32));
1808 let e = x_m.mul(&inv_sqrt2)?.scalar_unary(ScalarUnaryOp::Erf)?;
1809 let cdf = one.add(&e)?.mul(&half)?;
1810 let pdf = x_m.mul(&x_m)?.mul(&neg_half)?.exp().mul(&inv_sqrt_2pi)?;
1811 let factor = cdf.add(&x_m.mul(&pdf)?)?;
1812 g.mul(&factor)?
1813 }
1814 ScalarUnaryOp::GeluTanh => {
1815 let x_val = materialize_fallible(&parent_nodes, parent_ops, input)?;
1820 const C: f32 = 0.797_884_6;
1821 const K: f32 = 0.044_715;
1822 let t_host = |v: f32| (C * (v + K * v * v * v)).tanh();
1823 let sat = mask_from_pred(x_val, |v| {
1824 let t = t_host(v);
1825 1.0 - t * t == 0.0
1826 })?;
1827 let contiguous = x_val.contiguous();
1828 let limit_data: Vec<f32> = contiguous
1829 .as_slice()
1830 .map(|s| s.iter().map(|&v| 0.5 * (1.0 + t_host(v))).collect())
1831 .unwrap_or_default();
1832 let limit =
1833 child.var_no_grad(&Tensor::new(limit_data, contiguous.shape())?);
1834 let zeros = child.var_no_grad(&Tensor::zeros(&x_m.shape())?);
1835 let x_safe = Var::where_cond(&sat, &zeros, &x_m)?;
1836 let c = child.var_no_grad(&Tensor::scalar(C));
1837 let k = child.var_no_grad(&Tensor::scalar(K));
1838 let k3 = child.var_no_grad(&Tensor::scalar(3.0 * K));
1839 let half = child.var_no_grad(&Tensor::scalar(0.5f32));
1840 let one = child.var_no_grad(&Tensor::scalar(1.0f32));
1841 let x2 = x_safe.mul(&x_safe)?;
1842 let x3 = x2.mul(&x_safe)?;
1843 let u = x_safe.add(&k.mul(&x3)?)?.mul(&c)?;
1844 let t = u.tanh();
1845 let sech2 = one.sub(&t.mul(&t)?)?;
1846 let du = one.add(&k3.mul(&x2)?)?.mul(&c)?;
1847 let first = one.add(&t)?.mul(&half)?;
1848 let second = x_safe.mul(&sech2)?.mul(&du)?.mul(&half)?;
1849 let smooth = first.add(&second)?;
1850 let factor = Var::where_cond(&sat, &limit, &smooth)?;
1851 g.mul(&factor)?
1852 }
1853 _ => {
1854 return Err(AutodiffError::Backward(format!(
1855 "create_graph: build_cgrads: 未対応の ScalarUnaryOp {sop:?}\
1856 (内部契約違反。scalar_unary_replayable との判定不整合)"
1857 )));
1858 }
1859 };
1860 accumulate(&parent_nodes, &mut cgrads, input, da)?;
1861 }
1862 Op::ScalarBinary { op: sop, a, b } => {
1863 let a_shape = parent_nodes[a.0].shape.clone();
1864 let b_shape = parent_nodes[b.0].shape.clone();
1865 if sop.is_comparison() {
1866 let da = child.var_no_grad(&Tensor::zeros(&a_shape)?);
1871 let db = child.var_no_grad(&Tensor::zeros(&b_shape)?);
1872 accumulate(&parent_nodes, &mut cgrads, a, da)?;
1873 accumulate(&parent_nodes, &mut cgrads, b, db)?;
1874 } else {
1875 let out_shape = parent_nodes[id.0].shape.clone();
1876 let a_m = get_mirror(mirror, a)?;
1877 let b_m = get_mirror(mirror, b)?;
1878 let (da_full, db_full) = match sop {
1885 ScalarBinaryOp::Add => (g, g),
1886 ScalarBinaryOp::Sub => (g, g.neg()?),
1887 ScalarBinaryOp::Mul => {
1888 let a_bc = a_m.broadcast_to(&out_shape)?;
1889 let b_bc = b_m.broadcast_to(&out_shape)?;
1890 (g.mul(&b_bc)?, g.mul(&a_bc)?)
1891 }
1892 ScalarBinaryOp::Div => {
1893 let a_bc = a_m.broadcast_to(&out_shape)?;
1894 let b_bc = b_m.broadcast_to(&out_shape)?;
1895 let one = child.var_no_grad(&Tensor::scalar(1.0f32));
1896 let factor_a = one.div(&b_bc)?;
1897 let factor_b = a_bc.div(&b_bc)?.div(&b_bc)?.neg()?;
1898 (g.mul(&factor_a)?, g.mul(&factor_b)?)
1899 }
1900 ScalarBinaryOp::Pow => {
1901 let a_bc = a_m.broadcast_to(&out_shape)?;
1902 let b_bc = b_m.broadcast_to(&out_shape)?;
1903 let a_val = materialize_fallible(&parent_nodes, parent_ops, a)?;
1904 let b_val = materialize_fallible(&parent_nodes, parent_ops, b)?;
1905 let (a_bc_val, b_bc_val) =
1906 a_val.broadcast_with(b_val).map_err(AutodiffError::Shape)?;
1907 let mask_b0 = mask_from_pred(&b_bc_val, |v| v == 0.0)?;
1908 let mask_a0 = mask_from_pred(&a_bc_val, |v| v == 0.0)?;
1909 let zeros = child.var_no_grad(&Tensor::zeros(&out_shape)?);
1910 let out_c = get_mirror(mirror, id)?;
1911 let one = child.var_no_grad(&Tensor::scalar(1.0f32));
1912 let b_minus_one = b_bc.sub(&one)?;
1913 let safe_a_for_pow = Var::where_cond(&mask_b0, &one, &a_bc)?;
1934 let safe_a_for_log = Var::where_cond(&mask_a0, &one, &a_bc)?;
1935 let a_pow = safe_a_for_pow.pow(&b_minus_one)?;
1936 let da_factor = Var::where_cond(&mask_b0, &zeros, &b_bc.mul(&a_pow)?)?;
1937 let db_factor = Var::where_cond(
1938 &mask_a0,
1939 &zeros,
1940 &out_c.mul(&safe_a_for_log.log()?)?,
1941 )?;
1942 (g.mul(&da_factor)?, g.mul(&db_factor)?)
1943 }
1944 ScalarBinaryOp::Maximum | ScalarBinaryOp::Minimum => {
1945 let a_val = materialize_fallible(&parent_nodes, parent_ops, a)?;
1950 let b_val = materialize_fallible(&parent_nodes, parent_ops, b)?;
1951 let (a_bc_val, b_bc_val) =
1952 a_val.broadcast_with(b_val).map_err(AutodiffError::Shape)?;
1953 let a_favored = matches!(sop, ScalarBinaryOp::Maximum);
1954 let mask_a_wins = pairwise_mask(&a_bc_val, &b_bc_val, |x, y| {
1955 !(x.is_nan() || y.is_nan()) && if a_favored { x > y } else { x < y }
1956 })?;
1957 let mask_b_wins = pairwise_mask(&a_bc_val, &b_bc_val, |x, y| {
1958 !(x.is_nan() || y.is_nan()) && if a_favored { x < y } else { x > y }
1959 })?;
1960 let mask_tie = pairwise_mask(&a_bc_val, &b_bc_val, |x, y| {
1961 !(x.is_nan() || y.is_nan()) && x == y
1962 })?;
1963 let zeros = child.var_no_grad(&Tensor::zeros(&out_shape)?);
1964 let half = child.var_no_grad(&Tensor::scalar(0.5f32));
1965 let half_g = g.mul(&half)?;
1966 let tie_branch = Var::where_cond(&mask_tie, &half_g, &zeros)?;
1967 let da = Var::where_cond(&mask_a_wins, &g, &tie_branch)?;
1968 let db = Var::where_cond(&mask_b_wins, &g, &tie_branch)?;
1969 (da, db)
1970 }
1971 _ => {
1972 return Err(AutodiffError::Backward(format!(
1973 "create_graph: build_cgrads: 未対応の ScalarBinaryOp {sop:?}\
1974 (内部契約違反。scalar_binary_replayable との判定不整合)"
1975 )));
1976 }
1977 };
1978 let da = reduce_bias_grad_var(child, &da_full, &a_shape)?;
1979 let db = reduce_bias_grad_var(child, &db_full, &b_shape)?;
1980 accumulate(&parent_nodes, &mut cgrads, a, da)?;
1981 accumulate(&parent_nodes, &mut cgrads, b, db)?;
1982 }
1983 }
1984 Op::RmsNorm { input, weight, eps } => {
1985 let targets = crate::create_graph_norm::Targets {
1986 input,
1987 weight,
1988 bias: None,
1989 };
1990 let contributions = crate::create_graph_norm::rms_norm_cgrads(
1991 child,
1992 &parent_nodes,
1993 parent_ops,
1994 mirror,
1995 g,
1996 &targets,
1997 eps,
1998 )?;
1999 for (target, v) in contributions {
2000 accumulate(&parent_nodes, &mut cgrads, target, v)?;
2001 }
2002 }
2003 Op::LayerNorm {
2004 input,
2005 weight,
2006 bias,
2007 eps,
2008 } => {
2009 let targets = crate::create_graph_norm::Targets {
2010 input,
2011 weight,
2012 bias,
2013 };
2014 let contributions = crate::create_graph_norm::layer_norm_cgrads(
2015 child,
2016 &parent_nodes,
2017 parent_ops,
2018 mirror,
2019 g,
2020 &targets,
2021 eps,
2022 )?;
2023 for (target, v) in contributions {
2024 accumulate(&parent_nodes, &mut cgrads, target, v)?;
2025 }
2026 }
2027 Op::Conv2d {
2029 input,
2030 weight,
2031 bias,
2032 params,
2033 ..
2034 } => {
2035 let need = conv::ConvNeeds {
2036 input: parent_nodes[input.0].requires_grad,
2037 weight: parent_nodes[weight.0].requires_grad,
2038 bias: bias.is_some_and(|b| parent_nodes[b.0].requires_grad),
2039 };
2040 let bias_shape = bias.map(|b| parent_nodes[b.0].shape.clone());
2041 let r = conv::conv2d_vjp(
2042 child,
2043 &g,
2044 &get_mirror(mirror, input)?,
2045 &get_mirror(mirror, weight)?,
2046 bias_shape.as_deref(),
2047 ¶ms,
2048 need,
2049 )?;
2050 if let Some(d) = r.d_input {
2051 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2052 }
2053 if let Some(d) = r.d_weight {
2054 accumulate(&parent_nodes, &mut cgrads, weight, d)?;
2055 }
2056 if let (Some(b), Some(d)) = (bias, r.d_bias) {
2057 accumulate(&parent_nodes, &mut cgrads, b, d)?;
2058 }
2059 }
2060 Op::ConvTranspose2d {
2061 input,
2062 weight,
2063 bias,
2064 params,
2065 } => {
2066 let need = conv::ConvNeeds {
2067 input: parent_nodes[input.0].requires_grad,
2068 weight: parent_nodes[weight.0].requires_grad,
2069 bias: bias.is_some_and(|b| parent_nodes[b.0].requires_grad),
2070 };
2071 let bias_shape = bias.map(|b| parent_nodes[b.0].shape.clone());
2072 let r = conv::conv_transpose2d_vjp(
2073 child,
2074 &g,
2075 &get_mirror(mirror, input)?,
2076 &get_mirror(mirror, weight)?,
2077 bias_shape.as_deref(),
2078 ¶ms,
2079 need,
2080 )?;
2081 if let Some(d) = r.d_input {
2082 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2083 }
2084 if let Some(d) = r.d_weight {
2085 accumulate(&parent_nodes, &mut cgrads, weight, d)?;
2086 }
2087 if let (Some(b), Some(d)) = (bias, r.d_bias) {
2088 accumulate(&parent_nodes, &mut cgrads, b, d)?;
2089 }
2090 }
2091 Op::Unfold { input, params } => {
2092 let input_shape = parent_nodes[input.0].shape.clone();
2093 let d = conv::unfold_vjp(&g, &input_shape, ¶ms)?;
2094 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2095 }
2096 Op::Fold { input, params } => {
2097 let d = conv::fold_vjp(&g, ¶ms)?;
2098 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2099 }
2100 Op::MaxPool2d { input, index } | Op::MaxPool3d { input, index } => {
2102 if parent_nodes[input.0].requires_grad {
2103 let input_shape = parent_nodes[input.0].shape.clone();
2104 let d = pool::max_pool_vjp(child, &g, &input_shape, &index)?;
2105 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2106 }
2107 }
2108 Op::AvgPool2d {
2109 input,
2110 params,
2111 count_include_pad,
2112 } => {
2113 if parent_nodes[input.0].requires_grad {
2114 let input_shape = parent_nodes[input.0].shape.clone();
2115 let d = pool::avg_pool2d_cgrad(
2116 child,
2117 &g,
2118 &input_shape,
2119 ¶ms,
2120 count_include_pad,
2121 )?;
2122 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2123 }
2124 }
2125 Op::AvgPool3d {
2126 input,
2127 params,
2128 count_include_pad,
2129 } => {
2130 if parent_nodes[input.0].requires_grad {
2131 let input_shape = parent_nodes[input.0].shape.clone();
2132 let d = pool::avg_pool3d_cgrad(
2133 child,
2134 &g,
2135 &input_shape,
2136 ¶ms,
2137 count_include_pad,
2138 )?;
2139 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2140 }
2141 }
2142 Op::AdaptiveAvgPool2d { input } => {
2143 if parent_nodes[input.0].requires_grad {
2144 let input_shape = parent_nodes[input.0].shape.clone();
2145 let d = pool::adaptive_avg_pool2d_cgrad(child, &g, &input_shape)?;
2146 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2147 }
2148 }
2149 Op::Gather { input, dim, index }
2151 | Op::Sort { input, dim, index }
2152 | Op::Topk { input, dim, index }
2153 | Op::Cummax { input, dim, index }
2154 | Op::Cummin { input, dim, index } => {
2155 if parent_nodes[input.0].requires_grad {
2156 let input_shape = parent_nodes[input.0].shape.clone();
2157 let d = index_cum::gather_like_vjp(child, &g, &input_shape, dim, &index)?;
2158 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2159 }
2160 }
2161 Op::Pad { input, pads, .. } => {
2163 if parent_nodes[input.0].requires_grad {
2164 let input_shape = parent_nodes[input.0].shape.clone();
2165 let d = pad_mask::pad_cgrad(&g, &input_shape, &pads)?;
2166 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2167 }
2168 }
2169 Op::MaskedFill { input, mask, .. } => {
2170 if parent_nodes[input.0].requires_grad {
2171 let d = pad_mask::masked_fill_cgrad(child, &g, &f32_mask_to_bool(&mask)?)?;
2172 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2173 }
2174 }
2175 Op::PadMode { input, pads, mode } => {
2176 if parent_nodes[input.0].requires_grad {
2177 let input_shape = parent_nodes[input.0].shape.clone();
2178 let d = pad_mask::pad_mode_cgrad(child, &g, &input_shape, &pads, mode)?;
2179 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2180 }
2181 }
2182 Op::Embedding {
2183 weight,
2184 index,
2185 padding_idx,
2186 } => {
2187 if parent_nodes[weight.0].requires_grad {
2188 let weight_shape = parent_nodes[weight.0].shape.clone();
2189 let d =
2190 index_cum::embedding_vjp(child, &g, &weight_shape, &index, padding_idx)?;
2191 accumulate(&parent_nodes, &mut cgrads, weight, d)?;
2192 }
2193 }
2194 Op::Scatter {
2195 input,
2196 dim,
2197 index,
2198 src,
2199 reduce,
2200 } => {
2201 let input_shape = parent_nodes[input.0].shape.clone();
2202 let src_shape = parent_nodes[src.0].shape.clone();
2203 let r = index_cum::scatter_vjp(
2204 child,
2205 &g,
2206 &index_cum::ScatterSpec {
2207 input_shape: &input_shape,
2208 src_shape: &src_shape,
2209 dim,
2210 index: &index,
2211 reduce,
2212 },
2213 parent_nodes[input.0].requires_grad,
2214 parent_nodes[src.0].requires_grad,
2215 )?;
2216 if let Some(d) = r.d_input {
2217 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2218 }
2219 if let Some(d) = r.d_src {
2220 accumulate(&parent_nodes, &mut cgrads, src, d)?;
2221 }
2222 }
2223 Op::Cumsum { input, dim } => {
2224 if parent_nodes[input.0].requires_grad {
2225 let d = index_cum::cumsum_cgrad(child, &g, dim)?;
2226 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2227 }
2228 }
2229 Op::Rfft {
2231 input,
2232 n,
2233 dim,
2234 norm,
2235 }
2236 | Op::Irfft {
2237 input,
2238 n,
2239 dim,
2240 norm,
2241 }
2242 | Op::Fft {
2243 input,
2244 n,
2245 dim,
2246 norm,
2247 }
2248 | Op::Ifft {
2249 input,
2250 n,
2251 dim,
2252 norm,
2253 } => {
2254 if parent_nodes[input.0].requires_grad {
2255 let kind = match &node.op {
2256 Op::Rfft { .. } => linalg_fft::FftKind::Rfft,
2257 Op::Irfft { .. } => linalg_fft::FftKind::Irfft,
2258 Op::Fft { .. } => linalg_fft::FftKind::Fft,
2259 _ => linalg_fft::FftKind::Ifft,
2260 };
2261 let input_shape = parent_nodes[input.0].shape.clone();
2262 let d = linalg_fft::fft_vjp(child, &g, kind, &input_shape, n, dim, norm)?;
2263 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2264 }
2265 }
2266 Op::Inv { input } => {
2269 if parent_nodes[input.0].requires_grad {
2270 let a_val = materialize_fallible(&parent_nodes, parent_ops, input)?;
2271 linalg_fft::guard_inv(a_val, &value_of(child, &g)?)?;
2272 let x = get_mirror(mirror, id)?;
2273 let a_shape = parent_nodes[input.0].shape.clone();
2274 let d = linalg_fft::inv_cgrad(child, &g, &x, &a_shape)?;
2275 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2276 }
2277 }
2278 Op::Solve { a, b } => {
2279 if parent_nodes[a.0].requires_grad || parent_nodes[b.0].requires_grad {
2280 let a_val = materialize_fallible(&parent_nodes, parent_ops, a)?;
2281 let b_val = materialize_fallible(&parent_nodes, parent_ops, b)?;
2282 linalg_fft::guard_solve(a_val, b_val, &value_of(child, &g)?)?;
2283 let a_m = get_mirror(mirror, a)?;
2284 let x = get_mirror(mirror, id)?;
2285 let (da, db) = linalg_fft::solve_cgrad(child, &g, &a_m, &x)?;
2286 accumulate(&parent_nodes, &mut cgrads, a, da)?;
2287 accumulate(&parent_nodes, &mut cgrads, b, db)?;
2288 }
2289 }
2290 Op::Det { input } | Op::SlogdetLogAbsDet { input } => {
2291 if parent_nodes[input.0].requires_grad {
2292 let with_det = matches!(node.op, Op::Det { .. });
2293 let a_val = materialize_fallible(&parent_nodes, parent_ops, input)?;
2294 linalg_fft::guard_det_like(a_val, &value_of(child, &g)?, with_det)?;
2295 let a_m = get_mirror(mirror, input)?;
2296 let scale = if with_det {
2297 g.mul(&get_mirror(mirror, id)?)?
2298 } else {
2299 g
2300 };
2301 let d = linalg_fft::det_like_cgrad(child, &scale, &a_m)?;
2302 accumulate(&parent_nodes, &mut cgrads, input, d)?;
2303 }
2304 }
2305 _ => {
2306 return Err(AutodiffError::Backward(
2307 "create_graph: build_cgrads: supports_create_graph() が true の未対応 Op\
2308 (内部契約違反)"
2309 .into(),
2310 ));
2311 }
2312 }
2313 }
2314
2315 Ok(cgrads)
2316}
2317
2318fn accumulate<'c>(
2327 parent_nodes: &[TapeNode],
2328 cgrads: &mut [Option<Var<'c>>],
2329 target: NodeId,
2330 contribution: Var<'c>,
2331) -> Result<(), AutodiffError> {
2332 if !parent_nodes[target.0].requires_grad {
2333 return Ok(());
2334 }
2335 match cgrads[target.0] {
2336 Some(existing) => {
2337 cgrads[target.0] = Some(existing.add(&contribution)?);
2338 }
2339 None => {
2340 cgrads[target.0] = Some(contribution);
2341 }
2342 }
2343 Ok(())
2344}
2345
2346fn reduce_to<'c>(
2361 child: &'c Tape,
2362 v: &Var<'c>,
2363 target_shape: &[usize],
2364) -> Result<Var<'c>, AutodiffError> {
2365 let v_shape = v.shape();
2366 if v_shape == target_shape {
2367 return Ok(*v);
2368 }
2369 debug_assert!(
2370 v_shape.len() >= target_shape.len(),
2371 "reduce_to: broadcast 後 shape の rank は入力 rank 以上のはず(契約違反)"
2372 );
2373 let rank_diff = v_shape.len().saturating_sub(target_shape.len());
2374 let mut padded_target = vec![1usize; rank_diff];
2375 padded_target.extend_from_slice(target_shape);
2376
2377 let mut cur = *v;
2378 let mut cur_shape = v_shape;
2379 for axis in 0..cur_shape.len() {
2380 if padded_target[axis] == 1 && cur_shape[axis] != 1 {
2381 let axis_len = cur_shape[axis];
2382 if axis_len == 0 {
2383 let mut reduced_shape = cur_shape.clone();
2394 reduced_shape[axis] = 1;
2395 cur = child.var_no_grad(&Tensor::zeros(&reduced_shape)?);
2396 cur_shape[axis] = 1;
2397 continue;
2398 }
2399 let mut acc = cur.narrow(axis, 0, 1)?;
2404 for a in 1..axis_len {
2405 let slice = cur.narrow(axis, a, 1)?;
2406 acc = acc.add(&slice)?;
2407 }
2408 cur = acc;
2409 cur_shape[axis] = 1;
2410 }
2411 }
2412 if cur_shape == target_shape {
2413 Ok(cur)
2414 } else {
2415 cur.contiguous()?.reshape(target_shape)
2419 }
2420}
2421
2422fn value_of<'c>(child: &'c Tape, v: &Var<'c>) -> Result<Tensor<f32>, AutodiffError> {
2425 let nodes = child.nodes.borrow();
2426 Ok(materialize_fallible(&nodes, child.ops(), v.node_id())?.clone())
2427}
2428
2429fn reduce_bias_grad_var<'c>(
2476 child: &'c Tape,
2477 g: &Var<'c>,
2478 target_shape: &[usize],
2479) -> Result<Var<'c>, AutodiffError> {
2480 let g_shape = g.shape();
2481 let is_row_axis_reduction = g_shape.len() == 2
2482 && target_shape.last() == Some(&g_shape[1])
2483 && target_shape[..target_shape.len().saturating_sub(1)]
2484 .iter()
2485 .all(|&d| d == 1);
2486 if !is_row_axis_reduction {
2487 return reduce_to(child, g, target_shape);
2488 }
2489 let g_val = {
2490 let nodes = child.nodes.borrow();
2491 materialize_fallible(&nodes, child.ops(), g.node_id())?.clone()
2492 };
2493 let reduced = crate::eval::reduce_bias_grad_rows(&g_val);
2494 let value = Tensor::new(reduced, &[g_shape[1]])?;
2495 let id = child.push_eager(
2496 Op::Sum {
2497 input: g.node_id(),
2498 dim: Some(0),
2499 },
2500 value,
2501 );
2502 Var::from_raw(child, id).reshape(target_shape)
2503}
2504
2505fn positive_mask(t: &Tensor<f32>) -> Result<Tensor<bool>, AutodiffError> {
2510 mask_from_pred(t, |v| v > 0.0)
2511}
2512
2513fn mask_from_pred(
2519 t: &Tensor<f32>,
2520 pred: impl Fn(f32) -> bool,
2521) -> Result<Tensor<bool>, AutodiffError> {
2522 let contiguous = t.contiguous();
2523 let data: Vec<bool> = contiguous
2524 .as_slice()
2525 .map(|s| s.iter().map(|&v| pred(v)).collect())
2526 .unwrap_or_default();
2527 Tensor::new(data, contiguous.shape()).map_err(AutodiffError::from)
2528}
2529
2530fn pairwise_mask(
2537 a: &Tensor<f32>,
2538 b: &Tensor<f32>,
2539 pred: impl Fn(f32, f32) -> bool,
2540) -> Result<Tensor<bool>, AutodiffError> {
2541 let a_c = a.contiguous();
2542 let b_c = b.contiguous();
2543 let data: Vec<bool> = match (a_c.as_slice(), b_c.as_slice()) {
2544 (Some(sa), Some(sb)) => sa
2545 .iter()
2546 .zip(sb.iter())
2547 .map(|(&x, &y)| pred(x, y))
2548 .collect(),
2549 _ => Vec::new(),
2550 };
2551 Tensor::new(data, a_c.shape()).map_err(AutodiffError::from)
2552}
2553
2554fn loss_scale_bcast<'c>(
2560 child: &'c Tape,
2561 g: &Var<'c>,
2562 reduction: crate::var::Reduction,
2563 mse: bool,
2564 n: usize,
2565 shape: &[usize],
2566) -> Result<Var<'c>, AutodiffError> {
2567 let g2 = if mse {
2568 g.mul(&child.var_no_grad(&Tensor::scalar(2.0f32)))?
2569 } else {
2570 *g
2571 };
2572 let s = match reduction {
2573 crate::var::Reduction::Mean => g2.div(&child.var_no_grad(&Tensor::scalar(n as f32)))?,
2574 crate::var::Reduction::Sum => g2,
2575 };
2576 s.broadcast_to(shape)
2577}
2578
2579fn class_sample_count(shape: &[usize], class_dim: usize) -> usize {
2582 let outer: usize = shape[..class_dim].iter().product();
2583 let inner: usize = shape[class_dim + 1..].iter().product();
2584 outer * inner
2585}
2586
2587fn onehot_positions(targets: &Tensor<i32>, shape: &[usize], class_dim: usize) -> Vec<usize> {
2591 let outer: usize = shape[..class_dim].iter().product();
2592 let axis_len = shape[class_dim];
2593 let inner: usize = shape[class_dim + 1..].iter().product();
2594 let data = crate::eval::dense_vec_i32(targets);
2595 let mut out = Vec::with_capacity(outer * inner);
2596 for o in 0..outer {
2597 for i in 0..inner {
2598 let Some(&t) = data.get(o * inner + i) else {
2599 continue;
2600 };
2601 if t >= 0 && (t as usize) < axis_len {
2602 out.push((o * axis_len + t as usize) * inner + i);
2603 }
2604 }
2605 }
2606 out
2607}
2608
2609fn onehot_f32(
2611 targets: &Tensor<i32>,
2612 shape: &[usize],
2613 class_dim: usize,
2614) -> Result<Tensor<f32>, AutodiffError> {
2615 let mut data = vec![0f32; shape.iter().product()];
2616 for idx in onehot_positions(targets, shape, class_dim) {
2617 data[idx] = 1.0;
2618 }
2619 Tensor::new(data, shape).map_err(AutodiffError::from)
2620}
2621
2622fn onehot_mask(
2624 targets: &Tensor<i32>,
2625 shape: &[usize],
2626 class_dim: usize,
2627) -> Result<Tensor<bool>, AutodiffError> {
2628 let mut data = vec![false; shape.iter().product()];
2629 for idx in onehot_positions(targets, shape, class_dim) {
2630 data[idx] = true;
2631 }
2632 Tensor::new(data, shape).map_err(AutodiffError::from)
2633}
2634
2635fn f32_mask_to_bool(cond: &Tensor<f32>) -> Result<Tensor<bool>, AutodiffError> {
2640 mask_from_pred(cond, |v| v != 0.0)
2641}
2642
2643#[cfg(test)]
2644mod contiguous_tests {
2645 use super::*;
2646
2647 #[test]
2662 fn contiguous_replay_and_vjp_are_identity() {
2663 let tape = Tape::new_with_ops(crate::default_ops::naive_ops());
2664 let child = Tape::new_with_ops(crate::default_ops::naive_ops());
2665 let x = tape.var(&Tensor::new(vec![1.0f32, -2.0, 0.5, 3.0], &[2, 2]).unwrap());
2666 let y = x.permute(&[1, 0]).unwrap().contiguous().unwrap();
2667 let loss = y.mul(&y).unwrap().sum(None).unwrap();
2668
2669 let cg = tape
2670 .backward_create_graph(&loss, &child)
2671 .expect("backward_create_graph に失敗(Op::Contiguous は対象)");
2672 let gx = cg
2673 .grad(&x)
2674 .expect("grad: クロステープ検査は通るはず")
2675 .expect("x は loss に到達するはず(1 階勾配)");
2676
2677 let expected = cg.first_order().get(&x).unwrap().unwrap();
2680 assert_eq!(gx.to_tensor().as_slice(), expected.as_slice());
2681
2682 let cx = cg.child_var(&x).unwrap().unwrap();
2684 let row0 = child
2685 .backward(
2686 &gx.narrow(0, 0, 1)
2687 .unwrap()
2688 .narrow(1, 0, 1)
2689 .unwrap()
2690 .sum(None)
2691 .unwrap(),
2692 )
2693 .unwrap();
2694 let d2 = row0.get(&cx).unwrap().unwrap();
2695 assert_eq!(d2.get(&[0, 0]).unwrap(), 2.0);
2696 }
2697}
2698
2699#[cfg(test)]
2718mod maximum_minimum_tests {
2719 use super::*;
2720
2721 fn maximum_probe<'t>(tape: &'t Tape, x0: &[f32; 3]) -> (Tape, Var<'t>, Var<'t>) {
2722 let x = tape.var(&Tensor::new(x0.to_vec(), &[3]).unwrap());
2725 let a = x.mul(&x).unwrap();
2726 let b_data = vec![0.2f32, 30.0, x0[2] * x0[2]];
2727 let b = tape.var_no_grad(&Tensor::new(b_data, &[3]).unwrap());
2728 let loss = a
2729 .scalar_binary(&b, ScalarBinaryOp::Maximum)
2730 .unwrap()
2731 .sum(None)
2732 .unwrap();
2733 let child = Tape::new_with_ops(crate::default_ops::naive_ops());
2734 (child, x, loss)
2735 }
2736
2737 #[test]
2738 fn hessian_maximum_win_lose_tie_matches_expected_curvature() {
2739 let x0 = [3.0f32, 0.1, -2.0];
2740 let tape = Tape::new_with_ops(crate::default_ops::naive_ops());
2741 let (child, x, loss) = maximum_probe(&tape, &x0);
2742
2743 let cg = tape
2744 .backward_create_graph(&loss, &child)
2745 .expect("backward_create_graph に失敗(Op::ScalarBinary::Maximum は対象)");
2746 let gx = cg.grad(&x).unwrap().unwrap();
2747 let cx = cg.child_var(&x).unwrap().unwrap();
2748
2749 let mut diag = [0.0f32; 3];
2750 for (i, d) in diag.iter_mut().enumerate() {
2751 let gi = gx.narrow(0, i, 1).unwrap().sum(None).unwrap();
2752 let row = child.backward(&gi).unwrap();
2753 let d2 = row.get(&cx).unwrap().unwrap();
2754 *d = d2.get(&[i]).unwrap();
2755 }
2756 assert_eq!(diag[0], 2.0, "index0(a 勝ち)は a=x^2 の曲率 2 のはず");
2757 assert_eq!(
2758 diag[1], 0.0,
2759 "index1(b 勝ち。b は x 非依存)は曲率 0 のはず"
2760 );
2761 assert_eq!(diag[2], 1.0, "index2(タイ。0.5 分割)は 0.5*2=1 のはず");
2762 }
2763
2764 #[test]
2765 fn hessian_minimum_win_lose_tie_matches_expected_curvature() {
2766 let x0 = [0.1f32, 3.0, -2.0];
2770 let tape = Tape::new_with_ops(crate::default_ops::naive_ops());
2771 let x = tape.var(&Tensor::new(x0.to_vec(), &[3]).unwrap());
2772 let a = x.mul(&x).unwrap();
2773 let b_data = vec![30.0f32, 0.2, x0[2] * x0[2]];
2774 let b = tape.var_no_grad(&Tensor::new(b_data, &[3]).unwrap());
2775 let loss = a
2776 .scalar_binary(&b, ScalarBinaryOp::Minimum)
2777 .unwrap()
2778 .sum(None)
2779 .unwrap();
2780 let child = Tape::new_with_ops(crate::default_ops::naive_ops());
2781
2782 let cg = tape
2783 .backward_create_graph(&loss, &child)
2784 .expect("backward_create_graph に失敗(Op::ScalarBinary::Minimum は対象)");
2785 let gx = cg.grad(&x).unwrap().unwrap();
2786 let cx = cg.child_var(&x).unwrap().unwrap();
2787
2788 let mut diag = [0.0f32; 3];
2789 for (i, d) in diag.iter_mut().enumerate() {
2790 let gi = gx.narrow(0, i, 1).unwrap().sum(None).unwrap();
2791 let row = child.backward(&gi).unwrap();
2792 let d2 = row.get(&cx).unwrap().unwrap();
2793 *d = d2.get(&[i]).unwrap();
2794 }
2795 assert_eq!(diag[0], 2.0, "index0(a 勝ち)は a=x^2 の曲率 2 のはず");
2796 assert_eq!(
2797 diag[1], 0.0,
2798 "index1(b 勝ち。b は x 非依存)は曲率 0 のはず"
2799 );
2800 assert_eq!(diag[2], 1.0, "index2(タイ。0.5 分割)は 0.5*2=1 のはず");
2801 }
2802
2803 #[test]
2804 fn maximum_nan_element_grad_is_zero() {
2805 let tape = Tape::new_with_ops(crate::default_ops::naive_ops());
2815 let child = Tape::new_with_ops(crate::default_ops::naive_ops());
2816 let x = tape.var(&Tensor::new(vec![f32::NAN, 1.0, 2.0], &[3]).unwrap());
2817 let b = tape.var_no_grad(&Tensor::new(vec![1.0f32, 2.0, 0.5], &[3]).unwrap());
2818 let loss = x
2819 .scalar_binary(&b, ScalarBinaryOp::Maximum)
2820 .unwrap()
2821 .sum(None)
2822 .unwrap();
2823
2824 let cg = tape
2825 .backward_create_graph(&loss, &child)
2826 .expect("backward_create_graph に失敗");
2827 let gx = cg.grad(&x).unwrap().unwrap();
2828 let g0 = gx.to_tensor().get(&[0]).expect("shape 範囲内のはず");
2829 assert_eq!(g0, 0.0, "NaN 要素の 1 階勾配は 0 のはず");
2830
2831 let expected = cg.first_order().get(&x).unwrap().unwrap();
2840 assert_eq!(
2841 gx.to_tensor().as_slice(),
2842 expected.as_slice(),
2843 "子テープ上の 1 階勾配が既存 Tape::backward 経路と一致しない"
2844 );
2845 }
2846}
2847
2848#[cfg(test)]
2852mod new_2649_tests {
2853 use super::*;
2854
2855 #[test]
2856 fn gelu_family_kinds_are_replayable() {
2857 assert!(scalar_unary_replayable(ScalarUnaryOp::Gelu));
2859 assert!(scalar_unary_replayable(ScalarUnaryOp::GeluTanh));
2860 assert!(scalar_unary_replayable(ScalarUnaryOp::Erf));
2861 }
2862
2863 #[test]
2864 fn new_2649_kinds_are_not_replayable() {
2865 for op in [
2866 ScalarUnaryOp::Selu,
2867 ScalarUnaryOp::Celu { alpha: 1.0 },
2868 ScalarUnaryOp::Softsign,
2869 ScalarUnaryOp::Hardsigmoid,
2870 ScalarUnaryOp::LogSigmoid,
2871 ] {
2872 assert!(!scalar_unary_replayable(op), "{op:?}");
2873 }
2874 }
2875}
2876#[cfg(test)]
2881mod conv_low_precision_tests {
2882 use super::*;
2883 use fandhe_ai_tensor_core::ScalarDType;
2884
2885 #[test]
2886 fn low_precision_conv2d_is_rejected_and_child_stays_empty() {
2887 let tape = Tape::new_with_ops(crate::test_support::low_precision_test_ops());
2888 let child = Tape::new_with_ops(crate::default_ops::naive_ops());
2889 let x = tape.var(&Tensor::full(&[1, 1, 3, 3], 0.5f32).unwrap());
2890 let w = tape.var(&Tensor::full(&[1, 1, 2, 2], 0.25f32).unwrap());
2891 let y = x
2892 .conv2d_low_precision(&w, None, [1, 1], [0, 0], [1, 1], 1, ScalarDType::F16)
2893 .unwrap();
2894 let loss = y.sum(None).unwrap();
2895 let err = tape
2896 .backward_create_graph(&loss, &child)
2897 .expect_err("低精度 conv2d は fail-closed で拒否されるはず");
2898 assert!(matches!(err, AutodiffError::Backward(_)), "{err:?}");
2899 assert!(child.is_empty(), "入口検査は child へ書き込む前に拒否する");
2900 }
2901}