def _infer_elem_type(list_var):
    """
    Returns types.tensor. None if failed to infer element type.
    Example:

    Given:

    main(%update: (2,fp32)) {
      block0() {
        %list: List[unknown] = tf_make_list(...) # unknown elem type
        %while_loop_0:0: (i32), %while_loop_0:1: List[(2,fp32)] = while_loop(loop_vars=(...))
          while_loop_0_body(...) {
            %list_write_0: List[(2,fp32)] = list_write(index=..., ls=%list, value=%update)
          } -> (%add_0, %list_write_0)

        Result:

        main(%update: (2,fp32)) {
          block0() {
        %list: List[(2,fp32)] = tf_make_list(...) # Get the elem type from list_write
        %while_loop_0:0: (i32), %while_loop_0:1: List[(2,fp32)] = while_loop(loop_vars=(...))
          while_loop_0_body(...) {
            %list_write_0: List[(2,fp32)] = list_write(index=..., ls=%list, value=%update)
          } -> (%add_0, %list_write_0)
    """
    # Search for child op that have informative element types
    for o in list_var.child_ops:
        if o.op_type in ["list_write", "list_scatter"]:
            return o.outputs[0].elem_type
        if o.op_type == "while_loop":
            idx = list(o.loop_vars).index(list_var)
            block = o.blocks[0]
            # the corresponding Var in body block
            block_var = block.inputs[idx]
            elem_type = _infer_elem_type(block_var)
            if elem_type is not None:

                def _set_types_for_block_inputs(block):
                    block_var = block.inputs[idx]
                    new_block_var = ListVar(name=block_var.name, elem_type=elem_type,
                                            init_length=block_var.sym_type.T[1],
                                            dynamic_length=block_var.sym_type.T[2])
                    block._replace_var(block_var, new_block_var)

                _set_types_for_block_inputs(o.blocks[0])  # condition block
                _set_types_for_block_inputs(o.blocks[1])  # body block

                return elem_type
            # otherwise continue to other block_var (a list_var can be
            # passed into while_loop twice).
    return None

    Returns types.tensor. None if failed to infer element type.
    Example:

    Given:

    main(%update: (2,fp32)) {
      block0() {
        %list: List[unknown] = tf_make_list(...) # unknown elem type
        %while_loop_0:0: (i32), %while_loop_0:1: List[(2,fp32)] = while_loop(loop_vars=(...))
          while_loop_0_body(...) {
            %list_write_0: List[(2,fp32)] = list_write(index=..., ls=%list, value=%update)
          } -> (%add_0, %list_write_0)

        Result:

        main(%update: (2,fp32)) {
          block0() {
        %list: List[(2,fp32)] = tf_make_list(...) # Get the elem type from list_write
        %while_loop_0:0: (i32), %while_loop_0:1: List[(2,fp32)] = while_loop(loop_vars=(...))
          while_loop_0_body(...) {
            %list_write_0: List[(2,fp32)] = list_write(index=..., ls=%list, value=%update)
          } -> (%add_0, %list_write_0)
    
coremltools/converters/mil/frontend/tensorflow/ssa_passes/backfill_make_list_elem_type.py