Module Ocannl_tensor.Shape

Private (gh-ocannl-1506): fields are read and matched freely, but only Shape's guarded setters write them and only Shape constructs a shape, so the closed-shape guard cannot be bypassed. The privacy is shallow: a padding array read off a shape is still a mutable array.

Private (gh-ocannl-1573): read and matched freely, written only by Shape -- through guarded setters that refuse a reference a rewind discarded and count the write as a session mutation -- and built only by get_variable_ref. The Ir.Indexing.variable_ref it carries is solved only by Shape too: Ir.Indexing.set_solved_dim and set_solved_sym take the process's one Ir.Indexing.shape_writer, which Shape claims as it initializes (gh-ocannl-1624).

Returns a fully unset variable reference with the given label.

A reference passed through ~capture_dims to several operations names one dimension, in whichever order inference solves them (gh-ocannl-1545): within one inference (up to the next finish_inference) the operations' dimension variables are equated, as are their row variables, and a dimension variable with a row variable has the row's product. Across a closure the reference carries its solved value, for a row variable the product of its dimensions (see set_dim).

A reference whose closure recorded no value -- a row holding a symbolic extent has no numeric product -- has lost its constraint: capturing it again, or relating it by set_dim, set_sym_dim, set_equal or set_scale, raises Row.Shape_error naming the closure (gh-ocannl-1655).

Sets the dimension resp. total elements of the dim resp. row variable reference to the given value. This will propagate through shape inference.

For row variables, this means the product of the dimensions, via the Total_elems constraint.

Sets the dim variable reference to a symbolic extent (gh-490): during shape inference the axis is a rigid constant equal only to itself, and it is materialized at its declared maximum extent (static_range, which is required and must be positive) for Tnode dims and projections.

Row variables are not supported: a symbolic total-elements constraint would require symbolic arithmetic (raises Row.Shape_error). Conflicts with a prior set_dim, or a different symbol set on the same reference, also raise Row.Shape_error.

Sets the two variable references to be equal (in some sense). This will propagate through shape inference.

When both references are dimension variables or both are row variables, this means they are precisely equal: two rows get the same axes, not only the same product, whatever either reference has been solved to so far (gh-ocannl-1617). When one is a dimension variable and the other is a row variable, this means they have the same number of total elements.

A row reference bound by an inference that has since closed keeps only its row's product (solved_dim), so a Row-Row call involving it raises Row.Shape_error: relate the rows before the closure.

An unbound reference (from get_variable_ref, not yet passed to ~capture_dims) paired with a row reference takes or gives a number, the row's product. Binding it to a row afterwards would leave the two rows tied by product alone, so that binding raises Row.Shape_error (gh-ocannl-1655): bind it first to have the rows equal. Binding it to a dimension stands.

Asserts a multiplicative relation between two dimension variable references: the first ref equals factor times the second. Mirrors set_equal in spirit but expresses non-convolutional scaling -- the motivating use case is hourglass-style networks where one dimension is a constant multiple of another (e.g. an MLP hidden width that is twice the input width).

Concretely, if the first ref ends up bound to dimension d_large and the second to d_small, this asserts d_large = factor * d_small.

Restrictions:

