def matmul_values
bendtorch-core@0.1.0.0 · src/autograd/tensor_vjp_PROOF.bend:462
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/word_conversion.bend as Conversion
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/layout.bend as Layout
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
import bendtorch-core@0.1.0.0/src/tensor/transpose_spec.bend as Transpose
import bendtorch-core@0.1.0.0/src/tensor/permute_values_contract_LAWS.bend as PermuteLaws
import bendtorch-core@0.1.0.0/src/tensor/contracts/matmul_proof.bend as MatmulProof
462def matmul_values(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, +shape: Shape.Shape,
463 +right_inside: {True{} == Transpose.valid(right, 0, 1) : Bool},
464 +right_plan: Layout.Plan<right, Laws.swapped(right)>,
465 +right_laid: {Layout.permutation(right, Laws.swapped(right)) == Done{right_plan} : Result<&2, &2, Layout.Error, Layout.Plan<right, Laws.swapped(right)>>},
466 +output: Shape.Shape, +rows: U32, +inner: U32, +right_inner: U32, +columns: U32,
467 -le: {Shape.dimensions(shape) == [rows, inner] : List<&2, U32>},
468 -re: {Shape.dimensions(Laws.transposed(right, right_plan)) == [right_inner, columns] : List<&2, U32>},
469 -oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
470 -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
471 +first_planned: {Matrix.plan(shape, Laws.transposed(right, right_plan)) ==
472 Done{Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}} :
473 Result<&2, &2, Matrix.Error, Matrix.Plan<shape, Laws.transposed(right, right_plan)>>},
474 +first_same: {Shape.dimensions(output) == Shape.dimensions(left) : List<&2, U32>},
475 +left_inside: {True{} == Transpose.valid(left, 0, 1) : Bool},
476 +left_plan: Layout.Plan<left, Laws.swapped(left)>,
477 +left_laid: {Layout.permutation(left, Laws.swapped(left)) == Done{left_plan} : Result<&2, &2, Layout.Error, Layout.Plan<left, Laws.swapped(left)>>},
478 +back: Shape.Shape, +back_rows: U32, +back_inner: U32, +back_right_inner: U32, +back_columns: U32,
479 -back_le: {Shape.dimensions(Laws.transposed(left, left_plan)) == [back_rows, back_inner] : List<&2, U32>},
480 -back_re: {Shape.dimensions(shape) == [back_right_inner, back_columns] : List<&2, U32>},
481 -back_oe: {Shape.dimensions(back) == [back_rows, back_columns] : List<&2, U32>},
482 -back_compatible: {U32.to_nat(back_inner) == U32.to_nat(back_right_inner) : Nat},
483 +second_planned: {Matrix.plan(Laws.transposed(left, left_plan), shape) ==
484 Done{Matrix.Plan{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible}} :
485 Result<&2, &2, Matrix.Error, Matrix.Plan<Laws.transposed(left, left_plan), shape>>},
486 +second_same: {Shape.dimensions(back) == Shape.dimensions(right) : List<&2, U32>},
487 +fill: F32, +ops: List<&2, StorageLaws.Op>,
488 -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
489 +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
490 -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat},
491 +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
492 -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}) ->
493 MatmulGoal(a, b, left, right, shape, right_plan, Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible},
494 left_plan, Matrix.Plan{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible},
495 fill, ops, left_fill, left_ops, right_fill, right_ops,
496 Backward.vjp(Backward.Matmul{a, b, left, right},
497 Backward.Operands{Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
498 Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}},
499 Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})):
500 +flipped = Laws.transposed(right, right_plan)
501 +turned = Laws.transposed(left, left_plan)
502 +first : Matrix.Plan<shape, Laws.transposed(right, right_plan)> = Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}
503 +second : Matrix.Plan<Laws.transposed(left, left_plan), shape> =
504 Matrix.Plan{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible}
505 +front_facts : MatmulProof.MatrixFacts<shape, flipped, output, rows, inner, right_inner, columns> =
506 MatmulProof.plan_facts(shape, flipped, first, first_planned)
507 +back_facts : MatmulProof.MatrixFacts<turned, shape, back, back_rows, back_inner, back_right_inner, back_columns> =
508 MatmulProof.plan_facts(turned, shape, second, second_planned)
509 +front_re = facts_right(shape, flipped, output, rows, inner, right_inner, columns, front_facts)
510 +front_oe = MatmulProof.output_dimensions(shape, flipped, output, rows, inner, right_inner, columns, front_facts)
511 +front_compatible = facts_compatible(shape, flipped, output, rows, inner, right_inner, columns, front_facts)
512 +turned_le = facts_left(turned, shape, back, back_rows, back_inner, back_right_inner, back_columns, back_facts)
513 +turned_oe = MatmulProof.output_dimensions(turned, shape, back, back_rows, back_inner, back_right_inner, back_columns, back_facts)
514 -Goal : Backward.Pass -> Type = pass => MatmulGoal(a, b, left, right, shape, right_plan, first, left_plan, second,
515 fill, ops, left_fill, left_ops, right_fill, right_ops, pass)
516 -x = {Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}} : Core.Dynamic}
517 -y = {Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}} : Core.Dynamic}
518 +front = Replayed.snoc(ops, StorageLaws.First{})
519 SumToSizeFoldProof.clone_lift(shape, fill, ops, fits,
520 pair => Goal(Backward.matmul_first(a, b, left, right, x, y, Core.clone_dynamic_result(shape, pair))),
521 g_fits => copy_fits =>
522 -g = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, front), g_fits}} : Core.Dynamic}
523 -copy = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(ops, StorageLaws.Second{})), copy_fits}} : Core.Dynamic}
524 PermuteLaws.transpose_replayed(right, right_fill, right_ops, right_fits, 0, 1, right_inside, right_plan, right_laid,
525 pair => Goal(Backward.matmul_second(b, right, x, g, Backward.left_product(a, left, copy, pair))),
526 y_kept => y_words => y_fits => flipped_ops => flipped_moved => flipped_fits =>
527 -y_back = {Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, y_kept), y_fits}} : Core.Dynamic}
528 Contracts.matmul_replayed(shape, flipped, first, first_planned, fill, Replayed.snoc(ops, StorageLaws.Second{}), copy_fits,
529 0.0, flipped_ops, flipped_fits,
530 result => Goal(Backward.matmul_second(b, right, x, g, (y_back, Backward.multiplied(a, left, result)))),
531 copy_kept => copy_words => copy_kept_fits => flipped_kept => flipped_words => flipped_kept_fits => output_ops => front_cells => output_fits =>
532 -delivered = {Backward.Contribution{a, left, Core.Packed{output, Core.TTensor{Replayed.storage(output, 0.0, output_ops), output_fits}}} :
533 Backward.Contribution}
534 SumToSizeFoldProof.clone_lift(shape, fill, front, g_fits,
535 pair => Goal(Backward.matmul_third(b, right, Done{delivered}, y_back, x, Core.clone_dynamic_result(shape, pair))),
536 last_fits => again_fits =>
537 -last = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(front, StorageLaws.First{})), last_fits}} : Core.Dynamic}
538 -again = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(front, StorageLaws.Second{})), again_fits}} :
539 Core.Dynamic}
540 PermuteLaws.transpose_replayed(left, left_fill, left_ops, left_fits, 0, 1, left_inside, left_plan, left_laid,
541 pair => Goal(Backward.matmul_fourth(Done{delivered}, y_back, last, Backward.right_product(b, right, again, pair))),
542 x_kept => x_words => x_fits => turned_ops => turned_moved => turned_fits =>
543 -x_back = {Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, x_kept), x_fits}} : Core.Dynamic}
544 Contracts.matmul_replayed(turned, shape, second, second_planned, 0.0, turned_ops, turned_fits,
545 fill, Replayed.snoc(front, StorageLaws.Second{}), again_fits,
546 result => Goal(Backward.matmul_fourth(Done{delivered}, y_back, last, (x_back, Backward.multiplied(b, right, result)))),
547 turned_kept => turned_words => turned_kept_fits => again_kept => again_words => again_kept_fits =>
548 back_ops => back_cells => back_fits =>
549 Laws.MatmulPairContract{
550 {==},
551 Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions},
552 Shape.dimensions(output), Shape.dimensions(left), first_same),
553 row => row_below => column => column_below =>
554 +row_below = row_below
555 +column_below = column_below
556 product_cell(output, rows, columns, front_oe, output_ops, output_fits, row, row_below, column, column_below,
557 Contracts.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill,
558 Replayed.snoc(ops, StorageLaws.Second{}), flipped, 0.0, flipped_ops, U32.to_nat(inner)),
559 Laws.right_transposed_factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill, ops,
560 right, right_fill, right_ops, Laws.swapped(right), U32.to_nat(inner)),
561 front_cells(row, row_below, column, column_below),
562 first_factors(row, column, inner, right_inner, columns, shape, fill, ops, right, right_fill, right_ops,
563 right_plan, flipped_ops, flipped_moved, front_re, front_compatible, column_below,
564 U32.to_nat(inner), Conversion.le_reflexive(U32.to_nat(inner)))),
565 {==},
566 Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions},
567 Shape.dimensions(back), Shape.dimensions(right), second_same),
568 row => row_below => column => column_below =>
569 +row_below = row_below
570 +column_below = column_below
571 product_cell(back, back_rows, back_columns, turned_oe, back_ops, back_fits, row, row_below, column, column_below,
572 Contracts.factors(row, column, U32.to_nat(back_inner), U32.to_nat(back_columns), turned, 0.0, turned_ops,
573 shape, fill, Replayed.snoc(front, StorageLaws.Second{}), U32.to_nat(back_inner)),
574 Laws.left_transposed_factors(row, column, U32.to_nat(back_inner), U32.to_nat(back_columns),
575 left, left_fill, left_ops, Laws.swapped(left), shape, fill, ops, U32.to_nat(back_inner)),
576 back_cells(row, row_below, column, column_below),
577 second_factors(row, column, back_rows, back_inner, back_columns, left, left_fill, left_ops, left_plan,
578 turned_ops, turned_moved, turned_le, row_below, shape, fill, ops,
579 U32.to_nat(back_inner), Conversion.le_reflexive(U32.to_nat(back_inner))))}))))))