Does flax run modules implicitly in vmap? neither jax.lax.cond or flax.nn.cond works as expected!

[dhyani15]

I want to evaluate either function 1 or function 2 but not both, one is cheap and the other costly. But when I jit and run this flax vit implementation I get constant time of execution even when I run mostly cheap operations vs mostly expensive operations.

Here is the snippet of code, can you tell me if I am doing something incorrectly.

The ChooseAttention class is where the selection happening between the two functions. This is a modified implementation of flax vit of hugging face.

class ChooseAttention(nn.Module):
    config: ViTConfig
    dtype: jnp.dtype = jnp.float32
    def setup(self):
        self.alpha = self.param('alpha', nn.initializers.ones, (1, self.config.num_attention_heads, 1, 1))
    def choose_attention(self, alpha, x):
        return jax.lax.cond(alpha[0, 0, 0, 0] < 0.5, lambda x: x / x.shape[-1],lambda x: nn.relu(x)/(jnp.sum(nn.relu(x), axis=-1, keepdims=True) + 1e-5),x)
    def __call__(self, attn_weights):
        results = []
        for i in range(self.alpha.shape[1]):
            alpha_i = self.alpha[:, i:i+1, :, :]
            x_i = attn_weights[:, i:i+1, :, :]
            result_i = self.choose_attention(alpha_i, x_i)
            results.append(result_i)

        return jnp.concatenate(results, axis=1)

class CustomFlaxViTSelfAttention(FlaxViTSelfAttention):
    config: ViTConfig
    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
    
    def setup(self):
        # self.alpha = self.param('alpha', nn.initializers.ones, (1, self.config.num_attention_heads, 1, 1))
        self.chooseAttention = ChooseAttention(self.config)
        if self.config.hidden_size % self.config.num_attention_heads != 0:
            raise ValueError(
                "`config.hidden_size`: {self.config.hidden_size} has to be a multiple of `config.num_attention_heads`:"
                " {self.config.num_attention_heads}"
            )
        # self.qkv = nn.Dense(self.config.hidden_size * 3,dtype=self.dtype,kernel_init=jax.nn.initializers.glorot_uniform(),use_bias=False)
        self.query = nn.Dense(
            self.config.hidden_size,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.variance_scaling(
                self.config.initializer_range**2, mode="fan_in", distribution="truncated_normal"
            ),
            use_bias=self.config.qkv_bias,
        )
        self.key = nn.Dense(
            self.config.hidden_size,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.variance_scaling(
                self.config.initializer_range**2, mode="fan_in", distribution="truncated_normal"
            ),
            use_bias=self.config.qkv_bias,
        )
        self.value = nn.Dense(
            self.config.hidden_size,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.variance_scaling(
                self.config.initializer_range**2, mode="fan_in", distribution="truncated_normal"
            ),
            use_bias=self.config.qkv_bias,
        )
        
    
    def __call__(self, hidden_states, deterministic: bool = True, output_attentions: bool = False):
        head_dim = self.config.hidden_size // self.config.num_attention_heads # 1,257,192
        query_states = self.query(hidden_states).reshape( # query dense 192 -> 192 (12*16)
            hidden_states.shape[:2] + (self.config.num_attention_heads, head_dim)
        )
        
        value_states = self.value(hidden_states).reshape(
            hidden_states.shape[:2] + (self.config.num_attention_heads, head_dim)
        )
        
        key_states = self.key(hidden_states).reshape(
            hidden_states.shape[:2] + (self.config.num_attention_heads, head_dim)
        )
        

        dropout_rng = None
        if not deterministic and self.config.attention_probs_dropout_prob > 0.0:
            dropout_rng = self.make_rng("dropout")


        query, key = promote_dtype(query_states, key_states, dtype=self.dtype)
        dtype = query.dtype

        assert query.ndim == key.ndim, 'q, k must have same rank.'
        assert query.shape[:-3] == key.shape[:-3], 'q, k batch dims must match.'
        assert query.shape[-2] == key.shape[-2], 'q, k num_heads must match.'
        assert query.shape[-1] == key.shape[-1], 'q, k depths must match.'

        # calculate attention matrix
        depth = query.shape[-1]
        query = query / jnp.sqrt(depth).astype(dtype)
       
        attn_weights = jnp.einsum('...qhd,...khd->...hqk', query, key, precision='highest').astype(dtype)
        # vectorized_choose_attention = jax.vmap(choose_attention, in_axes=(1, 1), out_axes=1)
        # scalattn = (attn_weights/attn_weights.shape[2]).astype(dtype) # scale attention
        # attn_weights = nn.relu(attn_weights).astype(dtype) 
        # attn_weights = attn_weights / (jnp.sum(attn_weights, axis=-1, keepdims=True) + 1e-5).astype(dtype) # RElusoftmax
        # attn_weights = (self.alpha*attn_weights).astype(dtype) + ((1-self.alpha)*scalattn).astype(dtype) # weighted sum
        # attn_weights = vectorized_choose_attention(self.alpha,attn_weights)
        attn_weights = self.chooseAttention(attn_weights)


        if not deterministic and self.config.attention_probs_dropout_prob > 0.0:
            keep_prob = 1.0 - self.config.attention_probs_dropout_prob
            if True:
                # dropout is broadcast across the batch + head dimensions
                dropout_shape = tuple([1] * (key.ndim - 2)) + attn_weights.shape[-2:]
                keep = random.bernoulli(dropout_rng, keep_prob, dropout_shape)  # type: ignore
            else:
                keep = random.bernoulli(dropout_rng, keep_prob, attn_weights.shape)  # type: ignore
            multiplier = keep.astype(dtype) / jnp.asarray(keep_prob, dtype=dtype)
            attn_weights = attn_weights * multiplier

        attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value_states)
        attn_output = attn_output.reshape(attn_output.shape[:2] + (-1,))

        outputs = (attn_output, attn_weights) if output_attentions else (attn_output,)
        return outputs

[cgarciae]

Hey @dhyani15, the correct way to do this would be to expose a static argument in the jit-ted function to compile a different function for each mode of operation. Not sure how to do this here as you are only showing the model code.