code: 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)
  
file_path: insightface/reconstruction/ostec/external/stylegan2/dnnlib/tflib/ops/fused_bias_act.py
imports: import os
  import numpy
  import tensorflow
  from ..None import custom_ops
  from ...util import EasyDict
label_desc: Use image node_id 4 for calling a global function with example usage: _fused_bias_act_cuda(x, b, axis, act, alpha, gain) and returns: func_nonzero_2nd_grad, x, _fused_bias_act_ref, y, dx, db, d_dy, d_x, y,  grad, y,  grad_wrap, func_zero_2nd_grad, empty_tensor, dx,  db,  grad2, grad_impl, d_dy, dx,  db,  grad2, d_dy,  d_x