We learned that the squash function converts the length of the vector into probability, and it is given as follows:

The squash function can be defined as follows:

def squash(sj):
    
    sj_norm = tf.reduce_sum(tf.square(sj), -2, keep_dims=True)
    scalar_factor = sj_norm / (1 + sj_norm) / tf.sqrt(sj_norm + epsilon)

    vj = scalar_factor * sj 

    return vj