Module Ir.Assignments

Assignment computations over tensor nodes.

type init_data =
  1. | Reshape of Ndarray.t
  2. | Keep_shape_no_padding of Ndarray.t
  3. | 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
type comp = {
  1. asgns : t;
  2. embedded_nodes : Base.Set.M(Ir.Tnode).t;
}
val sexp_of_comp : comp -> Sexplib0.Sexp.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
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 : 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