#include <optional>
#include <utility>
#include <unordered_map>
#include <vector>

#include "dlpack/dlpack.h"

// Part of the Array API specification.
#define CUDNN_FRONTEND_DLPACK_CAPSULE_NAME "dltensor"
#define CUDNN_FRONTEND_DLPACK_USED_CAPSULE_NAME "used_dltensor"

#include "pybind11/pybind11.h"
#include "pybind11/cast.h"
#include "pybind11/stl.h"

#include "cudnn_frontend.h"
#include "pygraph.h"

namespace py = pybind11;
using namespace pybind11::literals;

namespace cudnn_frontend::python_bindings {

void
throw_if(bool const cond, cudnn_frontend::error_code_t const error_code, std::string const& error_msg);

void
init_pygraph_norm_submodule(py::class_<PyGraph>&);

void
init_pygraph_sdpa_submodule(py::class_<PyGraph>&);

void
init_pygraph_pointwise_submodule(py::class_<PyGraph>&);

cudnn_frontend::DataType_t
convert_to_cudnn_data_type(const DLDataType& dtype) {
    switch (dtype.code) {
        case DLDataTypeCode::kDLUInt:
            switch (dtype.bits) {
                case 8:
                    return cudnn_frontend::DataType_t::UINT8;
            }
            break;
        case DLDataTypeCode::kDLInt:
            switch (dtype.bits) {
                case 8:
                    return cudnn_frontend::DataType_t::INT8;
                case 32:
                    return cudnn_frontend::DataType_t::INT32;
                case 64:
                    return cudnn_frontend::DataType_t::INT64;
            }
            break;
        case DLDataTypeCode::kDLFloat:
            switch (dtype.bits) {
                case 16:
                    return cudnn_frontend::DataType_t::HALF;
                case 32:
                    return cudnn_frontend::DataType_t::FLOAT;
                case 64:
                    return cudnn_frontend::DataType_t::DOUBLE;
            }
            break;
        case DLDataTypeCode::kDLBfloat:
            switch (dtype.bits) {
                case 16:
                    return cudnn_frontend::DataType_t::BFLOAT16;
            }
            break;
        case DLDataTypeCode::kDLBool:
            switch (dtype.bits) {
                case 8:
                    return cudnn_frontend::DataType_t::BOOLEAN;
            }
            break;
    }
    return cudnn_frontend::DataType_t::NOT_SET;
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::tensor(std::vector<int64_t> const& dim,
                std::vector<int64_t> const& stride,
                cudnn_frontend::DataType_t const& data_type,
                bool const& is_virtual,
                bool const& is_pass_by_value,
                std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> const& ragged_offset,
                cudnn_frontend::TensorReordering_t const reordering_type,
                std::string const& name,
                int64_t const& uid,
                int64_t const& ragged_offset_multiplier) {
    auto props = cudnn_frontend::graph::Tensor_attributes()
                     .set_data_type(data_type)
                     .set_is_virtual(is_virtual)
                     .set_is_pass_by_value(is_pass_by_value)
                     .set_dim(dim)
                     .set_stride(stride)
                     .set_ragged_offset(ragged_offset)
                     .set_reordering_type(reordering_type)
                     .set_name(name)
                     .set_ragged_offset_multiplier(ragged_offset_multiplier);

    if (uid != -1) {
        props.set_uid(uid);
    }

    return graph->tensor(props);
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::tensor_like(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> const& tensor, std::string const& name) {
    return graph->tensor_like(tensor, name);
}

static std::intptr_t
extract_data_pointer(py::object const& obj) {
    throw_if(!py::hasattr(obj, "__dlpack__"),
             cudnn_frontend::error_code_t::INVALID_VARIANT_PACK,
             "Object does not have the __dlpack__() method");

    py::capsule capsule = obj.attr("__dlpack__")();
    throw_if(capsule.is_none(),
             cudnn_frontend::error_code_t::INVALID_VARIANT_PACK,
             "Failed to retrieve the DLPack capsule.");

    DLManagedTensor* managed =
        static_cast<DLManagedTensor*>(PyCapsule_GetPointer(capsule.ptr(), CUDNN_FRONTEND_DLPACK_CAPSULE_NAME));
    throw_if(managed == nullptr, cudnn_frontend::error_code_t::INVALID_VARIANT_PACK, "Invalid DLPack capsule.");

    DLDeviceType device_type = managed->dl_tensor.device.device_type;
    throw_if(
        device_type != kDLCPU && device_type != kDLCUDAHost && device_type != kDLCUDA && device_type != kDLCUDAManaged,
        cudnn_frontend::error_code_t::INVALID_VARIANT_PACK,
        "Invalid device type.");

    void* p     = (char*)managed->dl_tensor.data + managed->dl_tensor.byte_offset;
    auto result = reinterpret_cast<std::intptr_t>(p);
    return result;
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::tensor_like(py::object const& pyobj) {
    throw_if(!py::hasattr(pyobj, "__dlpack__"),
             cudnn_frontend::error_code_t::INVALID_VARIANT_PACK,
             "Object does not have the __dlpack__() method");

    py::capsule capsule = pyobj.attr("__dlpack__")();
    throw_if(capsule.is_none(),
             cudnn_frontend::error_code_t::INVALID_VARIANT_PACK,
             "Failed to retrieve the DLPack capsule.");

    DLManagedTensor* managed =
        static_cast<DLManagedTensor*>(PyCapsule_GetPointer(capsule.ptr(), CUDNN_FRONTEND_DLPACK_CAPSULE_NAME));
    throw_if(managed == nullptr, cudnn_frontend::error_code_t::INVALID_VARIANT_PACK, "Invalid DLPack capsule.");

    DLDeviceType device_type = managed->dl_tensor.device.device_type;
    throw_if(
        device_type != kDLCPU && device_type != kDLCUDAHost && device_type != kDLCUDA && device_type != kDLCUDAManaged,
        cudnn_frontend::error_code_t::INVALID_VARIANT_PACK,
        "Invalid device type.");

    auto ndim = managed->dl_tensor.ndim;
    std::vector<int64_t> dim(managed->dl_tensor.shape, managed->dl_tensor.shape + ndim);

    auto props = cudnn_frontend::graph::Tensor_attributes()
                     .set_data_type(convert_to_cudnn_data_type(managed->dl_tensor.dtype))
                     .set_is_virtual(false)
                     .set_is_pass_by_value(managed->dl_tensor.device.device_type == kDLCPU)
                     .set_dim(dim);

    if (managed->dl_tensor.strides == nullptr) {
        // dlpack says "can be NULL, indicating tensor is compact and row-majored"
        auto stride_order = detail::generate_row_major_stride_order(ndim);
        props.set_stride(detail::generate_stride(dim, stride_order));
    } else {
        std::vector<int64_t> stride(managed->dl_tensor.strides, managed->dl_tensor.strides + ndim);
        props.set_stride(stride);
    }

    return graph->tensor(props);
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::slice(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& input,
               std::vector<py::slice> const& slices,
               cudnn_frontend::DataType_t const& compute_data_type,
               std::string const& name) {
    auto input_dim = input->get_dim();

    std::vector<std::pair<int64_t, int64_t>> start_end_indices;
    std::vector<int64_t> steps;
    steps.reserve(slices.size());
    for (size_t i = 0; i < slices.size(); ++i) {
        int64_t start, stop, step, length;
        if (!slices[i].compute(input_dim[i], &start, &stop, &step, &length)) {
            throw std::runtime_error("Invalid slice");
        }
        CUDNN_FRONTEND_UNUSED(length);
        start_end_indices.push_back({start, stop});
        steps.push_back(step);
    }

    auto attributes = cudnn_frontend::graph::Slice_attributes()
                          .set_slices(start_end_indices)
                          .set_strides(steps)
                          .set_compute_data_type(compute_data_type)
                          .set_name(name);

    auto output = graph->slice(input, attributes);
    return output;
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::conv_fprop(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& image,
                    std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& weight,
                    std::vector<int64_t> const& pre_padding,
                    std::vector<int64_t> const& post_padding,
                    std::vector<int64_t> const& stride,
                    std::vector<int64_t> const& dilation,
                    cudnn_frontend::ConvolutionMode_t const& conv_mode,
                    cudnn_frontend::DataType_t const& compute_data_type,
                    std::string const& name) {
    auto attributes = cudnn_frontend::graph::Conv_fprop_attributes()
                          .set_pre_padding(pre_padding)
                          .set_post_padding(post_padding)
                          .set_stride(stride)
                          .set_dilation(dilation)
                          .set_convolution_mode(conv_mode)
                          .set_compute_data_type(compute_data_type)
                          .set_name(name);

    auto Y = graph->conv_fprop(image, weight, attributes);
    return Y;
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::conv_dgrad(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& loss,
                    std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& filter,
                    std::vector<int64_t> const& pre_padding,
                    std::vector<int64_t> const& post_padding,
                    std::vector<int64_t> const& stride,
                    std::vector<int64_t> const& dilation,
                    cudnn_frontend::ConvolutionMode_t const& conv_mode,
                    cudnn_frontend::DataType_t const& compute_data_type,
                    std::string const& name) {
    auto attributes = cudnn_frontend::graph::Conv_dgrad_attributes()
                          .set_pre_padding(pre_padding)
                          .set_post_padding(post_padding)
                          .set_stride(stride)
                          .set_dilation(dilation)
                          .set_convolution_mode(conv_mode)
                          .set_compute_data_type(compute_data_type)
                          .set_name(name);
    auto DX = graph->conv_dgrad(loss, filter, attributes);
    return DX;
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::conv_wgrad(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& image,
                    std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& loss,
                    std::vector<int64_t> const& pre_padding,
                    std::vector<int64_t> const& post_padding,
                    std::vector<int64_t> const& stride,
                    std::vector<int64_t> const& dilation,
                    cudnn_frontend::ConvolutionMode_t const& conv_mode,
                    cudnn_frontend::DataType_t const& compute_data_type,
                    std::string const& name) {
    auto attributes = cudnn_frontend::graph::Conv_wgrad_attributes()
                          .set_pre_padding(pre_padding)
                          .set_post_padding(post_padding)
                          .set_stride(stride)
                          .set_dilation(dilation)
                          .set_convolution_mode(conv_mode)
                          .set_compute_data_type(compute_data_type)
                          .set_name(name);
    auto DW = graph->conv_wgrad(loss, image, attributes);
    return DW;
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::block_scale_dequantize(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& input,
                                std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& descale,
                                std::vector<int32_t> const& block_size,
                                bool const is_negative_scale,
                                cudnn_frontend::DataType_t const& compute_data_type,
                                std::string const& name) {
    auto attributes = cudnn_frontend::graph::Block_scale_dequantize_attributes()
                          .set_block_size(block_size)
                          .set_is_negative_scale(is_negative_scale)
                          .set_compute_data_type(compute_data_type)
                          .set_name(name);
    if (compute_data_type != cudnn_frontend::DataType_t::NOT_SET) {
        attributes.set_compute_data_type(compute_data_type);
    }
    auto output = graph->block_scale_dequantize(input, descale, attributes);
    return output;
}

std::array<std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>, 2>
PyGraph::block_scale_quantize(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& input,
                              int32_t block_size,
                              std::optional<int64_t> axis,
                              bool transpose,
                              cudnn_frontend::DataType_t const& compute_data_type,
                              std::string const& name) {
    auto attributes = cudnn_frontend::graph::Block_scale_quantize_attributes()
                          .set_block_size(block_size)
                          .set_transpose(transpose)
                          .set_name(name);
    if (axis.has_value()) {
        attributes.set_axis(axis.value());
    }
    if (compute_data_type != cudnn_frontend::DataType_t::NOT_SET) {
        attributes.set_compute_data_type(compute_data_type);
    }
    auto outputs = graph->block_scale_quantize(input, attributes);
    return outputs;
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::matmul(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& A,
                std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& B,
                cudnn_frontend::DataType_t const& compute_data_type,
                double const padding,
                std::string const& name) {
    auto attributes = cudnn_frontend::graph::Matmul_attributes()
                          .set_compute_data_type(compute_data_type)
                          .set_name(name)
                          .set_padding(padding);

    auto C = graph->matmul(A, B, attributes);
    return C;
}

std::array<std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>, 2UL>
PyGraph::genstats(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& input,
                  cudnn_frontend::DataType_t const& compute_data_type,
                  std::string const& name) {
    auto attributes =
        cudnn_frontend::graph::Genstats_attributes().set_compute_data_type(compute_data_type).set_name(name);

    auto [SUM, SQ_SUM] = graph->genstats(input, attributes);
    return {SUM, SQ_SUM};
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::reduction(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& input,
                   cudnn_frontend::ReductionMode_t const mode,
                   cudnn_frontend::DataType_t const& compute_data_type,
                   std::string const& name,
                   std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> group_offset) {
    auto attributes = cudnn_frontend::graph::Reduction_attributes()
                          .set_mode(mode)
                          .set_compute_data_type(compute_data_type)
                          .set_name(name);
    if (group_offset != nullptr) {
        attributes.set_group_offset(std::move(group_offset));
    }

    auto OUT_0 = graph->reduction(input, attributes);
    return OUT_0;
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::reshape(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& input,
                 std::string const& name,
                 cudnn_frontend::ReshapeMode_t reshape_mode) {
    auto attributes = cudnn_frontend::graph::Reshape_attributes().set_name(name).set_reshape_mode(reshape_mode);

    auto OUT_0 = graph->reshape(input, attributes);
    return OUT_0;
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::transpose(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& input,
                   std::vector<int64_t> const& permutation,
                   cudnn_frontend::DataType_t const& compute_data_type,
                   std::string const& name) {
    auto attributes = cudnn_frontend::graph::Transpose_attributes()
                          .set_name(name)
                          .set_permutation(permutation)
                          .set_compute_data_type(compute_data_type);

    return graph->transpose(input, attributes);
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::concatenate(std::vector<std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>> inputs,
                     int64_t axis,
                     std::optional<int64_t> in_place_index,
                     std::string const& name) {
    auto attributes = cudnn_frontend::graph::Concatenate_attributes().set_axis(axis).set_name(name);
    if (in_place_index.has_value()) {
        attributes.set_in_place_index(in_place_index.value());
    }
    return graph->concatenate(std::move(inputs), attributes);
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::tensor_scalar(float const& value, cudnn_frontend::graph::ScalarType scalar_type) {
    return graph->tensor(value, scalar_type);
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::tensor_scalar(double const& value, cudnn_frontend::graph::ScalarType scalar_type) {
    return graph->tensor(value, scalar_type);
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::tensor_scalar(int32_t const& value, cudnn_frontend::graph::ScalarType scalar_type) {
    return graph->tensor(value, scalar_type);
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::tensor_scalar(int64_t const& value, cudnn_frontend::graph::ScalarType scalar_type) {
    return graph->tensor(value, scalar_type);
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::moe_grouped_matmul(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& token,
                            std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& weight,
                            std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& first_token_offset,
                            std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& token_index,
                            std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& token_ks,
                            cudnn_frontend::MoeGroupedMatmulMode_t const& mode,
                            cudnn_frontend::DataType_t const& compute_data_type,
                            int32_t const& top_k,
                            std::string const& name) {
    auto attributes = cudnn_frontend::graph::Moe_grouped_matmul_attributes()
                          .set_name(name)
                          .set_mode(mode)
                          .set_compute_data_type(compute_data_type)
                          .set_top_k(top_k);

    auto output = graph->moe_grouped_matmul(token, weight, first_token_offset, token_index, token_ks, attributes);
    return output;
}

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
PyGraph::moe_grouped_matmul_bwd(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& doutput,
                                std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& token,
                                std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& first_token_offset,
                                cudnn_frontend::DataType_t const& compute_data_type,
                                std::string const& name) {
    auto attributes = cudnn_frontend::graph::Moe_grouped_matmul_bwd_attributes().set_name(name).set_compute_data_type(
        compute_data_type);

    auto dweight = graph->moe_grouped_matmul_bwd(doutput, token, first_token_offset, attributes);
    return dweight;
}

void
PyGraph::validate() {
    auto status = graph->validate();
    throw_if(status.is_bad(), status.get_code(), status.get_message());
}

size_t
PyGraph::key() {
    return graph->key();
}

void
PyGraph::build_operation_graph() {
    auto status = graph->build_operation_graph(handle);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
}

std::vector<BehaviorNote_t>
PyGraph::get_behavior_notes() {
    std::vector<BehaviorNote_t> notes;
    auto status = graph->get_behavior_notes(notes);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
    return notes;
}

std::vector<BehaviorNote_t>
PyGraph::get_behavior_notes_for_plan_at_index(int64_t const index) {
    std::vector<BehaviorNote_t> notes;
    auto status = graph->get_behavior_notes_for_plan_at_index(index, notes);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
    return notes;
}

void
PyGraph::create_execution_plans(std::vector<cudnn_frontend::HeurMode_t> const& modes) {
    auto status = graph->create_execution_plans(modes);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
}

void
PyGraph::create_execution_plan(int64_t const engine_id, std::unordered_map<KnobType_t, int64_t> const& knobs) {
    auto status = graph->create_execution_plan(engine_id, knobs);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
}

int64_t
PyGraph::get_engine_count() {
    int64_t engine_count = 0;
    auto status          = graph->get_engine_count(engine_count);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
    return engine_count;
}

std::vector<Knob>
PyGraph::get_knobs_for_engine(int64_t const engine_id) {
    std::vector<Knob> knobs;
    auto status = graph->get_knobs_for_engine(engine_id, knobs);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
    return knobs;
}

void
PyGraph::build_plans(BuildPlanPolicy_t const policy) {
    // TODO: Add multithreaded support in python
    auto status = graph->build_plans(policy, false);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
}

void
PyGraph::build_plan_at_index(int64_t const index) {
    auto status = graph->build_plan_at_index(index);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
}

void
PyGraph::build(std::vector<cudnn_frontend::HeurMode_t> const& modes) {
    validate();
    build_operation_graph();
    create_execution_plans(modes);
    check_support();
    build_plans(cudnn_frontend::BuildPlanPolicy_t::HEURISTICS_CHOICE);
}

void
PyGraph::build() {
    validate();
    build_operation_graph();
}

void
PyGraph::check_support() {
    auto status = graph->check_support();
    throw_if(status.is_bad(), status.get_code(), status.get_message());
}

int64_t
PyGraph::get_workspace_size() {
    int64_t workspace = 0;

    auto status = graph->get_workspace_size(workspace);
    throw_if(status.is_bad(), status.get_code(), status.get_message());

    return workspace;
}

int64_t
PyGraph::get_workspace_size(std::optional<std::intptr_t> exec_handle,
                            py::object override_uids,
                            py::object override_shapes,
                            py::object override_strides) {
    std::vector<int64_t> override_uids_vec =
        override_uids.is_none() ? std::vector<int64_t>() : override_uids.cast<std::vector<int64_t>>();
    std::vector<std::vector<int64_t>> override_shapes_vec =
        override_shapes.is_none() ? std::vector<std::vector<int64_t>>()
                                  : override_shapes.cast<std::vector<std::vector<int64_t>>>();
    std::vector<std::vector<int64_t>> override_strides_vec =
        override_strides.is_none() ? std::vector<std::vector<int64_t>>()
                                   : override_strides.cast<std::vector<std::vector<int64_t>>>();

    int64_t workspace     = 0;
    cudnnHandle_t handle_ = exec_handle.has_value() ? static_cast<cudnnHandle_t>((void*)(exec_handle.value())) : handle;
    auto status =
        graph->get_workspace_size(handle_, workspace, override_uids_vec, override_shapes_vec, override_strides_vec);
    throw_if(status.is_bad(), status.get_code(), status.get_message());

    return workspace;
}

int64_t
PyGraph::get_workspace_size_plan_at_index(int64_t index) {
    int64_t workspace;

    auto status = graph->get_workspace_size_plan_at_index(index, workspace);
    throw_if(status.is_bad(), status.get_code(), status.get_message());

    return workspace;
}

int64_t
PyGraph::get_workspace_size_plan_at_index(int64_t index,
                                          std::optional<std::intptr_t> exec_handle,
                                          py::object override_uids,
                                          py::object override_shapes,
                                          py::object override_strides) {
    std::vector<int64_t> override_uids_vec =
        override_uids.is_none() ? std::vector<int64_t>() : override_uids.cast<std::vector<int64_t>>();
    std::vector<std::vector<int64_t>> override_shapes_vec =
        override_shapes.is_none() ? std::vector<std::vector<int64_t>>()
                                  : override_shapes.cast<std::vector<std::vector<int64_t>>>();
    std::vector<std::vector<int64_t>> override_strides_vec =
        override_strides.is_none() ? std::vector<std::vector<int64_t>>()
                                   : override_strides.cast<std::vector<std::vector<int64_t>>>();

    int64_t workspace     = 0;
    cudnnHandle_t handle_ = exec_handle.has_value() ? static_cast<cudnnHandle_t>((void*)(exec_handle.value())) : handle;
    auto status           = graph->get_workspace_size_plan_at_index(
        handle_, index, workspace, override_uids_vec, override_shapes_vec, override_strides_vec);
    throw_if(status.is_bad(), status.get_code(), status.get_message());

    return workspace;
}

std::vector<uint8_t>
PyGraph::serialize() const {
    std::vector<uint8_t> data;
    auto status = graph->serialize(data);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
    return data;
}

void
PyGraph::deserialize(std::optional<std::intptr_t> handle_, py::object const& pyobj, bool const enforce_precompiled) {
    if (py::isinstance<py::str>(pyobj)) {
        json j = json::parse(pyobj.cast<std::string>());

        auto status = graph->deserialize(j, enforce_precompiled);

        throw_if(status.is_bad(), status.get_code(), status.get_message());

    } else {
        // If handle is provided, use it (AoT compilation).
        cudnnHandle_t handle =
            handle_.has_value() ? static_cast<cudnnHandle_t>((void*)(handle_.value())) : this->handle;

        std::vector<uint8_t> data = pyobj.cast<std::vector<uint8_t>>();
        auto status               = graph->deserialize(handle, data, enforce_precompiled);

        throw_if(status.is_bad(), status.get_code(), status.get_message());
    }
}

void
PyGraph::deserialize(py::object const& pyobj, bool const enforce_precompiled) {
    // Call the overloaded version with default handle (nullopt)
    deserialize(std::nullopt, pyobj, enforce_precompiled);
}

void
PyGraph::update_cuda_graph(std::intptr_t handle,
                           std::unordered_map<cudnn_frontend::graph::Tensor_attributes::uid_t, std::intptr_t> var_pack,
                           std::intptr_t workspace,
                           std::intptr_t cuda_graph) {
    std::unordered_map<int64_t, void*> var_pack_;
    var_pack_.reserve(var_pack.size());
    for (auto const& [uid, device_pointer] : var_pack) {
        var_pack_.emplace(uid, (void*)device_pointer);
    }

    auto status = graph->update_cuda_graph(reinterpret_cast<cudnnHandle_t>(handle),
                                           var_pack_,
                                           reinterpret_cast<void*>(workspace),
                                           reinterpret_cast<cudaGraph_t>(cuda_graph));
    throw_if(status.is_bad(), status.get_code(), status.get_message());

    return;
}

void
PyGraph::populate_cuda_graph(
    std::intptr_t handle,
    std::unordered_map<cudnn_frontend::graph::Tensor_attributes::uid_t, std::intptr_t> var_pack,
    std::intptr_t workspace,
    std::intptr_t cuda_graph) {
    std::unordered_map<int64_t, void*> var_pack_;
    var_pack_.reserve(var_pack.size());
    for (auto const& [uid, device_pointer] : var_pack) {
        var_pack_.emplace(uid, (void*)device_pointer);
    }

    auto status = graph->populate_cuda_graph(reinterpret_cast<cudnnHandle_t>(handle),
                                             var_pack_,
                                             reinterpret_cast<void*>(workspace),
                                             reinterpret_cast<cudaGraph_t>(cuda_graph));
    throw_if(status.is_bad(), status.get_code(), status.get_message());

    return;
}

void
PyGraph::execute(std::unordered_map<int64_t, std::intptr_t> var_pack,
                 std::intptr_t workspace,
                 std::optional<std::intptr_t> exec_handle,
                 py::object override_uids,
                 py::object override_shapes,
                 py::object override_strides) {
    std::unordered_map<int64_t, void*> var_pack_;
    var_pack_.reserve(var_pack.size());
    for (auto const& [uid, device_pointer] : var_pack) {
        var_pack_.emplace(uid, (void*)device_pointer);
    }

    // Convert override_uids to a vector of int64_t (one-liner)
    std::vector<int64_t> override_uids_vec =
        override_uids.is_none() ? std::vector<int64_t>() : override_uids.cast<std::vector<int64_t>>();
    std::vector<std::vector<int64_t>> override_shapes_vec =
        override_shapes.is_none() ? std::vector<std::vector<int64_t>>()
                                  : override_shapes.cast<std::vector<std::vector<int64_t>>>();
    std::vector<std::vector<int64_t>> override_strides_vec =
        override_strides.is_none() ? std::vector<std::vector<int64_t>>()
                                   : override_strides.cast<std::vector<std::vector<int64_t>>>();

    auto workspace_ptr = (void*)workspace;

    cudnnHandle_t handle_ = exec_handle.has_value() ? static_cast<cudnnHandle_t>((void*)(exec_handle.value())) : handle;

    cudnn_frontend::error_t status = {error_code_t::OK, ""};
    if (override_uids_vec.empty()) {
        status = graph->execute(handle_, var_pack_, workspace_ptr);
    } else {
        status = graph->execute(
            handle_, var_pack_, workspace_ptr, override_uids_vec, override_shapes_vec, override_strides_vec);
    }
    throw_if(status.is_bad(), status.get_code(), status.get_message());
    return;
}

void
PyGraph::prepare_variant_pack_template() {
    auto status = graph->prepare_variant_pack_template();
    throw_if(status.is_bad(), status.get_code(), status.get_message());
}

void
PyGraph::execute_with_ptrs(std::vector<std::intptr_t> const& user_ptrs,
                           std::intptr_t workspace,
                           std::intptr_t exec_handle) {
    std::vector<void*> ptrs(user_ptrs.size());
    for (size_t i = 0; i < user_ptrs.size(); i++) {
        ptrs[i] = (void*)user_ptrs[i];
    }
    cudnnHandle_t h = exec_handle ? static_cast<cudnnHandle_t>((void*)exec_handle) : handle;
    auto status     = graph->execute(h, ptrs.data(), (int)ptrs.size(), (void*)workspace);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
}

void
PyGraph::execute_with_raw_ptrs(std::intptr_t user_ptrs_array,
                               int64_t n_user,
                               std::intptr_t workspace,
                               std::intptr_t exec_handle) {
    static_assert(sizeof(std::intptr_t) == sizeof(void*), "intptr_t and void* must be the same size");
    throw_if(n_user < 0, error_code_t::INVALID_VALUE, "n_user must be non-negative");
    throw_if(n_user > 0 && user_ptrs_array == 0, error_code_t::INVALID_VALUE, "user_ptrs_array is null");
    // user_ptrs_array points to a contiguous intptr_t[] of device pointers — zero copy
    void** ptrs     = reinterpret_cast<void**>(user_ptrs_array);
    cudnnHandle_t h = exec_handle ? static_cast<cudnnHandle_t>((void*)exec_handle) : handle;
    auto status     = graph->execute(h, ptrs, (int)n_user, (void*)workspace);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
}

void
PyGraph::execute_plan_at_index(std::unordered_map<int64_t, std::intptr_t> var_pack,
                               std::intptr_t workspace,
                               int64_t index,
                               std::optional<std::intptr_t> exec_handle,
                               py::object override_uids,
                               py::object override_shapes,
                               py::object override_strides) {
    std::unordered_map<int64_t, void*> var_pack_;
    for (auto const& [uid, device_pointer] : var_pack) {
        var_pack_.emplace(uid, (void*)device_pointer);
    }

    // Convert override_uids to a vector of int64_t (one-liner)
    std::vector<int64_t> override_uids_vec =
        override_uids.is_none() ? std::vector<int64_t>() : override_uids.cast<std::vector<int64_t>>();
    std::vector<std::vector<int64_t>> override_shapes_vec =
        override_shapes.is_none() ? std::vector<std::vector<int64_t>>()
                                  : override_shapes.cast<std::vector<std::vector<int64_t>>>();
    std::vector<std::vector<int64_t>> override_strides_vec =
        override_strides.is_none() ? std::vector<std::vector<int64_t>>()
                                   : override_strides.cast<std::vector<std::vector<int64_t>>>();

    auto workspace_ptr = (void*)workspace;

    cudnnHandle_t handle_ = exec_handle.has_value() ? static_cast<cudnnHandle_t>((void*)(exec_handle.value())) : handle;

    cudnn_frontend::error_t status = {error_code_t::OK, ""};
    if (override_uids_vec.empty()) {
        status = graph->execute_plan_at_index(handle_, var_pack_, workspace_ptr, index);
    } else {
        status = graph->execute_plan_at_index(
            handle_, var_pack_, workspace_ptr, index, override_uids_vec, override_shapes_vec, override_strides_vec);
    }
    throw_if(status.is_bad(), status.get_code(), status.get_message());
    return;
}

std::shared_ptr<graph::Tensor_attributes>
PyGraph::query_tensor_attributes_of_uid(int64_t const uid) const {
    graph::Tensor_attributes tensor;
    auto status = graph->query_tensor_attributes_of_uid(uid, tensor);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
    return std::make_shared<graph::Tensor_attributes>(tensor);
}

std::string
PyGraph::get_plan_name_at_index(int64_t index) {
    std::string plan_name;
    auto status = graph->get_plan_name_at_index(index, plan_name);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
    return plan_name;
}

std::pair<int64_t, std::unordered_map<KnobType_t, int64_t>>
PyGraph::get_engine_and_knobs_at_index(int64_t index) {
    int64_t engine_id = -1;
    std::unordered_map<KnobType_t, int64_t> knobs;
    auto status = graph->get_engine_and_knobs_at_index(index, engine_id, knobs);
    throw_if(status.is_bad(), status.get_code(), status.get_message());
    return {engine_id, knobs};
}

std::vector<int64_t>
default_vector(void) {
    return {};
}

void
init_pygraph_submodule(py::module_& m) {
    py::class_<PyGraph> pygraph_(m, "pygraph");
    pygraph_
        .def(py::init<std::string const&,
                      cudnn_frontend::DataType_t,
                      cudnn_frontend::DataType_t,
                      cudnn_frontend::DataType_t,
                      std::optional<std::intptr_t>,
                      py::object,
                      py::object,
                      std::shared_ptr<KernelCache>,
                      std::shared_ptr<cudnn_frontend::DeviceProperties>,
                      bool,
                      bool>(),
             py::arg_v("name", "test_graph"),
             py::arg_v("io_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("intermediate_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("handle", std::nullopt),
             py::arg_v("sm_count", py::none()),
             py::arg_v("sm_version", py::none()),
             py::arg_v("kernel_cache", nullptr),
             py::arg_v("device_property", nullptr),
             py::arg_v("is_dynamic_shape_enabled", false),
             py::arg_v("is_override_shape_enabled", false))
        .def("tensor_like",
             py::overload_cast<std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> const&, std::string const&>(
                 &PyGraph::tensor_like),
             py::arg("input"),
             py::arg_v("name", ""))
        .def("tensor_like", py::overload_cast<py::object const&>(&PyGraph::tensor_like))
        .def("_make_tensor",
             &PyGraph::tensor,
             py::arg{"dim"},
             py::arg{"stride"},
             py::arg_v("data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v{"is_virtual", false},
             py::arg_v{"is_pass_by_value", false},
             py::arg_v{"ragged_offset", nullptr},
             py::arg_v{"reordering_type", cudnn_frontend::TensorReordering_t::NONE},
             py::arg_v("name", ""),
             py::arg_v("uid", -1),
             py::arg_v("ragged_offset_multiplier", int64_t{1}))
        .def("genstats",
             &PyGraph::genstats,
             py::arg("input"),
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("name", ""))
        .def("slice",
             &PyGraph::slice,
             py::arg("input"),
             py::arg_v{"slices", default_vector()},
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("name", ""),
             R"pbdoc(
                Perform slice operation on the given input tensor.

                Args:
                    input (cudnn_tensor): The input tensor to be sliced.
                    slices (List[slice]): A list of Python slice objects, one for each dimension.
                        Per-axis step comes from each slice's ``step`` (default 1), after normalization
                        for the tensor shape (same semantics as indexing a sequence of that length).
                    compute_data_type (Optional[cudnn.data_type]): The data type for computation.
                        Default is NOT_SET.
                    name (Optional[str]): A name for the slice operation.

                Returns:
                    cudnn_tensor: The resulting sliced tensor.

                Example:
                    >>> input_tensor = graph.tensor([4, 8, 16])
                    >>> sliced_tensor = graph.slice(input_tensor, [slice(0, 2), slice(1, 8, 2), slice(0, 16)])
            )pbdoc")
        .def(
            "conv_fprop",
            [](PyGraph& self,
               std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& image,
               std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& weight,
               std::vector<int64_t> const& padding,
               std::vector<int64_t> const& stride,
               std::vector<int64_t> const& dilation,
               cudnn_frontend::ConvolutionMode_t const convolution_mode,
               cudnn_frontend::DataType_t const& compute_data_type,
               std::string const& name) {
                return self.conv_fprop(
                    image, weight, padding, padding, stride, dilation, convolution_mode, compute_data_type, name);
            },
            py::arg("image"),
            py::arg("weight"),
            py::arg_v{"padding", default_vector()},
            py::arg_v{"stride", default_vector()},
            py::arg_v{"dilation", default_vector()},
            py::arg_v{"convolution_mode", cudnn_frontend::ConvolutionMode_t::CROSS_CORRELATION},
            py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
            py::arg_v("name", ""))
        .def("conv_fprop",
             &PyGraph::conv_fprop,
             py::arg("image"),
             py::arg("weight"),
             py::arg_v{"pre_padding", default_vector()},
             py::arg_v{"post_padding", default_vector()},
             py::arg_v{"stride", default_vector()},
             py::arg_v{"dilation", default_vector()},
             py::arg_v{"convolution_mode", cudnn_frontend::ConvolutionMode_t::CROSS_CORRELATION},
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("name", ""),
             R"pbdoc(
                Perform convolution operation with the given inputs.

                Args:
                    image (cudnn_tensor): The image tensor.
                    weight (cudnn_tensor): The weight tensor.
                    pre_padding (Optional[List[int]]): The pre padding values for the operation. Default is an empty list.
                    post_padding (Optional[List[int]]): The post padding values for the operation. Default is an empty list.
                    stride (Optional[List[int]]): The stride values for the operation. Default is an empty list.
                    dilation (Optional[List[int]]): The dilation values for the operation. Default is an empty list.
                    compute_data_type (Optional[cudnn.data_type]): The data type for computation. Default is NOT_SET.
                    name (Optional[str]): A name for the operation to be performed.

                Returns:
                    cudnn_tensor: The created tensor.
            )pbdoc")
        .def(
            "conv_wgrad",
            [](PyGraph& self,
               std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& image,
               std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& loss,
               std::vector<int64_t> const& padding,
               std::vector<int64_t> const& stride,
               std::vector<int64_t> const& dilation,
               cudnn_frontend::ConvolutionMode_t const convolution_mode,
               cudnn_frontend::DataType_t const& compute_data_type,
               std::string const& name) {
                return self.conv_wgrad(
                    image, loss, padding, padding, stride, dilation, convolution_mode, compute_data_type, name);
            },
            py::arg("image"),
            py::arg("loss"),
            py::arg_v{"padding", default_vector()},
            py::arg_v{"stride", default_vector()},
            py::arg_v{"dilation", default_vector()},
            py::arg_v{"convolution_mode", cudnn_frontend::ConvolutionMode_t::CROSS_CORRELATION},
            py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
            py::arg_v("name", ""))
        .def("conv_wgrad",
             &PyGraph::conv_wgrad,
             py::arg("image"),
             py::arg("loss"),
             py::arg_v{"pre_padding", default_vector()},
             py::arg_v{"post_padding", default_vector()},
             py::arg_v{"stride", default_vector()},
             py::arg_v{"dilation", default_vector()},
             py::arg_v{"convolution_mode", cudnn_frontend::ConvolutionMode_t::CROSS_CORRELATION},
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("name", ""),
             R"pbdoc(
                Compute weight gradients using the given inputs and loss.

                Args:
                    image (cudnn_tensor): The image tensor.
                    loss (cudnn_tensor): The loss tensor.
                    pre_padding (Optional[List[int]]): The pre padding values for the operation. Default is an empty list.
                    post_padding (Optional[List[int]]): The post padding values for the operation. Default is an empty list.                    stride (Optional[List[int]]): The stride values for the operation. Default is an empty list.
                    dilation (Optional[List[int]]): The dilation values for the operation. Default is an empty list.
                    compute_data_type (Optional[cudnn.data_type]): The data type for computation. Default is NOT_SET.
                    name (Optional[str]): A name for the operation to be performed.

                Returns:
                    cudnn_tensor: The created tensor.
            )pbdoc")
        .def(
            "conv_dgrad",
            [](PyGraph& self,
               std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& loss,
               std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& filter,
               std::vector<int64_t> const& padding,
               std::vector<int64_t> const& stride,
               std::vector<int64_t> const& dilation,
               cudnn_frontend::ConvolutionMode_t const convolution_mode,
               cudnn_frontend::DataType_t const& compute_data_type,
               std::string const& name) {
                return self.conv_dgrad(
                    loss, filter, padding, padding, stride, dilation, convolution_mode, compute_data_type, name);
            },
            py::arg("loss"),
            py::arg("filter"),
            py::arg_v{"padding", default_vector()},
            py::arg_v{"stride", default_vector()},
            py::arg_v{"dilation", default_vector()},
            py::arg_v{"convolution_mode", cudnn_frontend::ConvolutionMode_t::CROSS_CORRELATION},
            py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
            py::arg_v("name", ""))
        .def("conv_dgrad",
             &PyGraph::conv_dgrad,
             py::arg("loss"),
             py::arg("filter"),
             py::arg_v{"pre_padding", default_vector()},
             py::arg_v{"post_padding", default_vector()},
             py::arg_v{"stride", default_vector()},
             py::arg_v{"dilation", default_vector()},
             py::arg_v{"convolution_mode", cudnn_frontend::ConvolutionMode_t::CROSS_CORRELATION},
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("name", ""),
             R"pbdoc(
                Compute filter gradients using the given inputs and loss.

                Args:
                    loss (cudnn_tensor): The loss tensor.
                    filter (cudnn_tensor): The filter tensor.
                    pre_padding (Optional[List[int]]): The pre padding values for the operation. Default is an empty list.
                    post_padding (Optional[List[int]]): The post padding values for the operation. Default is an empty list.
                    stride (Optional[List[int]]): The stride values for the operation. Default is an empty list.
                    dilation (Optional[List[int]]): The dilation values for the operation. Default is an empty list.
                    compute_data_type (Optional[pycudnn.data_type]): The data type for computation. Default is NOT_SET.
                    name (Optional[str]): A name for the operation to be performed.

                Returns:
                    cudnn_tensor: The created tensor.
            )pbdoc")
        .def("matmul",
             &PyGraph::matmul,
             py::arg("A"),
             py::arg("B"),
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("padding", 0.0),
             py::arg_v("name", ""),
             R"pbdoc(
                Perform matrix multiplication of two tensors A and B.

                Args:
                    A (cudnn_tensor): The first tensor.
                    B (cudnn_tensor): The second matrix tensor.
                    compute_data_type (Optional[cudnn.data_type]): The data type for computation. Default is NOT_SET.
                    name (Optional[str]): A name for the operation to be performed.

                Returns:
                    cudnn_tensor: The result of the matrix multiplication.
            )pbdoc")
        .def("reduction",
             &PyGraph::reduction,
             py::arg("input"),
             py::arg("mode"),
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("name", ""),
             py::arg_v("group_offset", nullptr),
             R"pbdoc(
                Reduce an input tensor along certain dimensions. These dimensions to reduce on are inferred from output tensor shape.

                Args:
                    input (cudnn_tensor): The input tensor.
                    mode (cudnn.reduction_mode): The mode to use to reduce along a dimension.
                    compute_data_type (Optional[cudnn.data_type]): The data type for computation. Default is NOT_SET.
                    name (Optional[str]): A name for the operation to be performed.
                    group_offset (Optional[cudnn_tensor]): The group offset tensor for grouped reduction.

                Returns:
                    cudnn_tensor: The result of reduction operation.
            )pbdoc")
        .def("block_scale_dequantize",
             &PyGraph::block_scale_dequantize,
             py::arg("input"),
             py::arg("descale"),
             py::arg("block_size"),
             py::arg_v("is_negative_scale", false),
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("name", ""),
             R"pbdoc(
                Dequantize an input tensor to other dimensions without changing the actual memory layout.
                
                Args:
                    input (cudnn_tensor): The input tensor to dequantize.
                    descale (cudnn_tensor): The scale tensor for dequantization.
                    block_size (List[int]): The block size for dequantization.
                    is_negative_scale (Optional[bool]): Whether the scale values can be negative. Default is False.
                    compute_data_type (Optional[cudnn.data_type]): The data type for computation. Default is NOT_SET.
                    name (Optional[str]): A name for the operation.
                
                Returns:
                    cudnn_tensor: The dequantized output tensor.
            )pbdoc")
        .def("block_scale_quantize",
             &PyGraph::block_scale_quantize,
             py::arg("input"),
             py::arg("block_size"),
             py::arg_v("axis", std::nullopt),
             py::arg_v("transpose", false),
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("name", ""),
             R"pbdoc(
                Quantize an input tensor with block scaling to produce quantized output and scale tensors.
                
                Args:
                    input (cudnn_tensor): The input tensor to quantize.
                    block_size (int): The block size for quantization.
                    axis (Optional[int]): The axis along which to quantize. Default is None.
                    transpose (Optional[bool]): Whether to transpose during quantization. Default is False.
                    compute_data_type (Optional[cudnn.data_type]): The data type for computation. Default is NOT_SET.
                    name (Optional[str]): A name for the operation.
                
                Returns:
                    Tuple[cudnn_tensor, cudnn_tensor]: A tuple of (quantized_output, scale) tensors.
            )pbdoc")
        .def("reshape",
             &PyGraph::reshape,
             py::arg("input"),
             py::arg_v("name", ""),
             py::arg_v("reshape_mode", cudnn_frontend::ReshapeMode_t::VIEW_ONLY),
             R"pbdoc(
                Reshape an input tensor to other dimensions without changing the actual memory layout.
                These dimensions to reshape to are inferred from output tensor shape.

                Args:
                    input (cudnn_tensor): The input tensor.
                    name (Optional[str]): A name for the operation to be performed.
                    reshape_mode (cudnn.reshape_mode): VIEW_ONLY (default) or LOGICAL for lexicographic logical reshapes.

                Returns:
                    cudnn_tensor: The result of reshape operation. Please set the dims for the output tensor.
            )pbdoc")
        .def("transpose",
             &PyGraph::transpose,
             py::arg("input"),
             py::arg("permutation"),
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET),
             py::arg_v("name", ""),
             R"pbdoc(
                Permute tensor dimensions using a permutation vector (output axis i reads input axis permutation[i]).

                Args:
                    input (cudnn_tensor): The input tensor.
                    permutation (List[int]): Permutation of axis indices.
                    compute_data_type (Optional[cudnn.data_type]): Optional compute type; default NOT_SET.
                    name (Optional[str]): Operation name.

                Returns:
                    cudnn_tensor: Transposed tensor (dims/strides inferred when not set).
            )pbdoc")
        .def("concatenate",
             &PyGraph::concatenate,
             py::arg("inputs"),
             py::arg("axis"),
             py::arg_v("in_place_index", std::optional<int64_t>{}),
             py::arg_v("name", ""),
             R"pbdoc(
                Concatenate tensors along an axis.

                Args:
                    inputs (List[cudnn_tensor]): Tensors to concatenate.
                    axis (int): Concatenation axis.
                    in_place_index (Optional[int]): When set, optional in-place concat per backend semantics.
                    name (Optional[str]): Operation name.

                Returns:
                    cudnn_tensor: Concatenated output tensor.
            )pbdoc")
        .def("tensor_scalar",
             py::overload_cast<float const&, cudnn_frontend::graph::ScalarType>(&PyGraph::tensor_scalar),
             py::arg("value"),
             py::arg("scalar_type"),
             R"pbdoc(
                Create a rank-1 scalar tensor from a Python float, marked runtime or compile-time.

                Args:
                    value (float): Scalar value.
                    scalar_type (cudnn.scalar_type): RUNTIME_PARAM or COMPILE_TIME_CONST.

                Returns:
                    cudnn_tensor: Scalar tensor (set dim/stride/name as needed for your graph).
            )pbdoc")
        .def("tensor_scalar",
             py::overload_cast<double const&, cudnn_frontend::graph::ScalarType>(&PyGraph::tensor_scalar),
             py::arg("value"),
             py::arg("scalar_type"))
        .def("tensor_scalar",
             py::overload_cast<int32_t const&, cudnn_frontend::graph::ScalarType>(&PyGraph::tensor_scalar),
             py::arg("value"),
             py::arg("scalar_type"))
        .def("tensor_scalar",
             py::overload_cast<int64_t const&, cudnn_frontend::graph::ScalarType>(&PyGraph::tensor_scalar),
             py::arg("value"),
             py::arg("scalar_type"))
        .def("moe_grouped_matmul",
             &PyGraph::moe_grouped_matmul,
             py::arg("token"),
             py::arg("weight"),
             py::arg("first_token_offset"),
             py::arg_v("token_index", nullptr),
             py::arg_v("token_ks", nullptr),
             py::arg_v("mode", cudnn_frontend::MoeGroupedMatmulMode_t::NONE),
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::FLOAT),
             py::arg_v("top_k", 0),
             py::arg_v("name", ""),
             R"pbdoc(
                Perform MoE Grouped Matmul operation.

                Args:
                    token (cudnn_tensor): The token tensor.
                    weight (cudnn_tensor): The weight tensor.
                    first_token_offset (cudnn_tensor): The first token offset tensor.
                    token_index (cudnn_tensor): The token index tensor or nullptr.
                    token_ks (cudnn_tensor): The token ks tensor or nullptr.
                    mode (cudnn.moe_grouped_matmul_mode): The mode of the operation.
                    compute_data_type (cudnn.data_type): The data type for computation.
                    top_k (int): The top k value.
                    name (str): The name of the operation.
            )pbdoc")
        .def("moe_grouped_matmul_bwd",
             &PyGraph::moe_grouped_matmul_bwd,
             py::arg("doutput"),
             py::arg("token"),
             py::arg("first_token_offset"),
             py::arg_v("compute_data_type", cudnn_frontend::DataType_t::FLOAT),
             py::arg_v("name", ""),
             R"pbdoc(
                Perform MoE Grouped Matmul Bwd operation.

                Args:
                    doutput (cudnn_tensor): The doutput tensor.
                    token (cudnn_tensor): The token tensor.
                    first_token_offset (cudnn_tensor): The first token offset tensor.
                    compute_data_type (cudnn.data_type): The data type for computation.
                    name (str): The name of the operation.
            )pbdoc")
        .def("get_behavior_notes", &PyGraph::get_behavior_notes)
        .def("get_behavior_notes_for_plan_at_index", &PyGraph::get_behavior_notes_for_plan_at_index)
        .def("deselect_engines", &PyGraph::deselect_engines)
        .def("deselect_numeric_notes", &PyGraph::deselect_numeric_notes)
        .def("deselect_behavior_notes", &PyGraph::deselect_behavior_notes)
        .def("select_numeric_notes", &PyGraph::select_numeric_notes)
        .def("select_behavior_notes", &PyGraph::select_behavior_notes)
        .def("deselect_workspace_greater_than", &PyGraph::deselect_workspace_greater_than)
        .def("validate", &PyGraph::validate)
        .def("key", &PyGraph::key)
        .def("build_operation_graph", &PyGraph::build_operation_graph)
        .def("create_execution_plans", &PyGraph::create_execution_plans)
        .def("create_execution_plan",
             &PyGraph::create_execution_plan,
             R"pbdoc(
                Gets the knob configurations available for the given engine.
                Args:
                    engine_id (int): The ID of the engine to create the execution plan on.
                    knobs (dict[Knob, int]): The map of knobs to knob values.
            )pbdoc")
        .def("get_engine_count", &PyGraph::get_engine_count)
        .def("get_knobs_for_engine",
             &PyGraph::get_knobs_for_engine,
             R"pbdoc(
                Gets the knob configurations available for the given engine.
                Args:
                    engine_id (int): The ID of the engine to query knob configurations for.
            )pbdoc")
        .def("check_support", &PyGraph::check_support)
        .def("build_plans",
             &PyGraph::build_plans,
             py::arg("policy") = cudnn_frontend::BuildPlanPolicy_t::HEURISTICS_CHOICE)
        .def("build_plan_at_index",
             &PyGraph::build_plan_at_index,
             py::arg("index"),
             R"pbdoc(
                Build a plan at the given index.
                Args:
                    index (int): The index of the plan to build.
            )pbdoc")
        .def("build", (void (PyGraph::*)(std::vector<cudnn_frontend::HeurMode_t> const&))&PyGraph::build)
        .def("build", (void (PyGraph::*)())&PyGraph::build)
        .def("get_execution_plan_count",
             &PyGraph::get_execution_plan_count,
             R"pbdoc(
                Get the number of execution plan candidates.
            )pbdoc")
        .def("get_workspace_size", (int64_t (PyGraph::*)())&PyGraph::get_workspace_size)
        .def("get_workspace_size",
             (int64_t (PyGraph::*)(
                 std::optional<std::intptr_t>, py::object, py::object, py::object))&PyGraph::get_workspace_size,
             py::arg("handle")           = std::nullopt,
             py::arg("override_uids")    = py::none(),
             py::arg("override_shapes")  = py::none(),
             py::arg("override_strides") = py::none())
        .def("get_workspace_size_plan_at_index",
             (int64_t (PyGraph::*)(int64_t))&PyGraph::get_workspace_size_plan_at_index,
             py::arg("index"),
             R"pbdoc(
                Get workspace for a plan at the given index.
                Args:
                    index (int): The index of the plan to get workspace from.
                    If the graph is not built at the index, this will return 0.
            )pbdoc")
        .def("get_workspace_size_plan_at_index",
             (int64_t (PyGraph::*)(int64_t,
                                   std::optional<std::intptr_t>,
                                   py::object,
                                   py::object,
                                   py::object))&PyGraph::get_workspace_size_plan_at_index,
             py::arg("index"),
             py::arg("handle")           = std::nullopt,
             py::arg("override_uids")    = py::none(),
             py::arg("override_shapes")  = py::none(),
             py::arg("override_strides") = py::none())
        .def("query_tensor_attributes_of_uid",
             &PyGraph::query_tensor_attributes_of_uid,
             py::arg("uid"),
             R"pbdoc(
                    Get tensor_attributes for a given UID
                    Args:
                    uid (int): The uid of tensor to be queried
                    If the graph does not have the UID, this will raise an error
                )pbdoc")
        .def("get_plan_name_at_index",
             &PyGraph::get_plan_name_at_index,
             py::arg("index"),
             R"pbdoc(
                    Get the name for a plan at the given index.
                    Args:
                    index (int): The index of the plan to get workspace from.
                )pbdoc")
        .def("get_engine_and_knobs_at_index",
             &PyGraph::get_engine_and_knobs_at_index,
             py::arg("index"),
             R"pbdoc(
                    Get the engine id and knob choices for a plan at the given index.
                    Structured counterpart of get_plan_name_at_index; the result can be
                    passed to create_execution_plan() to rebuild the same kernel.
                    Args:
                    index (int): The index of the plan to query.
                )pbdoc")
        .def("_execute",
             &PyGraph::execute,
             py::arg("var_pack"),
             py::arg("workspace"),
             py::arg("handle"),
             py::arg("override_uids")    = py::none(),
             py::arg("override_shapes")  = py::none(),
             py::arg("override_strides") = py::none())
        .def("_prepare_variant_pack_template", &PyGraph::prepare_variant_pack_template)
        .def("_get_variant_pack_uids_sorted", &PyGraph::get_variant_pack_uids_sorted)
        .def("_execute_with_ptrs",
             &PyGraph::execute_with_ptrs,
             py::arg("user_ptrs"),
             py::arg("workspace"),
             py::arg("handle"))
        .def("_execute_with_raw_ptrs",
             &PyGraph::execute_with_raw_ptrs,
             py::arg("user_ptrs_array"),
             py::arg("n_user"),
             py::arg("workspace"),
             py::arg("handle"))
        .def("populate_cuda_graph", &PyGraph::populate_cuda_graph)
        .def("update_cuda_graph", &PyGraph::update_cuda_graph)
        .def("serialize", &PyGraph::serialize)
        .def("deserialize",
             (void (PyGraph::*)(std::optional<std::intptr_t>, py::object const&, bool const))&PyGraph::deserialize,
             py::arg("handle_"),
             py::arg("pyobj"),
             py::arg("enforce_precompiled") = false)
        .def("deserialize",
             (void (PyGraph::*)(py::object const&, bool const))&PyGraph::deserialize,
             py::arg("pyobj"),
             py::arg("enforce_precompiled") = false)
        .def("_execute_plan_at_index",
             &PyGraph::execute_plan_at_index,
             py::arg("var_pack"),
             py::arg("workspace"),
             py::arg("index"),
             py::arg("handle"),
             py::arg("override_uids")    = py::none(),
             py::arg("override_shapes")  = py::none(),
             py::arg("override_strides") = py::none())
        .def("__repr__", [](PyGraph const& pygraph) {
            std::stringstream ss;
            json j = pygraph.graph;
            ss << j.dump(4);
            return ss.str();
        });

    m.def("_get_data_ptr", &extract_data_pointer);

    init_pygraph_norm_submodule(pygraph_);
    init_pygraph_sdpa_submodule(pygraph_);
    init_pygraph_pointwise_submodule(pygraph_);
}

}  // namespace cudnn_frontend::python_bindings
