XlaPlatformInfoFromDevice(ctx->device())));
    }
    TF_ASSIGN_OR_RETURN(auto transfer_manager,
                        xla::TransferManager::GetForPlatform(platform));

    xla::Shape output_device_shape = output.on_device_shape();
    TF_RETURN_IF_ERROR(transfer_manager->ReadDynamicShapes(
        stream, &output, &output_device_shape));