docs menu

Module Rune.Probe

Values observed inside functions.

A probe receives values that code anywhere inside a function emits, and the caller reads them when the function returns or raises: Probe.collect p f is f ()'s outcome beside the offers that reached it. Nothing reads an offer before collect returns, so an emission never changes a value the function computes.

Sites and offers. An emission's site is an emit call. Outside every traced function each call is a site run once, its offer the value with each array given a leading axis of one row. Inside a traced function each emit made as the function is traced is a site, and each call of the outermost construct or map inside the collect offers once per site that ran, one row per execution along a new axis 0. Enclosing rows, steps and lanes merge into that axis, the outer one major: a scan of n rows whose step emits a [3] value offers [n; 3]; a vmap of 4 lanes around it offers [4n; 3], lane 0's rows first. A step or branch not taken and a held lane (held lanes) offer no row, and a site with no row no offer.

Differentiation. Emissions offer primals, once per execution: under reverse mode those of the forward pass. A backward pass, a transpose, a remat's computing again, a root's derivative, and a custom rule's function derived again or applied to tangents offer nothing. A custom_jvp tangent map's emissions are offered under forward mode.

#type 'a t

The type for probes of 'a values.

#val make : ('a, 'd) Nx.Ptree.t -> 'a t

make s is a fresh probe of values of s's structure, distinct from every other.

#val emit : 'a t -> 'a -> unit

emit p v offers v to the innermost open collect of p. With none open it does nothing, at the cost of one effect, walking nothing; one that reaches a collect reshapes v's arrays.

Raises Invalid_argument while a collect of p is open: naming Rune.Probe.emit for a tangent under reverse mode, and naming the entry point in an iterate's until, whose tests transformations make apart.

val collect : 
  'a t ->
  (unit -> 'r) ->
  ('r, exn * Stdlib.Printexc.raw_backtrace) Stdlib.result * 'a list

collect p f is f ()'s result, or the exception it raised with its backtrace, and the offers to p that reached it. Outside every transformation they reach it in the order their emits and constructs return, a construct's in the order its sites were made; a transformation around a map may deliver a site's offer before another's. When f raises, the offers are those of the emits that ran before the raise: an Nx.check failing in row k of a scan's step offers the rows up to k, row k's emits before the check included.

Offers are arrays left unread on the sets their values lie on, except where lanes were held: collect then reads on the host which rows ran, once per offer, and gathers them.

Out_of_memory, Stack_overflow and Sys.Break from f are raised again, with their backtrace, and the offers dropped.

Raises Invalid_argument inside a traced function, a map or a differentiation, where a construct could be answered outside it.