aten/src/ATen/native/cuda/MultiLabelMarginCriterion.cu - platform/external/pytorch - Git at Google

 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
 #include <ATen/core/Tensor.h>
 #include <ATen/AccumulateType.h>
 #include <ATen/Dispatch.h>
 #include <c10/macros/Macros.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <ATen/native/cuda/block_reduce.cuh>

 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
 #include <ATen/CUDAFunctions.h>
 #include <ATen/NativeFunctions.h>
 #else
 #include <ATen/ops/empty.h>
 #include <ATen/ops/zeros_like.h>
 #include <ATen/ops/sum_cuda_dispatch.h>
 #include <ATen/ops/multilabel_margin_loss.h>
 #endif


 namespace at::native {

 namespace {
 const int MULTILABELMARGIN_THREADS = 128;

 void multilabel_margin_loss_shape_check(
     int64_t& nframe,
     int64_t& dim,
     const int64_t& ndims,
     const Tensor& input,
     const Tensor& target) {
     TORCH_CHECK(
         (ndims == 2 && input.size(1) != 0) || (ndims == 1 && input.size(0) != 0) || ndims == 0,
         "Expected non-empty vector or matrix with optional 0-dim batch size, but got: ",
         input.sizes());

     if (ndims <= 1) {
       nframe = 1;
       dim = ndims == 0 ? 1 : input.size(0);
       TORCH_CHECK(
           target.dim() <= 1 && target.numel() == dim,
           "inconsistent target size: ", target.sizes(), " for input of size: ",
           input.sizes());
     } else {
       nframe = input.size(0);
       dim = input.size(1);
       TORCH_CHECK(
           target.dim() == 2 && target.size(0) == nframe &&
           target.size(1) == dim,
           "inconsistent target size: ", target.sizes(), " for input of size: ",
           input.sizes());
     }
 }

 template <typename scalar_t, typename accscalar_t>
 C10_LAUNCH_BOUNDS_1(MULTILABELMARGIN_THREADS)
 __global__ void multilabel_margin_loss_forward_kernel(
     scalar_t* output,
     const scalar_t* input,
     const int64_t* target,
     scalar_t* is_target,
     int nframe,
     int dim,
     bool size_average) {

   // vectors:
   int k = blockIdx.x;
   const scalar_t* input_k = input + k * dim;
   const int64_t* target_k = target + k * dim;
   scalar_t* output_k = output + k;
   scalar_t* is_target_k = is_target + k * dim;

   // zero is_target
   for (int d = threadIdx.x; d < dim; d += blockDim.x) {
     is_target_k[d] = static_cast<scalar_t>(0);
   }
   __syncthreads();

   // mark targets in is_target
   if (threadIdx.x == 0) {
     for (int dt = 0; dt < dim; dt++) {
       int target_idx = target_k[dt];
       if (target_idx < 0) {
         break;
       }
       is_target_k[target_idx] = static_cast<scalar_t>(1);
     }
   }
   __syncthreads();

   // iterate over targets
   accscalar_t sum = 0;
   for (int dt = 0; dt < dim; dt++) {
     // next target:
     int target_idx = target_k[dt];
     if (target_idx < 0) {
       break;
     }

     // current value for target
     scalar_t input_target_k = input_k[target_idx];

     // compare to all inputs (multithreaded):
     for (int d = threadIdx.x; d < dim; d += blockDim.x) {
       // contribute to loss only if not a target
       if (!static_cast<int>(is_target_k[d])) {
         scalar_t z = 1 - input_target_k + input_k[d];
         if (z > 0) {
           sum += z;
         }
       }
     }
   }

   // Temporary sums (for mapreduce)
   __shared__ accscalar_t smem[MULTILABELMARGIN_THREADS];
   accscalar_t total_sum = cuda_utils::BlockReduceSum(sum, smem);
   if (threadIdx.x == 0) {
     if (size_average) {
       *output_k = static_cast<scalar_t>((total_sum / dim) / nframe);
     } else {
       *output_k = static_cast<scalar_t>(total_sum / dim);
     }
   }
 }

 template <typename scalar_t, typename accscalar_t>
 C10_LAUNCH_BOUNDS_1(MULTILABELMARGIN_THREADS)
 __global__ void multilabel_margin_loss_backward_kernel(
     scalar_t* grad_input,
     const scalar_t* grad_output,
     const scalar_t* input,
     const int64_t* target,
     const scalar_t* is_target,
     int nframe,
     int dim,
     bool size_average,
     bool reduce) {

   int k = blockIdx.x;
   const scalar_t* input_k = input + k * dim;
   scalar_t* grad_input_k = grad_input + k * dim;
   const int64_t* target_k = target + k * dim;
   const scalar_t* is_target_k = is_target + k * dim;

   const scalar_t* grad_output_k = grad_output;
   if (!reduce) {
     grad_output_k += k;
   }

   // gain:
   scalar_t g = static_cast<scalar_t>(
       size_average && reduce ? 1. / static_cast<accscalar_t>(nframe * dim)
                              : 1. / static_cast<accscalar_t>(dim));

   // zero gradients:
   for (int d = threadIdx.x; d < dim; d += blockDim.x) {
     grad_input_k[d] = static_cast<scalar_t>(0);
   }
   __syncthreads();

   // iterate over targets
   for (int dt = 0; dt < dim; dt++) {
     // next target:
     int target_idx = static_cast<int>(target_k[dt]);
     if (target_idx < 0) {
       break;
     }

     // current value for target
     scalar_t input_target_k = input_k[target_idx];

     // compare to all inputs (multithreaded):
     accscalar_t sum = 0;
     for (int d = threadIdx.x; d < dim; d += blockDim.x) {
       // contribute to loss only if not a target
       if (!static_cast<int>(is_target_k[d])) {
         scalar_t z = 1 - input_target_k + input_k[d];
         if (z > 0) {
           sum -= g;
           grad_input_k[d] += g;
         }
       }
     }
     __syncthreads();

     // Temporary sums (for mapreduce)
     __shared__ accscalar_t smem[MULTILABELMARGIN_THREADS];
     accscalar_t total_sum = cuda_utils::BlockReduceSum(sum, smem);
     if (threadIdx.x == 0) {
       grad_input_k[target_idx] += static_cast<scalar_t>(total_sum);
     }
   }

   for (int d = threadIdx.x; d < dim; d += blockDim.x) {
     grad_input_k[d] *= *grad_output_k;
   }
 }

 void multilabel_margin_loss_forward_out_cuda_template(
     const Tensor& input,
     const Tensor& target,
     int64_t reduction,
     Tensor& output,
     Tensor& is_target) {
   int64_t nframe, dim;
   const int64_t ndims = input.dim();
   multilabel_margin_loss_shape_check(nframe, dim, ndims, input, target);

   if (input.numel() == 0) {
     return;
   }

   auto input_ = input.contiguous();
   auto target_ = target.contiguous();
   auto is_target_ = is_target.contiguous();
   is_target_.resize_as_(target);

   if (input.dim() <= 1) {
     output.resize_({});

     dim3 blocks(1);
     dim3 threads(MULTILABELMARGIN_THREADS);

     AT_DISPATCH_FLOATING_TYPES_AND2(
         at::ScalarType::Half,
         at::ScalarType::BFloat16,
         input.scalar_type(),
         "multilabel_margin_loss_forward_kernel",
         [&] {
           using accscalar_t = at::acc_type<scalar_t, true>;
           multilabel_margin_loss_forward_kernel<scalar_t, accscalar_t>
               <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
                   output.mutable_data_ptr<scalar_t>(),
                   input_.const_data_ptr<scalar_t>(),
                   target_.const_data_ptr<int64_t>(),
                   is_target_.mutable_data_ptr<scalar_t>(),
                   1,
                   dim,
                   reduction == at::Reduction::Mean);
           C10_CUDA_KERNEL_LAUNCH_CHECK();
         });
   } else if (input.dim() == 2) {
     dim3 blocks(input.size(0));
     dim3 threads(MULTILABELMARGIN_THREADS);

     if (reduction != at::Reduction::None) {
       auto output_tmp = at::empty({input_.size(0)}, input_.options());
       output.resize_({});
       AT_DISPATCH_FLOATING_TYPES_AND2(
           at::ScalarType::Half,
           at::ScalarType::BFloat16,
           input.scalar_type(),
           "multilabel_margin_loss_forward_kernel",
           [&] {
             using accscalar_t = at::acc_type<scalar_t, true>;
             multilabel_margin_loss_forward_kernel<scalar_t, accscalar_t>
                 <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
                     output_tmp.mutable_data_ptr<scalar_t>(),
                     input_.const_data_ptr<scalar_t>(),
                     target_.const_data_ptr<int64_t>(),
                     is_target_.mutable_data_ptr<scalar_t>(),
                     nframe,
                     dim,
                     reduction == at::Reduction::Mean);
             C10_CUDA_KERNEL_LAUNCH_CHECK();
           });
       at::cuda::sum_out(
           output,
           output_tmp,
           at::IntArrayRef(std::vector<int64_t>{}),
           false,
           output.scalar_type());
     } else {
       output.resize_({input.size(0)});
       AT_DISPATCH_FLOATING_TYPES_AND2(
           at::ScalarType::Half,
           at::ScalarType::BFloat16,
           input.scalar_type(),
           "multilabel_margin_loss_forward_kernel",
           [&] {
             using accscalar_t = at::acc_type<scalar_t, true>;
             multilabel_margin_loss_forward_kernel<scalar_t, accscalar_t>
                 <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
                     output.mutable_data_ptr<scalar_t>(),
                     input_.const_data_ptr<scalar_t>(),
                     target_.const_data_ptr<int64_t>(),
                     is_target_.mutable_data_ptr<scalar_t>(),
                     nframe,
                     dim,
                     false);
             C10_CUDA_KERNEL_LAUNCH_CHECK();
           });
     }

   } else {
     TORCH_CHECK(
         false,
         "Expected 2D input with optional zero batch dim, or 1D input with non-zero dims, but got sizes: ",
         input.sizes());
   }
 }

 void multilabel_margin_loss_backward_cuda_out_template(
     const Tensor& grad_output,
     const Tensor& input,
     const Tensor& target,
     int64_t reduction,
     const Tensor& is_target,
     Tensor& grad_input) {
   int64_t nframe, dim;
   const int64_t ndims = input.dim();
   multilabel_margin_loss_shape_check(nframe, dim, ndims, input, target);

   if (input.numel() == 0) {
     return;
   }

   auto input_ = input.contiguous();
   auto target_ = target.contiguous();
   auto is_target_ = is_target.contiguous();
   auto grad_output_ = grad_output.contiguous();
   grad_input.resize_as_(input_);

   if (grad_input.dim() <= 1) {
     int target_size = target_.dim() == 0 ? 1 : target_.size(0);
     TORCH_CHECK(
         (target_.numel() != 0) && (target_.dim() <= 1) && (target_size == dim),
         "inconsistent target size");
     TORCH_CHECK(
         target_.sizes() == is_target_.sizes(), "inconsistent is_target size");
     dim3 blocks(1);
     dim3 threads(MULTILABELMARGIN_THREADS);

     AT_DISPATCH_FLOATING_TYPES_AND2(
         at::ScalarType::Half,
         at::ScalarType::BFloat16,
         input.scalar_type(),
         "multilabel_margin_loss_backward_kernel",
         [&] {
           using accscalar_t = at::acc_type<scalar_t, true>;
           multilabel_margin_loss_backward_kernel<scalar_t, accscalar_t>
               <<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
                   grad_input.mutable_data_ptr<scalar_t>(),
                   grad_output_.const_data_ptr<scalar_t>(),
                   input_.const_data_ptr<scalar_t>(),
                   target_.const_data_ptr<int64_t>(),
                   is_target_.const_data_ptr<scalar_t>(),
                   1,
                   dim,
                   reduction == at::Reduction::Mean,
                   reduction != at::Reduction::None);
           C10_CUDA_KERNEL_LAUNCH_CHECK();
         });
   } else if (grad_input.dim() == 2) {
     TORCH_CHECK(
         (input_.size(1) != 0) && (target_.dim() == 2) &&
             (target_.size(0) == nframe) && (target_.size(1) == dim),
         "inconsistent target size");
     TORCH_CHECK(target_.sizes() == is_target_.sizes(), "inconsistent is_target size");
     dim3 blocks(grad_input.size(0));
     dim3 threads(MULTILABELMARGIN_THREADS);

     AT_DISPATCH_FLOATING_TYPES_AND2(
         at::ScalarType::Half,
         at::ScalarType::BFloat16,
         input.scalar_type(),
         "multilabel_margin_loss_backward_kernel",
         [&] {
           using accscalar_t = at::acc_type<scalar_t, true>;
           multilabel_margin_loss_backward_kernel<scalar_t, accscalar_t>
               <<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
                   grad_input.mutable_data_ptr<scalar_t>(),
                   grad_output_.const_data_ptr<scalar_t>(),
                   input_.const_data_ptr<scalar_t>(),
                   target_.const_data_ptr<int64_t>(),
                   is_target_.const_data_ptr<scalar_t>(),
                   grad_input.size(0),
                   grad_input.size(1),
                   reduction == at::Reduction::Mean,
                   reduction != at::Reduction::None);
           C10_CUDA_KERNEL_LAUNCH_CHECK();
         });
   } else {
     TORCH_CHECK(
         false,
         "Expected 2D input with optional zero batch dim, or 1D input with non-zero dims, but got sizes: ",
         grad_input.sizes());
   }
 }

 } // namespace

 std::tuple<Tensor&, Tensor&> multilabel_margin_loss_forward_out_cuda(
     const Tensor& self,
     const Tensor& target,
     int64_t reduction,
     Tensor& output,
     Tensor& is_target) {
   multilabel_margin_loss_forward_out_cuda_template(
       self, target, reduction, output, is_target);
   return std::tuple<Tensor&, Tensor&>(output, is_target);
 }

 std::tuple<Tensor, Tensor> multilabel_margin_loss_forward_cuda(
     const Tensor& self,
     const Tensor& target,
     int64_t reduction) {
   auto output = at::empty({0}, self.options());
   auto is_target = at::empty({0}, self.options());
   multilabel_margin_loss_forward_out_cuda_template(
       self, target, reduction, output, is_target);
   return std::make_tuple(output, is_target);
 }

 Tensor& multilabel_margin_loss_backward_cuda_out(
     const Tensor& grad_output,
     const Tensor& self,
     const Tensor& target,
     int64_t reduction,
     const Tensor& is_target,
     Tensor& grad_input) {
   multilabel_margin_loss_backward_cuda_out_template(
       grad_output, self, target, reduction, is_target, grad_input);
   return grad_input;
 }

 Tensor multilabel_margin_loss_backward_cuda(
     const Tensor& grad_output,
     const Tensor& self,
     const Tensor& target,
     int64_t reduction,
     const Tensor& is_target) {
   auto grad_input = at::zeros_like(self, LEGACY_CONTIGUOUS_MEMORY_FORMAT);
   multilabel_margin_loss_backward_cuda_out_template(
       grad_output, self, target, reduction, is_target, grad_input);
   return grad_input;
 }

 } // namespace at::native
	#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
	#include <ATen/core/Tensor.h>
	#include <ATen/AccumulateType.h>
	#include <ATen/Dispatch.h>
	#include <c10/macros/Macros.h>
	#include <ATen/cuda/CUDAContext.h>
	#include <ATen/native/cuda/block_reduce.cuh>

	#ifndef AT_PER_OPERATOR_HEADERS
	#include <ATen/Functions.h>
	#include <ATen/CUDAFunctions.h>
	#include <ATen/NativeFunctions.h>
	#else
	#include <ATen/ops/empty.h>
	#include <ATen/ops/zeros_like.h>
	#include <ATen/ops/sum_cuda_dispatch.h>
	#include <ATen/ops/multilabel_margin_loss.h>
	#endif


	namespace at::native {

	namespace {
	const int MULTILABELMARGIN_THREADS = 128;

	void multilabel_margin_loss_shape_check(
	int64_t& nframe,
	int64_t& dim,
	const int64_t& ndims,
	const Tensor& input,
	const Tensor& target) {
	TORCH_CHECK(
	(ndims == 2 && input.size(1) != 0) \|\| (ndims == 1 && input.size(0) != 0) \|\| ndims == 0,
	"Expected non-empty vector or matrix with optional 0-dim batch size, but got: ",
	input.sizes());

	if (ndims <= 1) {
	nframe = 1;
	dim = ndims == 0 ? 1 : input.size(0);
	TORCH_CHECK(
	target.dim() <= 1 && target.numel() == dim,
	"inconsistent target size: ", target.sizes(), " for input of size: ",
	input.sizes());
	} else {
	nframe = input.size(0);
	dim = input.size(1);
	TORCH_CHECK(
	target.dim() == 2 && target.size(0) == nframe &&
	target.size(1) == dim,
	"inconsistent target size: ", target.sizes(), " for input of size: ",
	input.sizes());
	}
	}

	template <typename scalar_t, typename accscalar_t>
	C10_LAUNCH_BOUNDS_1(MULTILABELMARGIN_THREADS)
	__global__ void multilabel_margin_loss_forward_kernel(
	scalar_t* output,
	const scalar_t* input,
	const int64_t* target,
	scalar_t* is_target,
	int nframe,
	int dim,
	bool size_average) {

	// vectors:
	int k = blockIdx.x;
	const scalar_t* input_k = input + k * dim;
	const int64_t* target_k = target + k * dim;
	scalar_t* output_k = output + k;
	scalar_t* is_target_k = is_target + k * dim;

	// zero is_target
	for (int d = threadIdx.x; d < dim; d += blockDim.x) {
	is_target_k[d] = static_cast<scalar_t>(0);
	}
	__syncthreads();

	// mark targets in is_target
	if (threadIdx.x == 0) {
	for (int dt = 0; dt < dim; dt++) {
	int target_idx = target_k[dt];
	if (target_idx < 0) {
	break;
	}
	is_target_k[target_idx] = static_cast<scalar_t>(1);
	}
	}
	__syncthreads();

	// iterate over targets
	accscalar_t sum = 0;
	for (int dt = 0; dt < dim; dt++) {
	// next target:
	int target_idx = target_k[dt];
	if (target_idx < 0) {
	break;
	}

	// current value for target
	scalar_t input_target_k = input_k[target_idx];

	// compare to all inputs (multithreaded):
	for (int d = threadIdx.x; d < dim; d += blockDim.x) {
	// contribute to loss only if not a target
	if (!static_cast<int>(is_target_k[d])) {
	scalar_t z = 1 - input_target_k + input_k[d];
	if (z > 0) {
	sum += z;
	}
	}
	}
	}

	// Temporary sums (for mapreduce)
	__shared__ accscalar_t smem[MULTILABELMARGIN_THREADS];
	accscalar_t total_sum = cuda_utils::BlockReduceSum(sum, smem);
	if (threadIdx.x == 0) {
	if (size_average) {
	*output_k = static_cast<scalar_t>((total_sum / dim) / nframe);
	} else {
	*output_k = static_cast<scalar_t>(total_sum / dim);
	}
	}
	}

	template <typename scalar_t, typename accscalar_t>
	C10_LAUNCH_BOUNDS_1(MULTILABELMARGIN_THREADS)
	__global__ void multilabel_margin_loss_backward_kernel(
	scalar_t* grad_input,
	const scalar_t* grad_output,
	const scalar_t* input,
	const int64_t* target,
	const scalar_t* is_target,
	int nframe,
	int dim,
	bool size_average,
	bool reduce) {

	int k = blockIdx.x;
	const scalar_t* input_k = input + k * dim;
	scalar_t* grad_input_k = grad_input + k * dim;
	const int64_t* target_k = target + k * dim;
	const scalar_t* is_target_k = is_target + k * dim;

	const scalar_t* grad_output_k = grad_output;
	if (!reduce) {
	grad_output_k += k;
	}

	// gain:
	scalar_t g = static_cast<scalar_t>(
	size_average && reduce ? 1. / static_cast<accscalar_t>(nframe * dim)
	: 1. / static_cast<accscalar_t>(dim));

	// zero gradients:
	for (int d = threadIdx.x; d < dim; d += blockDim.x) {
	grad_input_k[d] = static_cast<scalar_t>(0);
	}
	__syncthreads();

	// iterate over targets
	for (int dt = 0; dt < dim; dt++) {
	// next target:
	int target_idx = static_cast<int>(target_k[dt]);
	if (target_idx < 0) {
	break;
	}

	// current value for target
	scalar_t input_target_k = input_k[target_idx];

	// compare to all inputs (multithreaded):
	accscalar_t sum = 0;
	for (int d = threadIdx.x; d < dim; d += blockDim.x) {
	// contribute to loss only if not a target
	if (!static_cast<int>(is_target_k[d])) {
	scalar_t z = 1 - input_target_k + input_k[d];
	if (z > 0) {
	sum -= g;
	grad_input_k[d] += g;
	}
	}
	}
	__syncthreads();

	// Temporary sums (for mapreduce)
	__shared__ accscalar_t smem[MULTILABELMARGIN_THREADS];
	accscalar_t total_sum = cuda_utils::BlockReduceSum(sum, smem);
	if (threadIdx.x == 0) {
	grad_input_k[target_idx] += static_cast<scalar_t>(total_sum);
	}
	}

	for (int d = threadIdx.x; d < dim; d += blockDim.x) {
	grad_input_k[d] = grad_output_k;
	}
	}

	void multilabel_margin_loss_forward_out_cuda_template(
	const Tensor& input,
	const Tensor& target,
	int64_t reduction,
	Tensor& output,
	Tensor& is_target) {
	int64_t nframe, dim;
	const int64_t ndims = input.dim();
	multilabel_margin_loss_shape_check(nframe, dim, ndims, input, target);

	if (input.numel() == 0) {
	return;
	}

	auto input_ = input.contiguous();
	auto target_ = target.contiguous();
	auto is_target_ = is_target.contiguous();
	is_target_.resize_as_(target);

	if (input.dim() <= 1) {
	output.resize_({});

	dim3 blocks(1);
	dim3 threads(MULTILABELMARGIN_THREADS);

	AT_DISPATCH_FLOATING_TYPES_AND2(
	at::ScalarType::Half,
	at::ScalarType::BFloat16,
	input.scalar_type(),
	"multilabel_margin_loss_forward_kernel",
	[&] {
	using accscalar_t = at::acc_type<scalar_t, true>;
	multilabel_margin_loss_forward_kernel<scalar_t, accscalar_t>
	<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
	output.mutable_data_ptr<scalar_t>(),
	input_.const_data_ptr<scalar_t>(),
	target_.const_data_ptr<int64_t>(),
	is_target_.mutable_data_ptr<scalar_t>(),
	1,
	dim,
	reduction == at::Reduction::Mean);
	C10_CUDA_KERNEL_LAUNCH_CHECK();
	});
	} else if (input.dim() == 2) {
	dim3 blocks(input.size(0));
	dim3 threads(MULTILABELMARGIN_THREADS);

	if (reduction != at::Reduction::None) {
	auto output_tmp = at::empty({input_.size(0)}, input_.options());
	output.resize_({});
	AT_DISPATCH_FLOATING_TYPES_AND2(
	at::ScalarType::Half,
	at::ScalarType::BFloat16,
	input.scalar_type(),
	"multilabel_margin_loss_forward_kernel",
	[&] {
	using accscalar_t = at::acc_type<scalar_t, true>;
	multilabel_margin_loss_forward_kernel<scalar_t, accscalar_t>
	<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
	output_tmp.mutable_data_ptr<scalar_t>(),
	input_.const_data_ptr<scalar_t>(),
	target_.const_data_ptr<int64_t>(),
	is_target_.mutable_data_ptr<scalar_t>(),
	nframe,
	dim,
	reduction == at::Reduction::Mean);
	C10_CUDA_KERNEL_LAUNCH_CHECK();
	});
	at::cuda::sum_out(
	output,
	output_tmp,
	at::IntArrayRef(std::vector<int64_t>{}),
	false,
	output.scalar_type());
	} else {
	output.resize_({input.size(0)});
	AT_DISPATCH_FLOATING_TYPES_AND2(
	at::ScalarType::Half,
	at::ScalarType::BFloat16,
	input.scalar_type(),
	"multilabel_margin_loss_forward_kernel",
	[&] {
	using accscalar_t = at::acc_type<scalar_t, true>;
	multilabel_margin_loss_forward_kernel<scalar_t, accscalar_t>
	<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
	output.mutable_data_ptr<scalar_t>(),
	input_.const_data_ptr<scalar_t>(),
	target_.const_data_ptr<int64_t>(),
	is_target_.mutable_data_ptr<scalar_t>(),
	nframe,
	dim,
	false);
	C10_CUDA_KERNEL_LAUNCH_CHECK();
	});
	}

	} else {
	TORCH_CHECK(
	false,
	"Expected 2D input with optional zero batch dim, or 1D input with non-zero dims, but got sizes: ",
	input.sizes());
	}
	}

	void multilabel_margin_loss_backward_cuda_out_template(
	const Tensor& grad_output,
	const Tensor& input,
	const Tensor& target,
	int64_t reduction,
	const Tensor& is_target,
	Tensor& grad_input) {
	int64_t nframe, dim;
	const int64_t ndims = input.dim();
	multilabel_margin_loss_shape_check(nframe, dim, ndims, input, target);

	if (input.numel() == 0) {
	return;
	}

	auto input_ = input.contiguous();
	auto target_ = target.contiguous();
	auto is_target_ = is_target.contiguous();
	auto grad_output_ = grad_output.contiguous();
	grad_input.resize_as_(input_);

	if (grad_input.dim() <= 1) {
	int target_size = target_.dim() == 0 ? 1 : target_.size(0);
	TORCH_CHECK(
	(target_.numel() != 0) && (target_.dim() <= 1) && (target_size == dim),
	"inconsistent target size");
	TORCH_CHECK(
	target_.sizes() == is_target_.sizes(), "inconsistent is_target size");
	dim3 blocks(1);
	dim3 threads(MULTILABELMARGIN_THREADS);

	AT_DISPATCH_FLOATING_TYPES_AND2(
	at::ScalarType::Half,
	at::ScalarType::BFloat16,
	input.scalar_type(),
	"multilabel_margin_loss_backward_kernel",
	[&] {
	using accscalar_t = at::acc_type<scalar_t, true>;
	multilabel_margin_loss_backward_kernel<scalar_t, accscalar_t>
	<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
	grad_input.mutable_data_ptr<scalar_t>(),
	grad_output_.const_data_ptr<scalar_t>(),
	input_.const_data_ptr<scalar_t>(),
	target_.const_data_ptr<int64_t>(),
	is_target_.const_data_ptr<scalar_t>(),
	1,
	dim,
	reduction == at::Reduction::Mean,
	reduction != at::Reduction::None);
	C10_CUDA_KERNEL_LAUNCH_CHECK();
	});
	} else if (grad_input.dim() == 2) {
	TORCH_CHECK(
	(input_.size(1) != 0) && (target_.dim() == 2) &&
	(target_.size(0) == nframe) && (target_.size(1) == dim),
	"inconsistent target size");
	TORCH_CHECK(target_.sizes() == is_target_.sizes(), "inconsistent is_target size");
	dim3 blocks(grad_input.size(0));
	dim3 threads(MULTILABELMARGIN_THREADS);

	AT_DISPATCH_FLOATING_TYPES_AND2(
	at::ScalarType::Half,
	at::ScalarType::BFloat16,
	input.scalar_type(),
	"multilabel_margin_loss_backward_kernel",
	[&] {
	using accscalar_t = at::acc_type<scalar_t, true>;
	multilabel_margin_loss_backward_kernel<scalar_t, accscalar_t>
	<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
	grad_input.mutable_data_ptr<scalar_t>(),
	grad_output_.const_data_ptr<scalar_t>(),
	input_.const_data_ptr<scalar_t>(),
	target_.const_data_ptr<int64_t>(),
	is_target_.const_data_ptr<scalar_t>(),
	grad_input.size(0),
	grad_input.size(1),
	reduction == at::Reduction::Mean,
	reduction != at::Reduction::None);
	C10_CUDA_KERNEL_LAUNCH_CHECK();
	});
	} else {
	TORCH_CHECK(
	false,
	"Expected 2D input with optional zero batch dim, or 1D input with non-zero dims, but got sizes: ",
	grad_input.sizes());
	}
	}

	} // namespace

	std::tuple<Tensor&, Tensor&> multilabel_margin_loss_forward_out_cuda(
	const Tensor& self,
	const Tensor& target,
	int64_t reduction,
	Tensor& output,
	Tensor& is_target) {
	multilabel_margin_loss_forward_out_cuda_template(
	self, target, reduction, output, is_target);
	return std::tuple<Tensor&, Tensor&>(output, is_target);
	}

	std::tuple<Tensor, Tensor> multilabel_margin_loss_forward_cuda(
	const Tensor& self,
	const Tensor& target,
	int64_t reduction) {
	auto output = at::empty({0}, self.options());
	auto is_target = at::empty({0}, self.options());
	multilabel_margin_loss_forward_out_cuda_template(
	self, target, reduction, output, is_target);
	return std::make_tuple(output, is_target);
	}

	Tensor& multilabel_margin_loss_backward_cuda_out(
	const Tensor& grad_output,
	const Tensor& self,
	const Tensor& target,
	int64_t reduction,
	const Tensor& is_target,
	Tensor& grad_input) {
	multilabel_margin_loss_backward_cuda_out_template(
	grad_output, self, target, reduction, is_target, grad_input);
	return grad_input;
	}

	Tensor multilabel_margin_loss_backward_cuda(
	const Tensor& grad_output,
	const Tensor& self,
	const Tensor& target,
	int64_t reduction,
	const Tensor& is_target) {
	auto grad_input = at::zeros_like(self, LEGACY_CONTIGUOUS_MEMORY_FORMAT);
	multilabel_margin_loss_backward_cuda_out_template(
	grad_output, self, target, reduction, is_target, grad_input);
	return grad_input;
	}

	} // namespace at::native