py::class_<BC_GatedDeltaNet, std::shared_ptr<BC_GatedDeltaNet>>(m, "BC_GatedDeltaNet").def
(
    py::init<
        at::Tensor,
        at::Tensor,
        at::Tensor,
        at::Tensor,
        at::Tensor,
        at::Tensor,
        at::Tensor,
        at::Tensor,
        at::Tensor,
        at::Tensor,
        std::shared_ptr<BC_LinearEXL3>,
        std::shared_ptr<BC_LinearFP16>,
        at::Tensor,
        at::Tensor,
        int,
        int,
        int,
        int,
        at::Tensor,
        c10::optional<at::Tensor>,
        std::shared_ptr<BC_GatedRMSNorm>,
        std::shared_ptr<BC_LinearEXL3>,
        float
    >(),
    py::arg("mixed_qkv"),
    py::arg("z"),
    py::arg("beta"),
    py::arg("g"),
    py::arg("qkvz"),
    py::arg("ba"),
    py::arg("conv_temp_a"),
    py::arg("conv_temp_b"),
    py::arg("core_attn_out"),
    py::arg("core_attn_out_f"),
    py::arg("qkvz_proj"),
    py::arg("ba_proj"),
    py::arg("dt_bias"),
    py::arg("a_log"),
    py::arg("num_k_heads"),
    py::arg("num_v_heads"),
    py::arg("k_head_dim"),
    py::arg("v_head_dim"),
    py::arg("conv1d_weight"),
    py::arg("conv1d_bias"),
    py::arg("norm"),
    py::arg("o_proj"),
    py::arg("beta_scale")
)
.def("run_bsz1_a", &BC_GatedDeltaNet::run_bsz1_a)
.def("run_bsz1_b", &BC_GatedDeltaNet::run_bsz1_b);
