def _make_first_chunk_prog(
    prog: _mil.Program,
    op_idx: int,
) -> _mil.Program:
    """Build first chunk by declaring early outputs and removing unused subgraph"""
    block = prog.functions["main"]
    boundary_vars = _get_first_chunk_outputs(block, op_idx)

    # Due to possible numerical issues, cast any fp16 var to fp32
    new_boundary_vars = _add_fp32_casts(block, boundary_vars)

    block.outputs.clear()
    block.set_outputs(new_boundary_vars)
    _PASS_REGISTRY["common::dead_code_elimination"](prog)
    return prog
Build first chunk by declaring early outputs and removing unused subgraph
coremltools/models/utils.py