#pragma once

#include <ATen/core/Tensor.h>

#include <cstdint>

namespace at::native {

Tensor& embedding_bag_byte_rowwise_offsets_out(

Tensor& output,

const Tensor& weight,

const Tensor& indices,

const std::optional<Tensor>& offsets_in,

const bool /* scale_grad_by_freq */,

const int64_t /* mode */,

bool pruned_weights,

const std::optional<Tensor>& per_sample_weights_,

const std::optional<Tensor>& compressed_indices_mapping,

bool include_last_offset);

Tensor& embedding_bag_4bit_rowwise_offsets_out(

Tensor& output,

const Tensor& weight,

const Tensor& indices,

const std::optional<Tensor>& offsets_in,

const bool /* scale_grad_by_freq */,

const int64_t /* mode */,

bool pruned_weights,

const std::optional<Tensor>& per_sample_weights_,

const std::optional<Tensor>& compressed_indices_mapping,

bool include_last_offset);

Tensor& qembeddingbag_byte_unpack_out(Tensor& output, const Tensor& packed_weight);

} // namespace at::native