Convert DenseGeneral to NNX

# Description
This commit converts `DenseGeneral` to NNX and creates a `dense_general` to interface with it through a Linen wrapper. `dense_general` contains all the same arguments as the Linen version but adds two additional ones:

* `input_shape`: the expected shape of the input.
* `in_features`: an int or tuple representing the input features.

Only one of them can be set at a time.

# Tests

# Checklist

Before submitting this PR, please make sure (put X in square brackets):
- [x] I have performed a self-review of my code.
- [x] I have necessary comments in my code, particularly in hard-to-understand areas.
- [x] I have run end-to-end tests tests and provided workload links above if applicable.
- [x] I have made or will make corresponding changes to the doc if needed.

PiperOrigin-RevId: 748311465

Author: Cristian Garcia

MaxText/layers/attentions.py

@@ -60,6 +60,7 @@ class AttentionType(enum.Enum):
 Mesh = common_types.Mesh
 PRNGKey = common_types.PRNGKey
 DenseGeneral = linears.DenseGeneral
+dense_general = linears.dense_general
 RMSNorm = linears.RMSNorm
 RotaryEmbedding = embeddings.RotaryEmbedding
 YarnRotaryEmbedding = embeddings.YarnRotaryEmbedding
@@ -1126,7 +1127,8 @@ def query_init(*args):
     kernel_axes = (
         (None, None, None) if self.config.ici_context_autoregressive_parallelism > 1 else ("embed", "q_heads", "kv")
     )
-    query_proj = DenseGeneral(
+    query_proj = dense_general(
+        inputs_shape=inputs_q.shape,
         features=(self.num_query_heads, self.head_dim),
         axis=-1,
         kernel_init=query_init,
@@ -1162,7 +1164,8 @@ def kv_projection(self, inputs_kv: Array, proj_name: str) -> Array:
         else ("embed", "kv_heads", "kv_head_dim")
     )
 
-    kv_proj = DenseGeneral(
+    kv_proj = dense_general(
+        inputs_shape=inputs_kv.shape,
         features=(self.num_kv_heads, self.head_dim),
         axis=-1,
         kernel_init=self.kernel_init,
@@ -1178,7 +1181,8 @@ def kv_projection(self, inputs_kv: Array, proj_name: str) -> Array:
   def qkv_projection(self, inputs: Array, proj_name: str):
     """Fused QKV projection"""
 
-    qkv_proj = DenseGeneral(
+    qkv_proj = dense_general(
+        inputs_shape=inputs.shape,
         features=(3, self.num_query_heads, self.head_dim),
         axis=-1,
         kernel_init=self.kernel_init,
@@ -1197,7 +1201,8 @@ def out_projection(self, output_dim: int, out: Array) -> Array:
     out_kernel_axis = (
         (None, None, None) if self.config.ici_context_autoregressive_parallelism > 1 else ("heads", "kv", "embed")
     )
-    out_proj = DenseGeneral(
+    out_proj = dense_general(
+        inputs_shape=out.shape,
         features=output_dim,
         axis=(-2, -1),
         kernel_init=self.kernel_init,
@@ -1455,7 +1460,8 @@ def setup(self):
 
     if self.q_lora_rank == 0:
       # Standard Q projection (without LoRA).
-      self.query_proj = DenseGeneral(
+      self.query_proj = dense_general(
+          in_features=self.config.emb_dim,
           features=(self.num_query_heads, self.qk_head_dim),
           axis=-1,
           kernel_init=self.kernel_init,
@@ -1468,7 +1474,8 @@ def setup(self):
       )
     else:
       # LoRA path for Q.
-      self.wq_a = DenseGeneral(
+      self.wq_a = dense_general(
+          in_features=self.config.emb_dim,
           features=self.q_lora_rank,
           axis=-1,
           kernel_init=self.kernel_init,
@@ -1486,7 +1493,8 @@ def setup(self):
           epsilon=self.config.normalization_layer_epsilon,
           kernel_axes=("norm",),
       )
-      self.wq_b = DenseGeneral(
+      self.wq_b = dense_general(
+          in_features=self.q_lora_rank,
           features=(self.num_query_heads, self.qk_head_dim),
           axis=-1,
           kernel_init=self.kernel_init,
@@ -1499,7 +1507,8 @@ def setup(self):
       )
 
     # KV LoRA path.
-    self.wkv_a = DenseGeneral(
+    self.wkv_a = dense_general(
+        in_features=self.config.emb_dim,
         features=self.kv_lora_rank + self.qk_rope_head_dim,
         axis=-1,
         kernel_init=self.kernel_init,
@@ -1517,8 +1526,12 @@ def setup(self):
         epsilon=self.config.normalization_layer_epsilon,
         kernel_axes=("norm",),
     )
-    self.wkv_b = DenseGeneral(
-        features=(self.num_query_heads, (self.qk_nope_head_dim + self.v_head_dim)),
+    self.wkv_b = dense_general(
+        in_features=self.kv_lora_rank,
+        features=(
+            self.num_query_heads,
+            (self.qk_nope_head_dim + self.v_head_dim),
+        ),
         axis=-1,
         kernel_init=self.kernel_init,
         kernel_axes=("kv_lora", "kv_heads", "kv_head_dim"),

MaxText/layers/gpt3.py

@@ -50,6 +50,7 @@
 EMBED = common_types.EMBED
 
 DenseGeneral = linears.DenseGeneral
+dense_general = linears.dense_general
 NdInitializer = initializers.NdInitializer
 Initializer = initializers.Initializer
 nd_dense_init = initializers.nd_dense_init
@@ -158,7 +159,8 @@ class Gpt3MultiHeadAttention(nn.Module):
   def qkv_projection(self, inputs: Array, proj_name: str):
     """Fused QKV projection"""
 
-    qkv_proj = DenseGeneral(
+    qkv_proj = dense_general(
+        inputs_shape=inputs.shape,
         features=(3, self.num_heads, self.head_dim),
         axis=-1,
         kernel_init=self.kernel_init,
@@ -176,7 +178,8 @@ def qkv_projection(self, inputs: Array, proj_name: str):
 
   def projection(self, inputs: Array, proj_name: str) -> Array:
     """individual projection for one of q, k and v."""
-    proj = DenseGeneral(
+    proj = dense_general(
+        inputs_shape=inputs.shape,
         features=(self.num_heads, self.head_dim),
         axis=-1,
         kernel_init=self.kernel_init,
@@ -192,7 +195,8 @@ def projection(self, inputs: Array, proj_name: str) -> Array:
 
   def out_projection(self, output_dim: int, out: Array) -> Array:
     """output projection"""
-    out_proj = DenseGeneral(
+    out_proj = dense_general(
+        inputs_shape=out.shape,
         features=output_dim,
         axis=(-2, -1),
         kernel_init=self.kernel_init,