Module Shape.Scheme

Shape schemes (gh-ocannl-404)

The shape-inference half of Tensor.Scheme: use that. A scheme is the constraints a traced function registers, read off the session before the trace's rewind, with the solution equations of their variables; an instantiation renames the variables the trace minted (the generalized ones), keeps the older ones (shared parameters' and every other pre-existing tensor's: fixed), and relates the scheme's boundary rows to given argument and result shapes. It runs no function body and registers no update step: the next closure applies the solution to the boundary shapes and closes them, and finalizes the scheme's internal rows (renamed) as it finalizes an update step's rows.

type shape := t
type t
type sabotage =
  1. | No_generalization
    (*

    Rename nothing: applications share the trace's variables.

    *)
  2. | Sever_shared
    (*

    Rename every variable and drop the anchors: no application stays tied to pre-existing shapes.

    *)
  3. | No_marks
    (*

    Re-add none of the trace's policy marks (instantiate only).

    *)
  4. | No_internal_rows
    (*

    Register none of the scheme's internal rows with the closure (instantiate only).

    *)
  5. | No_end_validation
    (*

    Skip the end-of-trace solve of the whole window (trace only).

    *)
  6. | No_ref_inputs
    (*

    Relate and bind none of the scheme's capture inputs (instantiate and check).

    *)

Negative controls for the tests: each disables one part of the mechanism.

val label : t -> Base.string
val arity : t -> Base.int
val num_results : t -> Base.int
val to_string : t -> Base.string

See Tensor.Scheme.to_string.

val trace : ?sabotage:sabotage -> var_mark:Base.int -> label:Base.string -> args:shape Base.list -> (Base.unit -> shape Base.list) -> t

trace ~var_mark ~label ~args f runs f, which returns the result shapes, and reads the scheme off the session. args are the tracers' shapes, made after var_mark was read with Row.var_mark (so their variables are generalized) and before this call (so their own constraints are not the scheme's). Must run inside a rewound window (Tensor.with_rewind), which undoes the trace. While f runs, finish_inference raises. Refuses, with Untraceable_scheme on what a scheme cannot carry: a precision-dependent total-elements constraint, a capture reference made outside the window and set or related in it (set_dim and the like), or forcing dims or projections; and Row.Shape_error on a contradiction among the window's constraints (solved at Stage1 once more at the end, so a trailing set_dim is checked too). Both are prefixed "tracing the shape scheme <label>: ".

A capture reference made outside the window that f's operations bind (~capture_dims) is an input of the scheme (gh-ocannl-1791): during the trace it reads as unbound and unsolved, so the scheme does not depend on its state then, and each application relates the variable of its first binding in f to its state at that application -- an open binding, or a solution held across a closure -- as the body's own first binding of it would.

val instantiate : ?sabotage:sabotage -> t -> application:Base.string -> args:shape Base.list -> results:shape Base.list -> Base.unit

Adds the scheme's renamed constraints and the equalities of its boundary rows with args and results (and of its anchored pre-existing shapes with their current rows), solves them at Stage1, and registers the boundary shapes and the internal rows with the next closure. Every added constraint carries an origin naming the scheme and application, so a closure-time failure names them too. Raises Row.Shape_error prefixed "applying the shape scheme <label> to <application>: ": before anything is written, on an arity mismatch, a boundary or anchored shape discarded by a rewind, or a closure since the trace when the scheme depends on open state no anchored shape holds; after registering (tainting the session, as a failed construction does), on a contradiction. The scheme's capture inputs are related to the references' current state and bound, to the variable of their last binding in the trace, so their solutions are published -- after the Stage1 solve and at the closure -- as the body's would be.

val check : ?sabotage:sabotage -> t -> application:Base.string -> args:shape Base.list -> Base.unit

instantiate against args with no results, binding no capture reference (its capture inputs are only related to their current state), then puts back everything it wrote (the session and solver globals, the taint, the rows of the arguments and anchored shapes): an argument-dependent shape error raises as instantiate's does, and the session is unchanged either way. Costs what instantiate does, independent of the open frontier.

val stale : t -> Base.bool

Whether instantiate would refuse the scheme for a closure since its trace (the scheme depends on open state no anchored shape holds): trace it again.