} else if (auto xla_variadic_reduce_v2_op =
                   dyn_cast<TF::XlaVariadicReduceV2Op>(op)) {
      return propagate_shape_to(xla_variadic_reduce_v2_op.getReducer());
    } else if (auto xla_variadic_sort_op =
                   dyn_cast<TF::XlaVariadicSortOp>(op)) {
      return propagate_shape_to(xla_variadic_sort_op.getComparator());
    }
  }