&else_branch.front(), module, buffer_to_size,
          decomposed_partitioned_call_callees)))
    return failure();

  auto output_buffer_to_size = AddTensorListSizesToTerminator<TF::YieldOp>(
      then_branch.front(), *buffer_to_size);
  AddTensorListSizesToTerminator<TF::YieldOp>(else_branch.front(),
                                              *buffer_to_size);

  if (output_buffer_to_size.empty()) return success();