Kernels:

kernels-community
/

megablocks

Kernel card Files Files and versions

megablocks / csrc /new_cumsum.h

drbh

Migrated from kernels-community/megablocks

3db0cde unverified 17 days ago

history blame contribute delete

283 Bytes

	#pragma once

	#include <torch/all.h>

	namespace megablocks {

	// Forward declarations for the public interface functions
	void exclusive_cumsum(torch::Tensor x, int dim, torch::Tensor out);
	void inclusive_cumsum(torch::Tensor x, int dim, torch::Tensor out);

	} // namespace megablocks