80dd4096ed186e706775adbe3663b46d719efdb0,capsulelayers.py,CapsuleLayer,call,#CapsuleLayer#,94

Before Change

            return [i-1, b, outputs]

        cond = lambda i, b, inputs_hat: i > 0
        loop_vars = [K.constant(self.num_routing), self.bias, K.sum(inputs_hat, 1, keepdims=True)]
        _, self.bias, outputs = tf.while_loop(cond, body, loop_vars)

        
        // Routing algorithm V2. Seems not right. This may duplicate tensors by self.num_routing times.
        for _ in range(self.num_routing):

After Change

        _, self.bias, outputs = tf.while_loop(cond, body, loop_vars)

        // Routing algorithm V2. Use for iteration. V2 and V1 both work without much difference on performance
        for _ in range(self.num_routing):
            c = K.softmax(self.bias)
            c_expand = K.expand_dims(K.expand_dims(K.expand_dims(c, 2), 2), 0)
            outputs = K.sum(c_expand * inputs_hat, 1, keepdims=True)
            outputs = squash(outputs)
            self.bias = K.update(self.bias, self.bias + K.sum(inputs_hat * outputs, [0, -2, -1]))

        // Handling with no routing scenario. Prior bias will always be zero.
        if self.num_routing == 0:
            c = K.softmax(self.bias)
            c_expand = K.expand_dims(K.expand_dims(K.expand_dims(c, 2), 2), 0)
            outputs = squash(K.sum(c_expand * inputs_hat, 1, keepdims=True))

Instances