Module Ocannl_tensor.Tensor

The tensor's forward owner cell (gh-ocannl-1796): whether its forward code is owned by an unconsumed forward root, kept apart from the session's weak root map so that consume_forward_code's refusal does not depend on what the GC has reclaimed.

A mutable cell shared by every copy of the record, so the marker follows the tensor; it is what lets the consume functions say "already consumed" rather than "not a root".

Information needed for compositional code generation.

init_params ?skip t collects into a single sequence the forward code of t.params, and transitively the initializations of the parameters of the parameters. If skip is provided, it is used to filter out the parameters belonging to skip (e.g. skip can be a set of parameters that are already initialized). NOTE: it always outputs code with a block comment, even if the params set is empty.

Whether t's forward code is still unconsumed: no consumer embedded it and it was not handed out. The session holds its forward roots WEAKLY, checkpoints included (gh-ocannl-1796), as it does its backprop roots: a graph built and dropped without ever being evaluated is reclaimed by the GC, as is a root only the last closure's checkpoint captured. consume_forward_code's refusal of a tensor whose reads another forward root computes does not read the map, so it does not depend on the GC: a root dropped before its code was consumed keeps refusing its readers.

Whether t's backprop code is still unconsumed. The session holds its backprop roots WEAKLY, checkpoints included (gh-ocannl-1745). The entry lives while any record sharing t's consumption cell is reachable from outside the session's root maps: t itself, a { t with ... } copy that keeps ownership, or a tensor that holds t as a child. A differentiable graph whose forward code was handed out (Train.forward, Train.forward_once), and which the caller then dropped, is reclaimed by the GC. A loss built over a tensor the caller still holds embeds its backprop as before. A dropped root no longer vetoes consume_backprop_code; a dependence on its code is refused by the compilation lineage's requirement check (Context.run).

The unchecked half of consume_forward_code, for a tensor the caller has established is a forward root (is_fwd_root): removes it from the forward roots, marks its forward code as taken, and returns the code. It is an EMBEDDING, not a handout: rewind_to_last_closure treats an embedding made since the last closure as discarded work and makes the tensor a root again; code a caller keeps is handed out with consume_forward_code. This is the handout every path that embeds a tensor's forward code goes through -- consume_forward_code and the %cd operands' embedding alike -- so that a later consume_forward_code on the same tensor returns the same comp.

Runs f and restores the forward and backprop root maps afterwards, including the consumption markers of the roots that were live at entry.

The default precision for the value node of terminal (i.e. non-composite) tensors.

Note: the precision of a node can be set arbitrarily via Ir.Tnode.update_prec. The default precision for value nodes of composite tensors is the maximum of precisions of the value nodes of sub-tensors.

The default precision for the gradient node of terminal (i.e. non-composite) tensors.

Note: the precision of a node can be set arbitrarily via Ir.Tnode.update_prec. The default precision for gradient nodes of composite tensors is the maximum of precisions of the gradient nodes of sub-tensors.

Labels are collected in tensor construction order, with more specific information first.

The defaults are pointwise operations. The grad_asn function receives the non-differentiable variant of the tensor as an argument, which can be used to access the tensor's value in a tensor expression. op_label is prepended to label of the op_fun signature.

The operation's neutral element is read off the assignments op_asn returns, and its shape update step is created afterwards, so op_asn must not force projections.projections, nor finalize inference by any other route (Shape.to_dims/to_padding): the construction is rejected (a Session_error, resp. a Row.Shape_error raised at the finalizer, before anything is committed). The rejection is a bug report, not a recovery API -- shape inference is not transactional (gh-ocannl-903), so the session is reset rather than continued. projections.product_shape and *_pspace intermediates are fine. An operation that needs shape information uses Shape.set_dim or the einsum capture syntax instead; captured references are bound at registration, after op_asn, so relating them (Shape.set_equal/set_scale) belongs at the operation-expression level.

See binop -- same comments apply.

See binop -- same comments apply.

