if (isa<tensor::CastOp>(op)) return InferShapeForCast(op);
  return false;
}

// Finds element type to be used for result from operand, with special handling
// for handle types.
Type GetElementTypeFromOperand(TensorType operand_type,
                               TensorType result_type) {
  auto operand_handle_type =
      mlir::dyn_cast<TensorFlowTypeWithSubtype>(operand_type.getElementType());