def bias_relu()

in blocksparse/ewops.py [0:0]

• 16 lines of code
• 11 McCabe index (conditional complexity)

def bias_relu(x, b, axis=-1, relu=False, fast_gelu=False, atomics=True, bench=0, use_tf=False):

    if relu and fast_gelu:
        raise ValueError("relu and fast_gelu can not both be enabled.")

    dev = x.op.device.lower()
    if use_tf or not dev or "cpu" in dev:
        if b.shape.ndims > 1:
            y = x + b
        else:
            y = tf.nn.bias_add(x, b)
        if relu:
            y = tf.nn.relu(y)
        elif fast_gelu:
            y = y * tf.nn.sigmoid(1.702 * y)

        return y

    relu = 1 if relu else (2 if fast_gelu else 0)

    return bias_relu_op(x, b, axis=axis, relu=relu, bench=bench, atomics=atomics)