caffe2/operators/softmax_shared.h


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

#ifndef CAFFE2_OPERATORS_SOFTMAX_SHARED_H_
#define CAFFE2_OPERATORS_SOFTMAX_SHARED_H_

#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"

namespace caffe2 {

void SoftmaxCPU(
    CPUContext& context,
    const int N,
    const int D,
    const Tensor<CPUContext>& X,
    float* Ydata,
    Tensor<CPUContext>& scale,
    Tensor<CPUContext>& sum_multiplier);
} // namespace caffe2

#endif // #define CAFFE2_OPERATORS_SOFTMAX_SHARED_H_