diff --git a/paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc index 9c5d03c17afbadfda1a7d2fd618e35205f6c10ad..44edc224795706621050fb5b2c22037db86fc9bd 100644 --- a/paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc +++ b/paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc @@ -282,7 +282,6 @@ namespace ops = paddle::operators; __macro(swish, SwishMKLDNNFunctor, SwishMKLDNNGradFunctor); \ __macro(hard_swish, HardSwishMKLDNNFunctor, HardSwishMKLDNNGradFunctor); \ __macro(tanh, TanhMKLDNNFunctor, TanhMKLDNNGradFunctor); \ - __macro(sqrt, SqrtMKLDNNFunctor, SqrtMKLDNNGradFunctor); \ __macro(abs, AbsMKLDNNFunctor, AbsMKLDNNGradFunctor); \ __macro(elu, EluMKLDNNFunctor, EluMKLDNNGradFunctor); @@ -293,6 +292,8 @@ REGISTER_ACTIVATION_MKLDNN_BF16_KERNEL(gelu, GeluMKLDNNFunctor, GeluMKLDNNGradFunctor); REGISTER_ACTIVATION_MKLDNN_BF16_KERNEL(sigmoid, SigmoidMKLDNNFunctor, SigmoidMKLDNNGradFunctor); +REGISTER_ACTIVATION_MKLDNN_BF16_KERNEL(sqrt, SqrtMKLDNNFunctor, + SqrtMKLDNNGradFunctor); namespace ops = paddle::operators; REGISTER_OP_KERNEL( diff --git a/python/paddle/fluid/tests/unittests/mkldnn/test_activation_bf16_mkldnn_op.py b/python/paddle/fluid/tests/unittests/mkldnn/test_activation_bf16_mkldnn_op.py index cd9987b3c8e8245c9e8fc73a8599e5d9ae1241ac..c421a6d117e458689aa019dd8671cf51c7380606 100644 --- a/python/paddle/fluid/tests/unittests/mkldnn/test_activation_bf16_mkldnn_op.py +++ b/python/paddle/fluid/tests/unittests/mkldnn/test_activation_bf16_mkldnn_op.py @@ -82,6 +82,20 @@ class TestMKLDNNSigmoidBF16Op(MKLDNNBF16ActivationOp, TestActivation): return dout * self.op_forward(x) * (1 - self.op_forward(x)) +class TestMKLDNNSqrtBF16Op(MKLDNNBF16ActivationOp, TestActivation): + def config(self): + self.op_type = "sqrt" + + def init_data(self): + self.x = np.random.uniform(1, 2, [2, 4, 3, 5]).astype(np.float32) + + def op_forward(self, x): + return np.sqrt(x) + + def op_grad(self, dout, x): + return dout / (2 * np.sqrt(x)) + + class TestMKLDNNGeluErfBF16Op(MKLDNNBF16ActivationOp, TestActivation): def config(self): self.op_type = "gelu"