#pragma once

#include <torch/torch.h>

void poly_norm(torch::Tensor &out, torch::Tensor &input, torch::Tensor &weights, torch::Tensor &bias, double eps);
void poly_norm_backward(torch::Tensor& input_grad, torch::Tensor& weight_grad, torch::Tensor& bias_grad, torch::Tensor& output_grad, torch::Tensor& input, torch::Tensor& weight, double eps);