用Keras写一个non local block，用def方法写，不要class

好的，下面是一个用 Keras 实现的非局部块（Non-local Block）的例子，使用 def 方法写：

from keras import backend as K
from keras.layers import Conv2D, Activation, Lambda, add, multiply

def non_local_block(ip, computation='dot', compression=2, mode='embedded'):
    '''
    Non-local block implementation in Keras.
    
    Args:
    - ip: Input tensor to the non-local block.
    - computation: Type of computation to be used for similarity calculation. 
                   Valid options: 'dot' or 'embedded'.
                   Default: 'dot'.
    - compression: Factor by which to reduce the number of channels before computing 
                   similarities. 
                   Default: 2.
    - mode: Whether to use embedded Gaussian function ('embedded') or standard Gaussian 
            function ('gaussian') for computing weights.
            Default: 'embedded'.
    
    Returns:
    - Output tensor of the non-local block.
    '''
    
    # Compute the input shape of the tensor
    input_shape = K.int_shape(ip)
    
    # Compute the dimensions of the space in which the non-local operation is performed
    spatial_dims = len(input_shape) - 2
    
    # Compute the number of channels after compression
    reduced_channels = max(input_shape[spatial_dims + 1] // compression, 1)
    
    # Define the theta path
    theta = Conv2D(reduced_channels, kernel_size=1, strides=1, padding='same')(ip)
    
    # Define the phi path
    phi = Conv2D(reduced_channels, kernel_size=1, strides=1, padding='same')(ip)
    
    # Define the g path
    g = Conv2D(reduced_channels, kernel_size=1, strides=1, padding='same')(ip)
    
    if computation == 'dot':
        # Perform dot product between theta and phi
        theta_phi = Lambda(lambda x: K.batch_dot(x[0], x[1], axes=[3, 3]))([theta, phi])
    elif computation == 'embedded':
        # Compute the similarity between theta and phi using embedded Gaussian function
        theta_phi = Conv2D(reduced_channels, kernel_size=1, strides=1, padding='same')(theta)
        theta_phi = Lambda(lambda x: K.batch_dot(x[0], x[1], axes=[3, 3]))([theta_phi, phi])
        theta_phi = Activation('softmax')(theta_phi)
    else:
        raise ValueError("Invalid value for 'computation' argument. Valid options are 'dot' or 'embedded'.")
    
    # Define the output path
    if mode == 'embedded':
        # Compute the weight using embedded Gaussian function
        theta_phi_g = Conv2D(input_shape[spatial_dims + 1], kernel_size=1, strides=1, padding='same')(g)
        theta_phi_g = Lambda(lambda x: K.batch_dot(x[0], x[1], axes=[3, 2]))([theta_phi, theta_phi_g])
        theta_phi_g = Activation('softmax')(theta_phi_g)
    elif mode == 'gaussian':
        # Compute the weight using standard Gaussian function
        theta_phi_g = Lambda(lambda x: K.softmax(x / K.sqrt(K.cast(reduced_channels, dtype=K.floatx()))))(theta_phi)
    else:
        raise ValueError("Invalid value for 'mode' argument. Valid options are 'embedded' or 'gaussian'.")
    
    # Apply the weight to the output of g path
    g_theta_phi = Conv2D(input_shape[spatial_dims + 1], kernel_size=1, strides=1, padding='same')(g)
    g_theta_phi = multiply([g_theta_phi, theta_phi_g])
    
    # Add the output of g path to the input tensor to complete the non-local block
    y = add([ip, g_theta_phi])
    
    return y

这里实现了一个非局部块，包括了 dot-based 和 embedded Gaussian-based 两种计算相似度的方式，以及使用 embedded Gaussian function 或 standard Gaussian function 计算权重的两种方式。可以根据需要进行修改。