def _fused_bias_act_cuda(x, b, axis, act, alpha, gain):
    """Fast CUDA implementation of `fused_bias_act()` using custom ops."""

    # Validate arguments.
    x = tf.convert_to_tensor(x)
    empty_tensor = tf.constant([], dtype=x.dtype)
    b = tf.convert_to_tensor(b) if b is not None else empty_tensor
    act_spec = activation_funcs[act]
    assert b.shape.rank == 1 and (b.shape[0] == 0 or b.shape[0] == x.shape[axis])
    assert b.shape[0] == 0 or 0 <= axis < x.shape.rank
    if alpha is None:
        alpha = act_spec.def_alpha
    if gain is None:
        gain = act_spec.def_gain

    # Special cases.
    if act == 'linear' and b is None and gain == 1.0:
        return x
    if act_spec.cuda_idx is None:
        return _fused_bias_act_ref(x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain)

    # CUDA kernel.
    cuda_kernel = _get_plugin().fused_bias_act
    cuda_kwargs = dict(axis=axis, act=act_spec.cuda_idx, alpha=alpha, gain=gain)

    # Forward pass: y = func(x, b).
    def func_y(x, b):
        y = cuda_kernel(x=x, b=b, ref=empty_tensor, grad=0, **cuda_kwargs)
        y.set_shape(x.shape)
        return y

    # Backward pass: dx, db = grad(dy, x, y)
    def grad_dx(dy, x, y):
        ref = {'x': x, 'y': y}[act_spec.ref]
        dx = cuda_kernel(x=dy, b=empty_tensor, ref=ref, grad=1, **cuda_kwargs)
        dx.set_shape(x.shape)
        return dx
    def grad_db(dx):
        if b.shape[0] == 0:
            return empty_tensor
        db = dx
        if axis < x.shape.rank - 1:
            db = tf.reduce_sum(db, list(range(axis + 1, x.shape.rank)))
        if axis > 0:
            db = tf.reduce_sum(db, list(range(axis)))
        db.set_shape(b.shape)
        return db

    # Second order gradients: d_dy, d_x = grad2(d_dx, d_db, x, y)
    def grad2_d_dy(d_dx, d_db, x, y):
        ref = {'x': x, 'y': y}[act_spec.ref]
        d_dy = cuda_kernel(x=d_dx, b=d_db, ref=ref, grad=1, **cuda_kwargs)
        d_dy.set_shape(x.shape)
        return d_dy
    def grad2_d_x(d_dx, d_db, x, y):
        ref = {'x': x, 'y': y}[act_spec.ref]
        d_x = cuda_kernel(x=d_dx, b=d_db, ref=ref, grad=2, **cuda_kwargs)
        d_x.set_shape(x.shape)
        return d_x

    # Fast version for piecewise-linear activation funcs.
    @tf.custom_gradient
    def func_zero_2nd_grad(x, b):
        y = func_y(x, b)
        @tf.custom_gradient
        def grad(dy):
            dx = grad_dx(dy, x, y)
            db = grad_db(dx)
            def grad2(d_dx, d_db):
                d_dy = grad2_d_dy(d_dx, d_db, x, y)
                return d_dy
            return (dx, db), grad2
        return y, grad

    # Slow version for general activation funcs.
    @tf.custom_gradient
    def func_nonzero_2nd_grad(x, b):
        y = func_y(x, b)
        def grad_wrap(dy):
            @tf.custom_gradient
            def grad_impl(dy, x):
                dx = grad_dx(dy, x, y)
                db = grad_db(dx)
                def grad2(d_dx, d_db):
                    d_dy = grad2_d_dy(d_dx, d_db, x, y)
                    d_x = grad2_d_x(d_dx, d_db, x, y)
                    return d_dy, d_x
                return (dx, db), grad2
            return grad_impl(dy, x)
        return y, grad_wrap

    # Which version to use?
    if act_spec.zero_2nd_grad:
        return func_zero_2nd_grad(x, b)
    return func_nonzero_2nd_grad(x, b)
Fast CUDA implementation of `fused_bias_act()` using custom ops.
converters/stylegan2_official/dnnlib/tflib/ops/fused_bias_act.py