Like binop, but uses block/concatenation shape logic over an array of RHS tensors.

A terminal: a constant, a parameter, an input of the model. The semantics of shape specification is the same as in Shape.make, and by default the shape will be inferred. At most one of init_data or fetch_op should be provided. If init_data is provided, it is used to initialize the tensor's value node. If fetch_op is provided, it is used to generate the tensor's forward code. If init_data is provided, it is also used to verify the shape of the tensor's value node: Broadcast fills the inferred shape with a scalar without constraining its element count, Reshape (the default) if the data is not padded and both the tensor's shape and padding are inferred, Keep_shape_no_padding if the tensor should not be padded and the shape is as given by the ndarray, and Padded if the data is already padded as given, and the shape is as given by the ndarray.

With ~grad_spec:Require_grad the term is a differentiable leaf that is not a parameter: it stays out of params (nothing initializes it; its value is its data or its forward code), and its gradient carries observation intent (Ir.Tnode.set_observable), as a parameter's does -- the host and later routines are its only readers, so it can be read with Context.get_values or stepped through the optimizers' ?params without a materialization hint. Its value carries the intent too, since a recomputed gradient reads it. A leaf whose value comes from fetch_op is re-established by every forward run, so an optimized input should be data-backed (init_data) or a parameter.

A number: a tensor with a single axis of one dimension, initialized to the given value. axis_basis sets the dimension basis (semantic annotation) of that axis. grad_spec is by default Prohibit_grad.

A number with exact bit representation: a tensor with a single axis of one dimension, initialized to the given int64 value. Useful for initializing uint4x32 tensors where exact bit patterns matter. grad_spec is by default Prohibit_grad.

A tensor with an explicit shape, initialized to the given values. Omitted shape rows default to no axes. grad_spec is by default Prohibit_grad. If strict is true (the default), the given values must fill the tensor's value node precisely; otherwise, the values will be looped over to populate the value node.

param is intended for all tensors with initialization; these typically are parameters, but param is also used for non-differentiable tensors that need to be initialized. For proper parameters, pass require_grad:true. t should produce a tensor with no batch axes; input and output axes should by default be inferred. t's label is the passed string, appended by more_label if any, other parameters are forwarded to t. This function returns t's result with the field params replaced by a singleton set containing that result, and it also updates the memory modes. If require_grad is true, any gradient structure inherited from the initialization expression is replaced by a fresh gradient for the final parameter value only. The result has fresh consumption state, independent of any prior handouts of t’s result.

Post-processing hook applied by param to each fully-constructed parameter before it is returned to user code; defaults to the identity. Mixed-precision recipes install a wrapper here that returns a reduced-precision "cast twin" consuming the parameter (gh-ocannl-492 master weights): the graph then reads the twin while the optimizer keeps updating the master parameter, which remains the sole member of the result's params. Reset by unsafe_reinitialize.

A term wrapper that sets up the value node initialization (it generalizes ndarray to tensors with inferred shapes).

A forward root is a tensor that is not (currently) used to compute another tensor. On the first handout, checks that non-embedded reads are not owned by other unconsumed roots -- including a root the caller dropped before consuming it, whose code can then never run, so its readers must be rebuilt (gh-ocannl-1796) -- then removes t from the forward roots and returns its comp. Subsequent calls return the same comp, including after take_forward_code (the %cd embedding path), without checking or modifying roots constructed since the first handout. This is equivalent to retaining the original comp; the caller still owns execution order and the availability of non-embedded inputs. Parameters and subtensors embedded in a tensor consumer remain rejected.

The conflicting-roots check carries correctness, not just an early diagnosis (gh-ocannl-1741): no lineage check backs it, since a computation's requires names backprop fragments only, and the run-time initialization check accepts a read the owner never computed when the node is zero-valid or was written by other means (set_values, an earlier routine).

