venv/lib/python3.11/site-packages/deepspeed/ops/csrc/compile/z1.h

// Copyright (c) Microsoft Corporation.
// SPDX-License-Identifier: Apache-2.0

// DeepSpeed Team

#include "deepcompile.h"

#pragma once

namespace dc {

void register_graph_z1(long graph_id, const std::vector<long>& ds_ids);
void register_z1_param(long ds_id,
                       const std::vector<int64_t>& ds_shape,
                       at::Tensor ds_tensor,
                       at::Tensor grad_buffer,
                       int64_t offset);
}  // namespace dc