feat(nnx): add preferred_element_type to attention.py and recurrent.py, out_sharding to recurrent.py#

Author: aarushisingh04

flax/nnx/nn/attention.py

@@ -63,6 +63,7 @@ def dot_product_attention_weights(
   module: Module | None = None,
   promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,
   is_causal: bool = False,
+  preferred_element_type: Dtype | None = None,
 ):
   """Computes dot-product attention weights given query and key.
 
@@ -100,6 +101,9 @@ def dot_product_attention_weights(
       the logits to mask out the non-causal parts of the attention matrix,
       but other implementations like cudnn will avoid computing the
       non-causal regions, providing speedups.
+    preferred_element_type: Optional parameter controls the data type output by
+      the dot product. This argument is passed to ``dot_general`` function.
+      See ``jax.lax.dot`` for details.
 
   Returns:
     Output of shape `[batch..., num_heads, q_length, kv_length]`.
@@ -117,7 +121,11 @@ def dot_product_attention_weights(
   query = query / jnp.sqrt(depth).astype(dtype)
   # attn weight shape is (batch..., num_heads, q_length, kv_length)
   attn_weights = jnp.einsum(
-    '...qhd,...khd->...hqk', query, key, precision=precision
+    '...qhd,...khd->...hqk',
+    query,
+    key,
+    precision=precision,
+    preferred_element_type=preferred_element_type,
   )
 
   # apply attention bias: masking, dropout, proximity bias, etc.
@@ -172,6 +180,7 @@ def dot_product_attention(
   module: Module | None = None,
   promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,
   is_causal: bool = False,
+  preferred_element_type: Dtype | None = None,
 ):
   """Computes dot-product attention given query, key, and value.
 
@@ -218,6 +227,9 @@ def dot_product_attention(
       the logits to mask out the non-causal parts of the attention matrix,
       but other implementations like cudnn will avoid computing the
       non-causal regions, providing speedups.
+    preferred_element_type: Optional parameter controls the data type output by
+      the dot product. This argument is passed to ``dot_general`` function.
+      See ``jax.lax.dot`` for details.
 
   Returns:
     Output of shape `[batch..., q_length, num_heads, v_depth_per_head]`.
@@ -244,10 +256,17 @@ def reshape_4d(x):
         reshape_4d, (query, key, value, bias, mask))
     if mask is not None:
       mask = mask.astype(jnp.bool)
-    out = jax.nn.dot_product_attention(query, key, value, bias, mask, is_causal=is_causal)
+    out = jax.nn.dot_product_attention(
+      query,
+      key,
+      value,
+      bias,
+      mask,
+      is_causal=is_causal,
+    )
     if len(query_shape) > 4:
       out = jnp.reshape(out, query_shape)
-    return out
+    return out.astype(preferred_element_type)
 
   # compute attention weights
   attn_weights = dot_product_attention_weights(
@@ -264,11 +283,16 @@ def reshape_4d(x):
     module,
     promote_dtype,
     is_causal,
+    preferred_element_type=preferred_element_type,
   )
 
   # return weighted sum over values for each query position
   return jnp.einsum(
-    '...hqk,...khd->...qhd', attn_weights, value, precision=precision
+    '...hqk,...khd->...qhd',
+    attn_weights,
+    value,
+    precision=precision,
+    preferred_element_type=preferred_element_type,
   )
 
 
@@ -351,6 +375,8 @@ class MultiHeadAttention(Module):
       the scale of the query layer norm layer.
     key_ln_scale_metadata: Optional metadata dictionary to set when initializing
       the scale of the key layer norm layer.
+    preferred_element_type: numerical precision of the computation, see
+      `jax.lax.dot_general` for details.
   """
 
   def __init__(
@@ -378,6 +404,7 @@ def __init__(
     qkv_promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,
     out_promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,
     ln_promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,
+    preferred_element_type: Dtype | None = None,
     # Deprecated, will be removed.
     qkv_dot_general: DotGeneralT | None = None,
     out_dot_general: DotGeneralT | None = None,
@@ -412,6 +439,7 @@ def __init__(
     self.use_bias = use_bias
     self.attention_fn = attention_fn
     self.decode = decode
+    self.preferred_element_type = preferred_element_type
     self.normalize_qk = normalize_qk
     self.qkv_promote_dtype = qkv_promote_dtype
     self.out_promote_dtype = out_promote_dtype
@@ -441,6 +469,7 @@ def __init__(
       promote_dtype=self.qkv_promote_dtype,
       dot_general=self.qkv_dot_general,
       dot_general_cls=self.qkv_dot_general_cls,
+      preferred_element_type=self.preferred_element_type,
       kernel_metadata=kernel_metadata,
       bias_metadata=bias_metadata,
     )
@@ -490,6 +519,7 @@ def __init__(
       promote_dtype=self.out_promote_dtype,
       dot_general=self.out_dot_general,
       dot_general_cls=self.out_dot_general_cls,
+      preferred_element_type=self.preferred_element_type,
       rngs=rngs,
       kernel_metadata=out_kernel_metadata or kernel_metadata,
       bias_metadata=out_bias_metadata or bias_metadata,
@@ -665,6 +695,7 @@ def __call__(
       deterministic=deterministic,
       dtype=self.dtype,
       precision=self.precision,
+      preferred_element_type=self.preferred_element_type,
       module=self if sow_weights else None,
     )
     # back to the original inputs dimensions