A backprop root is a differentiable tensor not currently receiving gradients from another tensor. Checks non-embedded reads against other live backprop roots on the first handout and removes t from the backprop roots. Subsequent calls return the same comp without checking or modifying newer roots, as for consume_forward_code; compiling the replayed code a second time in one compilation lineage is refused by Context.compile (gh-ocannl-1724). Non-differentiable tensors and embedded subtensors remain rejected. The returned code does not include zero_grads or seeding the output gradient; callers must still arrange those for each run.

The conflicting-roots check is an early diagnostic (gh-ocannl-1741): a gradient write into another root's propagation is a backprop fragment the returned comp requires, which Context.run refuses unless the routine's own lineage owns it (gh-ocannl-1724). It stays because it fires at the handout, naming the root to consume first, and refuses the order outright where the lineage check would wait for the run.

iter_embedded t iterates over all descendant nodes that are embedded, i.e. are members of t.forward.embedded_nodes or 't.diff.backprop.embedded_nodes' (if any). Note: iter_embedded should only be called after shape inference finishes.

bump_next_id id ensures the next allocated tensor ID will be strictly greater than id. Called after loading tensors from a checkpoint file to prevent ID collisions with subsequently created tensors.

get_next_id () returns the next tensor ID that will be allocated by the session.

Restores a usable session after a failed construction, without discarding closed or compiled work (gh-ocannl-1301). A construction that fails after mutating the inference session taints it: from then on, constructions, Shape.finish_inference (so forcing dims or projections), code handout and every Context compile raise, naming the first failure. Rewinding discards everything opened since the last successful Shape.finish_inference -- the failed construction and every construction since that closure, successful or not -- and clears the taint:

• the solver's open state is emptied (closed shapes carry no variables, and their committed rows, paddings and neutral elements cannot have changed: a write that would change them is refused);
• the forward and backprop root maps are restored to their state at the closure, with the handout markers of the restored roots reset; code handed out to callers since then stays handed out;
• the tensors allocated since the closure are discarded (exactly the nodes the frontend allocated, by uid: a node created in that window outside the frontend is not): using one as an operand, handing out or compiling its code, printing it or transferring its node raises Session_error naming the rewind's provenance ("discarded by the rewind after ..."), as do their update steps' projections;
• variable references the discarded work bound or set are refused by Shape.set_dim and the other reference setters, and by any later einsum they are passed to;
• the cached random seed is restored to the closure's.

Counters are never reset, so later tensors never collide with discarded ones. Contexts, routines and buffers are untouched, and nothing executed is undone. Without a taint it simply discards the open work (reason names it in the stale-handle errors). Raises when called from within an op_asn. The explicit checkpoints taken since the closure can no longer be rewound to.

A snapshot of the open inference session.

Snapshots the open frontier: the solver's and the graph's session state (persistent values, recorded as pointers), and every mutable field of the open work's shapes, update steps and tensor nodes and of every live capture reference (Shape.get_variable_ref). The cost is linear in that frontier, so take one per phrase or trace, not per operation. label names it in messages (default "checkpoint #n"). Raises when the session is tainted (rewind the failure first) and when called from within an op_asn.

Puts the session back as it was at the checkpoint, and clears the taint:

• every object open at the checkpoint gets its fields back -- the shapes' rows, paddings and neutral elements (rewritten in place by later operations), the capture references' bindings and solutions (a reference made before the checkpoint and solved after it is unsolved again), the open nodes' precision links, memory-mode intents, ownership and observability (a Tensor.param over an open initializer included);
• the solver's environment, constraints, active steps, inference sets and rank graph, the graph-root maps (with the handout markers of the restored roots reset) and the random seed are the checkpoint's; code handed out to callers since stays handed out;
• what the window created is discarded, as by rewind_to_last_closure: its tensors (exactly the nodes the frontend allocated), update steps and the capture references it made and bound raise Session_error / Shape_error naming the rewind's provenance when used.

