Module Ir.Assignments

Assignment computations over tensor nodes.

type init_data =
  1. | Broadcast of Base.float
    (*

    Host-backed scalar fill of the inferred shape, including its committed padding. Unlike Reshape, this does not constrain the number of elements.

    *)
  2. | Reshape of Ndarray.t
  3. | Keep_shape_no_padding of Ndarray.t
  4. | Padded of {
    1. data : Ndarray.t;
    2. padding : Ops.axis_padding Base.array;
    3. padded_value : Base.float;
    }
val sexp_of_init_data : init_data -> Sexplib0.Sexp.t
val equal_init_data : init_data -> init_data -> Base.bool
type buffer =
  1. | Node of Tnode.t
  2. | Merge_buffer of Tnode.t
val sexp_of_buffer : buffer -> Sexplib0.Sexp.t
val equal_buffer : buffer -> buffer -> Base.bool
type fetch_op =
  1. | Constant of Base.float
  2. | Constant_bits of Base.int64
  3. | Constant_fill of Base.float Base.array
  4. | Range_over_offsets
  5. | Slice of {
    1. batch_idx : Indexing.static_symbol;
    2. sliced : Tnode.t;
    }
  6. | Embed_symbol of Indexing.static_symbol
  7. | Embed_self_id
  8. | Embed_dim of Indexing.variable_ref
val sexp_of_fetch_op : fetch_op -> Sexplib0.Sexp.t
val equal_fetch_op : fetch_op -> fetch_op -> Base.bool
type accum_rhs =
  1. | Ternop of {
    1. op : Ops.ternop;
    2. rhs1 : buffer;
    3. rhs2 : buffer;
    4. rhs3 : buffer;
    }
  2. | Binop of {
    1. op : Ops.binop;
    2. rhs1 : buffer;
    3. rhs2 : buffer;
    }
  3. | Unop of {
    1. op : Ops.unop;
    2. rhs : buffer;
    }
  4. | Block of {
    1. op : Ops.unop;
    2. rhses : buffer Base.array;
    }
  5. | Rev_sides of {
    1. op : Ops.unop;
    2. lhses : buffer Base.array;
    }
val sexp_of_accum_rhs : accum_rhs -> Sexplib0.Sexp.t
val equal_accum_rhs : accum_rhs -> accum_rhs -> Base.bool
type t =
  1. | Noop
  2. | Seq of t * t
  3. | Block_comment of Base.string * t
  4. | Accum_op of {
    1. initialize_neutral : Base.bool;
    2. accum : Ops.binop;
    3. lhs : Tnode.t;
    4. rhs : accum_rhs;
    5. projections : Indexing.projections Utils.Lazy.t;
    6. projections_debug : Base.string;
    }
  5. | Set_vec_unop of {
    1. op : Ops.vec_unop;
    2. lhs : Tnode.t;
    3. rhs : buffer;
    4. projections : Indexing.projections Utils.Lazy.t;
    5. projections_debug : Base.string;
    }
  6. | Fetch of {
    1. array : Tnode.t;
    2. fetch_op : fetch_op;
    3. dims : Base.int Base.array Utils.Lazy.t;
    }
val sexp_of_t : t -> Sexplib0.Sexp.t
val fold_leaves : t -> init:'a -> f:('a -> t -> 'a) -> 'a

Folds f over the LEAF statements (Accum_op, Set_vec_unop, Fetch) in execution order, skipping the Noop / Seq / Block_comment scaffolding. Descending is what this module's queries and the tests' finders used to each write for themselves, and forgetting to descend through a Block_comment loses statements silently -- so the recursion is written once. f still matches on t (the leaf kinds mean different things to different queries), and its scaffolding arm is unreachable (gh-ocannl-818 would give the leaves a type of their own).

val iter_leaves : t -> f:(t -> Base.unit) -> Base.unit

fold_leaves for a f that only has effects.

type side =
  1. | Fwd
  2. | Bprop
val sexp_of_side : side -> Sexplib0.Sexp.t
val side_of_sexp : Sexplib0.Sexp.t -> side
val compare_side : side -> side -> Base.int
val equal_side : side -> side -> Base.bool
module Fragment : sig ... end

A code fragment generated from a tensor, identified by its owner node and side (gh-ocannl-1724): (t.value, Fwd) is t's forward code; (t.grad, Bprop) is t's backprop code, the code propagating from t.grad into its children's gradients.