• Both refs must be (or resolve to) dimension variables. Row-variable and mixed dim/row variants are not currently supported and raise Row.Shape_error regardless of which side is solved (no silent routing through set_dim's row total-elements path).
• factor must be a positive integer (>= 1). factor = 1 is accepted and behaves exactly like set_equal; callers should prefer set_equal in that case for clarity.
• Convolutional scaling (kernel/dilation/padding) is expressed via einsum convolution notation, not via this function.
• As with set_equal, calling on two fresh refs from get_variable_ref (both still `Not_set_yet) raises Row.Shape_error; use set_dim on one side first, or capture both refs through einsum notation (e.g. +*) so that shape propagation can bind them to dimension variables before this call.

On a both-solved mismatch, on a one-solved-larger case where the solved dimension is not divisible by factor, on insufficient information (both refs are `Not_set_yet), or when either ref resolves to a row variable, this raises Row.Shape_error.

Extracts any available shape information from the initialization or fetch.

Creates a shape. id should be the id the associated tensor (if any). At most one of the pairs batch_dims, batch_axes etc. should be given: if none, the corresponding row will be inferred. batch_axes etc. provide bases (semantic annotations) for the dimensions of the corresponding axes. Note that these are dimension bases and not axis labels: they need not be unique for a row, are inferred when provided, and must match whenever the axis sizes must match.

Bring global state to its initialization values. This invalidates any unfinished inference.

Mark the shape as terminal, so that its rows can be closed to Greatest Lower Bounds (GLBs). This function is only intended for parameters shapes, which would otherwise not be terminal because of the initialization expressions of the parameters.

Mark the shape's rows as resolved by their use sites (gh-ocannl-544): at the final solver stages their row variables may close to the GLB of the use sites' rows, where operation results close to an empty remainder by default. Used by Tensor.param for parameter-initialization cones and by the stretch operation; terminal shapes (leaves, parameters) are marked automatically.

How to propagate shape updates and do the last update of Tensor.t.shape when finalizing the tensor. Axes are broadcast-expanded on a bottom-up update to fit the incoming shape.

Converts a shape logic to its string specification for debugging/display purposes.

Data required for a shape inference update step. Ideally, an update should be performed at least twice, the second time after all the other relevant updates have been performed for the first time. In OCANNL, this is achieved by performing updates both as the tensors are constructed, and via lazy callbacks as the corresponding Ir.Indexing dimensions and projections are first accessed.

Private (gh-ocannl-1573): read and matched freely; built only by make_update_step, and only Shape's projection derivation publishes unsafe_projections, once.

A fresh update step with a new id and no projections yet. It is not registered: see propagate_shapes.

Uses the matrix convention of putting the input axes last.

The product-space proxy shape of an einsum-family operation (gh-512): the result's axes followed by the reduced-over (contracted) axes, in the product-space order of the derived projections. The proxy shares the operation's shape-inference solution (its constraints are emitted on construction), so infer_equal-ing a %cd *_pspace intermediate against it gives the intermediate the dimensions of the operation's full product space (e.g. the (output x kernel) pair space of a windowed reduction). Must be called before the operation's inference completes. Raises Row.Shape_error when the product space is not expressible as a shape: a non-einsum operation, more than one reduced-over row variable, or a reduced-over row variable combined with an open result input row.

Returns the padding of the shape, if any. The inner float option is the padded value: Some v when all operations use the same neutral element, None when different operations require different neutral elements (margin must be reset before each operation). Uses the matrix convention of putting the input axes last.

Runs f with the construction-in-progress depth raised. While it is positive, every entry point that finalizes inference (to_dims, to_padding, get_projections) raises Row.Shape_error at the finalizer, before anything is committed: Tensor.op runs its op_asn under it, since the operation has not registered its constraints yet and a finalization there would close its operands without them (gh-ocannl-830).

Registers the step and solves its constraints eagerly. Refused while the session is tainted; a failure after registering taints it.

Raises Row.Shape_error when the session is tainted.

The first failure's provenance, when the session is tainted.

Records a taint; the first one is kept.

The message of a Row.Shape_error or Utils.User_error, else the printed exception.

The base of the reserved range from which a window that is always rewound -- a shape-scheme trace, a definition check -- allocates its tensors' ids, above every session id (gh-ocannl-1751). Tensor allocates from it; error_to_string_hum recognizes such ids.

error_to_string_hum msg trace renders a Row.Shape_error (msg, trace) for a person (gh-ocannl-1799): the message, then one indented line per finding -- the conflicting sizes or rows (or the one axis or row the message is about), the shapes involved, the failed constraint (with ⊑ for "refines", as in Tensor.Scheme.to_string), and the tensors it relates with their axis kinds and the operation that related them. Variables are named by first appearance ('a, ..'b.., followed by :<name> for a named dimension variable), and a tensor of a trace or a check, whose id comes from the reserved range, by first appearance too (n'1, #'1), so the text depends on neither the session's id counters nor the reserved ones; a session tensor keeps its id. This is what Printexc.to_string -- and so the toplevel -- prints for a Row.Shape_error (prefixed "Shape_error: "), and what its sexp converter (Exn.to_string) carries, one atom per line; the setting shape_error_full_sexp=true restores the full constraint sexp to both.

Runs f; if it raises after the session was mutated, taints the session with describe exn, then re-raises with the original backtrace.

Installs the hook called at the end of every successful closure (Tensor records its checkpoint there).

Raises when called within a construction window; what names the caller in the message.

Whether an explicit checkpoint can be taken now: outside a construction window and untainted.

A shape-scheme trace refused by design (gh-ocannl-1751): the function is valid, but a scheme cannot carry what it does -- a precision-dependent total-elements constraint (random number generation), set_dim, set_sym_dim, set_equal or set_scale on a capture reference made outside the trace (gh-ocannl-1791; binding one is an input), or forcing dims or projections during the trace. Raised by Scheme.trace, prefixed "tracing the shape scheme <label>: "; a contradiction raises Row.Shape_error instead.

Whether a shape-scheme trace is running.

as_tracing ~label f runs f as part of the shape-scheme trace label: its update steps are not counted or timed as inline registrations. Scheme.trace runs its function this way, and Tensor.Scheme makes a trace's tracers this way just before the trace starts.

Registration statistics (gh-ocannl-1752), process-wide since the last reset: checked %op applications, and every registration outside a trace that ran get_inequalities. Projection derivation is not counted.

The constraint counts (gh-ocannl-1752) say what each source adds to the open constraint store: a scheme instantiation's renamed trace constraints (C'), its renamed solution equations (E'), and its boundary and anchor equations -- checked instantiations included, though their window is rewound -- and every registration's own inequalities outside a trace. closure_constraints sums the store's size at the entry of each finish_inference (closures counts them).

rank_search_visits is the rank-cycle search's work counter (gh-ocannl-1802): the states it visited, for every rank fact recorded since the reset, inside traces too.

Where shape-registration wall time goes (gh-ocannl-1752), process-wide since the last reset, in seconds: tracing schemes (the body on tracers, and the rewind), instantiating them (checked applications and explicit instantiations), every step registered outside a trace, and finish_inference. A phase nested in another (an instantiation inside a closure built mid-construction, say) is charged to the inner phase alone. Construction work that is not shape registration (building tensors, their code) is in none of them.

Runs the function, charging its wall time, less that of nested timed phases, to the phase.

An explicit checkpoint; id_mark is the next tensor id. Refuse it while tainted or within a construction window first (Tensor.checkpoint does).

The last closure's checkpoint (empty frontier), which stays valid until the next closure.

The records the session holds of what rewinds discarded, by kind: a diagnostic for tests.

Returns the projections for this update step, computing them if not already done. This triggers finish_inference and then retrieves the projections from unsafe_projections. Use this instead of derive_projections directly.

The per-axis dimension bases (semantic annotations); the basis is total, so unannotated axes read back as "default" and the broadcast top as "bcast_if_1". Uses the matrix convention of putting the input axes last.

Per-kind dimension bases (batch, input, output), so callers can check which row a basis landed on rather than only that it appears somewhere in the flattened to_bases list.

Parse a N5_layout priority string (e.g., "0,1,2") into display indices. Only supports integer labels (Fixed_index).