Module Uop.Arg

type t =
  1. | Empty
  2. | Int of int
  3. | Ints of int list
  4. | Bools of bool list
  5. | String of string
  6. | Value of Const.t
  7. | Op of Ops.t
  8. | Range_info of {
    1. axis : int;
    2. sub : int list;
    3. kind : Axis_type.t;
    }
  9. | Param_arg of param_arg
  10. | Reduce_arg of reduce_arg
    (*

    For Reduce: reduction op plus tensor axes.

    *)
  11. | Device of device
    (*

    For Copy: device placement.

    *)
  12. | Op_device of Ops.t * device
    (*

    For Allreduce: reduction op and device group.

    *)
  13. | Stage_info of stage_opts
  14. | Opts of Opt.t list
    (*

    For Contiguous: schedule options attached to the boundary.

    *)
  15. | Kernel_info of kernel_info
  16. | Call_info of call_info
  17. | Program_info of program_info
  18. | Wmma_info of wmma_info
  19. | Shaped_wmma_info of shaped_wmma_info
val equal : t -> t -> bool

equal a b is structural equality on payloads. Call_info values are compared with physical equality on their grad_fxn closure, and structural equality on the rest.

val compare : t -> t -> int

compare is Stdlib.compare on payloads. Used only for total ordering; for semantic comparison use equal.

Payload accessors

One accessor per generic payload variant, returning None when the variant does not match. Use Option.is_some for a boolean check. The richer named-record payloads are destructured through the view accessors on nodes.

val as_int : t -> int option
val as_ints : t -> int list option
val as_bools : t -> bool list option
val as_string : t -> string option
val as_value : t -> Const.t option
val as_op : t -> Ops.t option
val as_param_arg : t -> param_arg option
val as_reduce_arg : t -> reduce_arg option
val as_device : t -> device option
val as_opts : t -> Opt.t list option
val as_stage_info : t -> stage_opts option
val as_program_info : t -> program_info option