Module Tensor.Scheme

Shape schemes (gh-ocannl-404)

The low-level, explicit layer of shape schemes, which the %op integration (gh-ocannl-1751) wraps. A scheme is the shape constraints a tensor function imposes, inferred once by a trace and instantiated at each application:

What is shared and what is per-application follows ML's value restriction: the dimensions of a tensor created before the trace -- a shared parameter, lifted to the function's unit-parameter scope -- stay fixed across applications, and everything the function creates (its intermediates, a parameter created inside it, a capture reference made inside it) is fresh per application. So create the shared parameters before Scheme.trace. The design note is the gh-ocannl-404 issue's reconciled-plan comment; docs/shape_inference.md has the summary.

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

to_string s prints s's boundary: <label> : <arguments> => <results>, each shape as batch|input->output under the trace's solution, with ..v.. for a row variable and () for an empty row. Variables are named by first appearance, so the text does not depend on the session's id counters: a generalized variable (renamed per application) as 'a, 'b, ...; a fixed one (shared with what the function closes over, such as a parameter lifted to its () scope) as '_a, '_b, ..., the way OCaml prints weak type variables. Then, one per indented line, the broadcasting inequalities r ⊑ o (the result side refines the operand side) among the boundary's variables and the fixed ones, and the lower bounds of the boundary's variables, under the trace's final environment, with the function's intermediate variables eliminated along monotone chains. Other constraints (totals, concatenations) are not printed. test/operations/reusable_tensor_programs.ml shows one.

val trace : label:Base.string -> Base.int -> (tensor Base.list -> tensor Base.list) -> t

trace ~label n f applies f to n tracing tensors whose shapes are fresh row variables, inside with_rewind, and returns the scheme of the constraints it registered over them and over its results (a flat list: flatten a structure). The trace leaves nothing in the session. Raises Row.Shape_error prefixed "tracing the shape scheme <label>: " on a shape error in f (including one only the end-of-trace check of f's trailing Shape.set_dim and the like finds), when f forces dims or projections (inference cannot close during a trace), and for what a scheme cannot carry: a precision-dependent total-elements constraint (random number generation inside f), or Shape.set_dim and the like on a capture reference made outside f. A capture reference made outside f that f's operations bind is an input of the scheme (gh-ocannl-1791): each application relates it to the reference's state then, and instantiate binds it, as the body would.

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

instantiate s ~args ~results adds s's constraints for one application of its function to args with results (fresh, inferred shapes, or the shapes of tensors computing the results), without running the function. An argument-dependent shape error raises Row.Shape_error prefixed "applying the shape scheme <label> to <application>: " and taints the session, as a failed construction does; the constraints carry the application's name, so a failure at the closure names it too. The next closure applies the solution to args and results and closes them. Valid across closures, unless the scheme depends on open state no shape it used holds (then refused after a closure: trace again).

Automatic schemes of %op functions (gh-ocannl-1751)

%op gives a function a scheme when every parameter of the function it returns is a tensor parameter (docs/syntax_extensions.md says which qualify). The closure built at the () application -- or, for a function without a () stage, its first application -- traces the scheme. How each application uses it is the op_shape_schemes setting:

A shape error that does not depend on the arguments raises at the trace, prefixed "tracing the shape scheme <function>: "; one that does raises at the application, before the body runs, prefixed "applying the shape scheme <function> to <arguments>: ". A trace a scheme cannot carry (random number generation, setting a capture reference made outside the function, forcing dims, or a body that cannot run on tracers) leaves the function unschemed, as it was before, once its rewind succeeded, and says so on stderr -- once per function and reason in a process (gh-ocannl-1791). A capture reference made outside the function that its operations bind is a scheme input, checked against each application's state of it. Ids are unaffected: traces and checks allocate tensor ids, node uids and shape-variable ids from reserved ranges.

Option-valued tensor parameters (gh-ocannl-1792), annotated ~(p : Tensor.t option) or ?(p : Tensor.t option): the function gets one scheme per presence pattern -- which of them are Some -- each traced, checked and refused on its own, under the label <function> [p, no q] (just <function> without such parameters). The closure built at () traces every pattern; a function without a () stage traces every pattern at its first application; afterwards each pattern re-traces on its own when stale.

type handle

A function's scheme slots: one handle per closure, one slot per presence pattern.

val status : handle -> Base.string

"pending", "schemed", or "unschemed: <reason>"; with option-valued parameters, one "<label>: <status>" per pattern, separated by "; ".

type event =
  1. | Traced of Base.string
    (*

    A function's scheme (a pattern's, by its label) was traced.

    *)
  2. | Unschemed of Base.string * Base.string
    (*

    A function (a pattern, by its label) was left unschemed, with the reason.

    *)
  3. | Checked of Base.string
    (*

    An application was checked against its pattern's scheme.

    *)
val registration_counts : Base.unit -> Shape.registration_counts

Process-wide registration statistics since the last reset (gh-ocannl-1752's attribution): see Shape.registration_counts.

val reset_registration_counts : Base.unit -> Base.unit
val registration_times : Base.unit -> Shape.registration_times

Process-wide shape-registration wall time since the last reset (gh-ocannl-1752): see Shape.registration_times.

val reset_registration_times : Base.unit -> Base.unit
val check_definition : name:Base.string -> (Base.unit -> Base.unit) -> Base.string Base.list

check_definition ~name (fun () -> ignore (block conf ())) runs a trial application of a %op block definition inside a window that is always rewound and allocates no session ids, so the () traces of the block and of the sub-blocks it builds run, and with them any shape error that does not depend on the arguments, prefixed "checking the definition <name>: ". Returns the labels of the schemes traced (every presence pattern of a closure built at () in the trial: see above), innermost first; the session is unchanged afterwards. A function without a () stage that the trial only builds is not traced. It catches errors that hold for every configuration and those of the trial configuration only; it costs the trial's construction and traces. Prefer it to a bare block conf (), which leaves the trial's parameters in the session.

module For_ppx : sig ... end

The runtime the %op expansion calls. Not for direct use.