onnxruntime/onnxruntime/contrib_ops/cuda/layer_norm.h
Yufeng Li a1f8fe1463
Add layernorm operator (#1967)
* Add LayerNormalization operator
It's a fusion of bert model's layer normalization.
2019-10-03 11:32:13 -07:00

29 lines
648 B
C++

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "core/providers/cuda/cuda_common.h"
namespace onnxruntime {
namespace contrib {
namespace cuda {
using namespace onnxruntime::cuda;
template <typename T, typename U>
class LayerNorm final : public CudaKernel {
public:
LayerNorm(const OpKernelInfo& op_kernel_info);
Status ComputeInternal(OpKernelContext* ctx) const override;
private:
int64_t axis_;
double epsilon_;
};
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime