def nt_values
bendtorch-core@0.1.0.0 · src/autograd/tensor_vjp_PROOF.bend:672
import bendtorch-core@0.1.0.0/src/autograd/tensor_vjp_PROOF.bend as TensorVjpPROOF
import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
import bendtorch-core@0.1.0.0/src/autograd/sum_to_size_fold_PROOF.bend as SumToSizeFoldProof
import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward
import bendtorch-core@0.1.0.0/src/autograd/tensor_vjp_LAWS.bend as Laws
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
import bendtorch-core@0.1.0.0/src/tensor/contracts/matmul_proof.bend as MatmulProof
import bendtorch-core@0.1.0.0/src/tensor/contracts/product_proof.bend as ProductProof
672def nt_values(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, +shape: Shape.Shape,
673 +output: Shape.Shape, +rows: U32, +inner: U32, +right_inner: U32, +columns: U32,
674 -le: {Shape.dimensions(shape) == [rows, inner] : List<&2, U32>},
675 -re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
676 -oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
677 -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
678 +first_planned: {Matrix.plan(shape, right) == Done{Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}} :
679 Result<&2, &2, Matrix.Error, Matrix.Plan<shape, right>>},
680 +first_same: {Shape.dimensions(output) == Shape.dimensions(left) : List<&2, U32>},
681 +back: Shape.Shape, +back_rows: U32, +back_inner: U32, +back_right_inner: U32, +back_columns: U32,
682 -back_le: {Shape.dimensions(shape) == Matrix.arranged(Matrix.Swapped{}, back_rows, back_inner) : List<&2, U32>},
683 -back_re: {Shape.dimensions(left) == Matrix.arranged(Matrix.Swapped{}, back_columns, back_right_inner) : List<&2, U32>},
684 -back_oe: {Shape.dimensions(back) == [back_rows, back_columns] : List<&2, U32>},
685 -back_compatible: {U32.to_nat(back_inner) == U32.to_nat(back_right_inner) : Nat},
686 +second_planned: {Matrix.product_plan(Matrix.Swapped{}, Matrix.Swapped{}, shape, left) ==
687 Done{Matrix.Product{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible}} :
688 Result<&2, &2, Matrix.Error, Matrix.Product<Matrix.Swapped{}, Matrix.Swapped{}, shape, left>>},
689 +second_same: {Shape.dimensions(back) == Shape.dimensions(right) : List<&2, U32>},
690 +fill: F32, +ops: List<&2, StorageLaws.Op>,
691 -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
692 +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
693 -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat},
694 +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
695 -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}) ->
696 NTGoal(a, b, left, right, shape, Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible},
697 Matrix.Product{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible},
698 fill, ops, left_fill, left_ops, right_fill, right_ops,
699 Backward.vjp(Backward.MatmulNT{a, b, left, right},
700 Backward.Operands{Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
701 Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}},
702 Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})):
703 +first : Matrix.Plan<shape, right> = Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}
704 +second : Matrix.Product<Matrix.Swapped{}, Matrix.Swapped{}, shape, left> =
705 Matrix.Product{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible}
706 +front_facts : MatmulProof.MatrixFacts<shape, right, output, rows, inner, right_inner, columns> =
707 MatmulProof.plan_facts(shape, right, first, first_planned)
708 +back_facts : ProductProof.ProductFacts<Matrix.Swapped{}, Matrix.Swapped{}, shape, left, back, back_rows, back_inner, back_right_inner, back_columns> =
709 ProductProof.plan_facts(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second, second_planned)
710 +front_oe = MatmulProof.output_dimensions(shape, right, output, rows, inner, right_inner, columns, front_facts)
711 +back_oe_live = ProductProof.output_dimensions(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, back, back_rows, back_inner, back_right_inner,
712 back_columns, back_facts)
713 -Goal : Backward.Pass -> Type = pass => NTGoal(a, b, left, right, shape, first, second,
714 fill, ops, left_fill, left_ops, right_fill, right_ops, pass)
715 -x = {Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}} : Core.Dynamic}
716 -y = {Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}} : Core.Dynamic}
717 +front = Replayed.snoc(ops, StorageLaws.First{})
718 SumToSizeFoldProof.clone_lift(shape, fill, ops, fits,
719 pair => Goal(Backward.nt_first(a, b, left, right, x, y, Core.clone_dynamic_result(shape, pair))),
720 g_fits => copy_fits =>
721 -g = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, front), g_fits}} : Core.Dynamic}
722 -copy = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(ops, StorageLaws.Second{})), copy_fits}} : Core.Dynamic}
723 SumToSizeFoldProof.clone_lift(right, right_fill, right_ops, right_fits,
724 pair => Goal(Backward.nt_second(b, right, x, g, Backward.nt_left(a, left, copy, Core.clone_dynamic_result(right, pair)))),
725 y_fits => flipped_fits =>
726 -y_back = {Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, Replayed.snoc(right_ops, StorageLaws.First{})), y_fits}} : Core.Dynamic}
727 Contracts.matmul_replayed(shape, right, first, first_planned, fill, Replayed.snoc(ops, StorageLaws.Second{}), copy_fits,
728 right_fill, Replayed.snoc(right_ops, StorageLaws.Second{}), flipped_fits,
729 result => Goal(Backward.nt_second(b, right, x, g, (y_back, Backward.multiplied(a, left, result)))),
730 copy_kept => copy_words => copy_kept_fits => flipped_kept => flipped_words => flipped_kept_fits => output_ops => front_cells => output_fits =>
731 -delivered = {Backward.Contribution{a, left, Core.Packed{output, Core.TTensor{Replayed.storage(output, 0.0, output_ops), output_fits}}} :
732 Backward.Contribution}
733 SumToSizeFoldProof.clone_lift(shape, fill, front, g_fits,
734 pair => Goal(Backward.nt_third(b, right, Done{delivered}, y_back, x, Core.clone_dynamic_result(shape, pair))),
735 last_fits => again_fits =>
736 -last = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(front, StorageLaws.First{})), last_fits}} : Core.Dynamic}
737 -again = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(front, StorageLaws.Second{})), again_fits}} :
738 Core.Dynamic}
739 SumToSizeFoldProof.clone_lift(left, left_fill, left_ops, left_fits,
740 pair => Goal(Backward.nt_fourth(Done{delivered}, y_back, last, Backward.nt_right(b, right, again, Core.clone_dynamic_result(left, pair)))),
741 x_fits => turned_fits =>
742 -x_back = {Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, Replayed.snoc(left_ops, StorageLaws.First{})), x_fits}} :
743 Core.Dynamic}
744 Contracts.product_replayed(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second, second_planned,
745 fill, Replayed.snoc(front, StorageLaws.Second{}), again_fits,
746 left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), turned_fits,
747 result => Goal(Backward.nt_fourth(Done{delivered}, y_back, last, (x_back, Backward.multiplied(b, right, result)))),
748 again_kept => again_words => again_kept_fits => turned_kept => turned_words => turned_kept_fits =>
749 back_ops => back_cells => back_fits =>
750 Laws.MatmulPairContract{
751 {==},
752 Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions},
753 Shape.dimensions(output), Shape.dimensions(left), first_same),
754 row => row_below => column => column_below =>
755 +row_below = row_below
756 +column_below = column_below
757 product_cell(output, rows, columns, front_oe, output_ops, output_fits, row, row_below, column, column_below,
758 Contracts.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill,
759 Replayed.snoc(ops, StorageLaws.Second{}), right, right_fill, Replayed.snoc(right_ops, StorageLaws.Second{}),
760 U32.to_nat(inner)),
761 Contracts.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill, ops,
762 right, right_fill, right_ops, U32.to_nat(inner)),
763 front_cells(row, row_below, column, column_below),
764 nt_first_factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill, ops,
765 right, right_fill, right_ops, U32.to_nat(inner))),
766 {==},
767 Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions},
768 Shape.dimensions(back), Shape.dimensions(right), second_same),
769 row => row_below => column => column_below =>
770 +row_below = row_below
771 +column_below = column_below
772 product_cell(back, back_rows, back_columns, back_oe_live, back_ops, back_fits, row, row_below, column, column_below,
773 Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, U32.to_nat(back_inner),
774 U32.to_nat(back_rows), U32.to_nat(back_columns), shape, fill, Replayed.snoc(front, StorageLaws.Second{}),
775 left, left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), U32.to_nat(back_inner)),
776 Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, U32.to_nat(back_inner),
777 U32.to_nat(back_rows), U32.to_nat(back_columns), shape, fill, ops, left, left_fill, left_ops, U32.to_nat(back_inner)),
778 back_cells(row, row_below, column, column_below),
779 nt_second_factors(row, column, U32.to_nat(back_inner), U32.to_nat(back_rows), U32.to_nat(back_columns),
780 shape, fill, ops, left, left_fill, left_ops, U32.to_nat(back_inner)))}))))))