#!/usr/bin/env python3
"""First-order derivatives and second-order derivatives for gelu function."""
import tensorflow
from tensorflow.python.framework import (
    ops,
)

from deepmd.env import (
    op_module,
)

try:
    gelu = tensorflow.nn.gelu
except AttributeError:

    @ops.RegisterGradient("Gelu")
    def _gelu_cc(op, dy):
        return op_module.gelu_grad_custom(dy, op.inputs[0])

    @ops.RegisterGradient("GeluGrad")
    def _gelu_grad_cc(op, dy):
        return [
            op_module.gelu_grad_custom(dy, op.inputs[1]),
            op_module.gelu_grad_grad_custom(dy, op.inputs[0], op.inputs[1]),
        ]


@ops.RegisterGradient("GeluCustom")
def _gelu_custom_cc(op, dy):
    return op_module.gelu_grad_custom(dy, op.inputs[0])


@ops.RegisterGradient("GeluGradCustom")
def _gelu_grad_custom_cc(op, dy):
    return [
        op_module.gelu_grad_custom(dy, op.inputs[1]),
        op_module.gelu_grad_grad_custom(dy, op.inputs[0], op.inputs[1]),
    ]