Counters are never reset. Closed and compiled work is untouched, and nothing executed is undone. The checkpoint stays valid, so it can be rewound to again; the checkpoints taken after it can not. Raises Session_error, before anything is written, when the checkpoint cannot be rewound to: inference closed since it was taken (a closure commits); a rewind past it, rewind_to_last_closure or unsafe_reinitialize ran since; or the window memoized state of the frontier that no rewind can reset -- a dims, padding or projections read of an open shape that raised (its lazy keeps the exception), or an open node's precision that began settling. rewind_to_last_closure is always available. A Tensor.take_forward_code called directly on a tensor that was not a root at the checkpoint leaves its marker taken. Raises when called from within an op_asn.

with_checkpoint f runs f; when f raises, rewinds to a checkpoint taken at entry and re-raises f's exception with its backtrace. When that checkpoint can no longer be rewound to (say f closed inference), f's exception still wins, and the session keeps any taint f left. A normal return keeps f's work.

with_rewind f runs f and always rewinds to a checkpoint taken at entry -- a trace that leaves nothing in the session -- returning f's result. Read what you need of the open state inside f; tensors f creates are discarded handles afterwards, so create the ones meant to outlive the trace (shared parameters) before calling it. When f returns but its checkpoint can no longer be rewound to, the rewind's Session_error is raised; when f raises, f's exception wins, as in with_checkpoint.

Diagnostics and negative controls for the checkpoint machinery's tests.

Sets an operand's precision ahead of an operation that requires it; refused (raises Session_error) when it would change a closed operand's precision.

update_operand_prec over several operands, all validated before any is written.

Raises Session_error when the tensor's value or gradient node was discarded by a rewind.

Bring global state to its initialization values. This invalidates any previously defined tensors and tensor nodes. Also reinitializes the modules: Shape, Ir.Tnode.

namespace sets the namespace stamped on subsequently created tensor nodes (gh-ocannl-372); this is the only way to change it. Defaults to Ir.Tnode.default_namespace (ocannl) — also when omitted, i.e. reinitializing never preserves a previously set namespace. Namespaces must match A-Za-z_A-Za-z0-9_* so they can qualify identifiers in generated code and debug names (rendered as ns__n42; the default namespace is elided).

While this function is intended for testing, using it can prevent unintentional session state pollution errors.

Creates the random seed tensor. If seed is provided, it is used to set the random seed. Otherwise, the seed is taken from the settings.

Returns a tensor with the current random seed. Lazily initialized using set_random_seed and reset when unsafe_reinitialize is called. IMPORTANT: all sites using the same global random seed, e.g. using get_random_seed () not separated by a call to unsafe_reinitialize, must descend from the first caller's optimization context.

The seeds of the random-seed singletons this session created (set_random_seed, directly or lazily through get_random_seed), oldest first, without one a rewind discarded: every seed the session's graphs can draw from, including those built under with_saved_random_seed.

Runs f, restoring the global random-seed singleton (the one get_random_seed returns) to its prior value afterwards, even on exception. Use this to scope set_random_seed mutations — e.g. building several graphs with distinct per-graph seeds — without perturbing the caller's seed. Tensors built inside f keep referencing the seed tensors that were current when they were constructed; only the global pointer is restored.

Converts ID, label and the dimensions of a node to a string.

Logs debug information about the tensor on the default ppx_minidebug runtime.

We print out up to 5 axes when printing a tensor, as a grid (outer rectangle) of (inner) rectangles, possibly repeated (screens).

to_printbox is the box print_tree prints, laid out by depth the same way, with with_id and with_value defaulting to false and true.

print renders a tensor's metadata and, when an explicit ?ctx is supplied, its values via an on-demand device-to-host transfer (gh-ocannl-333). Without ?ctx (the default) values are shown as a placeholder.

print_tree renders the tensor's graph: its outermost depth levels side by side in boxes, the levels below them as an indented tree. depth shapes the layout only, so every node of the graph is shown. With ?ctx, the values of all the nodes it shows are read together through Context.to_host_all, so the virtual ones are recomputed by one throwaway routine (gh-ocannl-1737); an unreadable node is shown as a placeholder naming why (Context.observability_tag). print reads the value and gradient it shows the same way.