type comp = {
  1. asgns : t;
  2. embedded_nodes : Base.Set.M(Ir.Tnode).t;
    (*

    The nodes in asgns that are not in embedded_nodes need to already be in contexts linked with the comp -- an initialization claim, not ownership.

    *)
  3. owned : Base.Set.M(Fragment).t;
    (*

    The tensor-generated code fragments this computation carries. Context.compile refuses a computation owning a fragment some routine of its compilation lineage already owns.

    *)
  4. requires : Base.Set.M(Fragment).t;
    (*

    The backprop fragments its gradient writes flow into that it does not carry itself. Context.run refuses the routine unless a routine of its own lineage (never a sibling) owns each one.

    *)
  5. zeroes : Base.Set.M(Ir.Tnode).t;
    (*

    The gradients whose tensor-generated zeroing this computation carries (gh-ocannl-1744): a tensor's zero_grads and whatever is sequenced from it.

    *)
  6. accumulates : Base.Set.M(Ir.Tnode).t;
    (*

    The gradients its tensor-generated backprop accumulates into (gh-ocannl-1744): the operands' gradients a tensor's own gradient code writes, and its embedded children's. Context.compile refuses a computation accumulating into a gradient that another routine of its lineage accumulates into, unless a zeroing compiled before both, and carried by neither, resets it; and one accumulating into a gradient whose propagation (its Bprop fragment) an earlier routine of the lineage already carries.

    *)
  7. accumulation_runs : Base.int Base.Map.M(Ir.Tnode).t;
    (*

    For a gradient in accumulates and not in zeroes, the runs of this computation one reset of the gradient deliberately spans (gh-ocannl-1781); absent means one. Context.run refuses a routine that has already accumulated into such a gradient's buffer that many times since the buffer was last reset -- by a run of a routine zeroing it, or by a write from outside a run. Train.grad_update ~accum_steps:k states k: the explicit opt-in for gradient accumulation over micro-steps.

    *)
}

Computations based on assignments. Tensor construction fills in all five sets; zeroing owns and requires nothing (it fills zeroes only), and initialization and hand-written %cd assignments fill none: a hand-written write to a gradient neither contributes nor resets.

val sexp_of_comp : comp -> Sexplib0.Sexp.t
val comp : ?owned:Base.Set.M(Fragment).t -> ?requires:Base.Set.M(Fragment).t -> ?zeroes:Base.Set.M(Ir.Tnode).t -> ?accumulates:Base.Set.M(Ir.Tnode).t -> ?accumulation_runs:Base.int Base.Map.M(Ir.Tnode).t -> embedded:Base.Set.M(Ir.Tnode).t -> t -> comp

The computation constructor; owned, requires, zeroes, accumulates and accumulation_runs default to empty.

val no_fragments : Base.Set.M(Fragment).t
val to_comp : t -> comp
val empty_comp : comp
val context_nodes : plc:Tnode.Placements.t -> t -> Tnode.t_set
val collect_nodes_guess_output : t -> Tnode.t_set * Tnode.t_set
val collect_written : t -> Tnode.t_set
val sequence : comp Base.list -> comp

Runs the computations in order, uniting their nodes, fragments and gradient sets; a gradient two members state different accumulation_runs for keeps the smaller, a member accumulating into it without zeroing it and stating nothing counting as one. Raises Utils.User_error, naming the fragment, when two members own the same fragment (gh-ocannl-1746): owned is a set, so the copy would otherwise run twice unseen -- for a backprop fragment, a double-counted gradient.

val collect_neutral_elem : t -> Base.float Base.option
val to_low_level : ?static_indices:Indexing.static_symbol Base.list -> t -> Low_level.t
val to_doc : ?name:Base.string -> ?static_indices:Indexing.static_symbol Base.list -> Base.unit -> t -> PPrint.document
val get_name_exn : t -> Base.string
val lower : ?rewrite_target:Rewrites.target -> ?blends_value_selects:Base.bool -> Low_level.optimize_ctx -> unoptim_ll_source:(PPrint.document -> Base.unit) Base.option -> ll_source:(PPrint.document -> Base.unit) Base.option -> cd_source:(PPrint.document -> Base.unit) Base.option -> name:Base.string -> Indexing.static_symbol Base.list -> t -> Low_level.optimized

Lowers and optimizes t. blends_value_selects is the rendering backend's fact, passed on to Low_level.optimize (gh-ocannl-1775); Backends.lower_assignments always states it.