File size: 666 Bytes
c1af2fa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
#pragma once

namespace at {
struct CUDAGeneratorImpl;
struct TensorIteratorBase;
class TensorBase;

namespace native {

void launch_poisson_cuda_kernel(

    const TensorBase &ret, const TensorBase &lambda, CUDAGeneratorImpl *gen);

void launch_gamma_kernel(

    const TensorBase &ret, const TensorBase &alpha, CUDAGeneratorImpl *gen);

void launch_binomial_cuda_kernel(

    TensorIteratorBase &iter, CUDAGeneratorImpl *gen);

void launch_dirichlet_kernel(TensorIteratorBase &iter);

void launch_standard_gamma_grad_kernel(TensorIteratorBase &iter);

void launch_dirichlet_grad_kernel(TensorIteratorBase &iter);

}}  // namespace at::native