Skip to main content

Module shape_infer

Module shape_infer 

Source
Expand description

Output-shape and output-dtype inference for StdTensorOp.

Called during StdTensorOp -> ExecProgram lowering to populate ExecInstruction::output_shapes and ExecInstruction::dtype.

Functionsยง

infer_output_dtype
Infer output dtype for a single instruction given its op and input dtypes.
infer_output_shapes
Infer output shapes for a single instruction.