def DescendWith

bendtorch-core@0.1.0.0 · src/autograd/sgd_LAWS.bend:34

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/elementwise.bend as Elementwise
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/tensor/contracts/binary_LAWS.bend as Binary
import bendtorch-core@0.1.0.0/src/autograd/gradient.bend as Gradient
import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward
import bendtorch-core@0.1.0.0/src/autograd/sgd.bend as Sgd

34# What the transport form of sgd_step_values asks of P: P holds for the entry
35# with its node, shape and saved values, whose value is replayed over +0.0
36# with any history whose words are the cells of Descend{lr} over the old
37# value p and the gradient g, and whose slot still holds g with its words,
38# next to Done.
39def DescendWith(+lr: F32, node: Backward.Node, +shape: Shape.Shape, saved: Backward.Saved, +plan: Shape.BroadcastPlan,
40  +fill: F32, +ops: List<&2, StorageLaws.Op>,
41  +g_shape: Shape.Shape, +g_fill: F32, +g_ops: List<&2, StorageLaws.Op>,
42  -P: Backward.Entry & Result<&2, &2, Sgd.Error, Unit> -> Type) -> Type:
43  @+kept: List<&2, StorageLaws.Op> ->
44  @kept_words: {Contracts.words(shape, fill, kept, Shape.logical_numel(Shape.dimensions(shape))) ==
45    Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))) : List<&2, F32>} ->
46  @-kept_fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, kept)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} ->
47  @+g_kept: List<&2, StorageLaws.Op> ->
48  @g_words: {Contracts.words(g_shape, g_fill, g_kept, Shape.logical_numel(Shape.dimensions(g_shape))) ==
49    Contracts.words(g_shape, g_fill, g_ops, Shape.logical_numel(Shape.dimensions(g_shape))) : List<&2, F32>} ->
50  @-g_fits: {Lifecycle.logical_length(Replayed.storage(g_shape, g_fill, g_kept)) == Shape.logical_numel(Shape.dimensions(g_shape)) : Nat} ->
51  @+output_ops: List<&2, StorageLaws.Op> ->
52  @values: {Contracts.words(Shape.output_shape(plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) ==
53    Binary.cells(Elementwise.Descend{lr}, plan, shape, fill, ops, g_shape, g_fill, g_ops,
54      Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) : List<&2, F32>} ->
55  @-output_fits: {Lifecycle.logical_length(Replayed.storage(Shape.output_shape(plan), 0.0, output_ops)) ==
56    Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat} ->
57    P((Backward.Entry{node, shape,
58      Core.Packed{Shape.output_shape(plan), Core.TTensor{Replayed.storage(Shape.output_shape(plan), 0.0, output_ops), output_fits}},
59      saved, Gradient.Held{Core.Packed{g_shape, Core.TTensor{Replayed.storage(g_shape, g_fill, g_kept), g_fits}}}}, Done{Unit{}}))