diff --git a/src/MaxText/layers/decoders.py b/src/MaxText/layers/decoders.py
index 249a65b9da..94579a6d34 100644
--- a/src/MaxText/layers/decoders.py
+++ b/src/MaxText/layers/decoders.py
@@ -399,14 +399,10 @@ def setup(self):
     self.decoder_layer = self.get_decoder_layers()
     self.norm_layer = self.get_norm_layer(num_features=self.config.emb_dim)
     if self.config.using_pipeline_parallelism:
-      nnx_decoder_classes = self.get_nnx_decoder_layers()
-      if nnx_decoder_classes is not None:
-        pipeline_stage_module = self.get_pipeline_stage_module(nnx_decoder_classes, use_nnx=True)
-      else:
-        pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer, use_nnx=False)
+      pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer)
       remat_policy = self.get_remat_policy()
       self.pipeline_module = pipeline.create_pipeline(
-          config=self.config, mesh=self.mesh, layers=pipeline_stage_module, remat_policy=remat_policy, use_nnx=(nnx_decoder_classes is not None)
+          config=self.config, mesh=self.mesh, layers=pipeline_stage_module, remat_policy=remat_policy,
       )
 
   def minimal_policy(self, with_context=False):
@@ -685,7 +681,7 @@ def scan_decoder_layers(self, cfg, decoder_layer, length, metadata_axis_name, me
         config=cfg, mesh=mesh, name=metadata_axis_name, quant=self.quant, **kwargs  # pytype: disable=wrong-keyword-args
     )
 
-  def get_pipeline_stage_module(self, decoder_blocks, use_nnx=False):
+  def get_pipeline_stage_module(self, decoder_blocks):
     """get pipeline stage module
 
     Args:
@@ -703,20 +699,6 @@ def get_layer_to_pipeline(blocks, cfg):
     cfg = self.config
     base_stage = get_layer_to_pipeline(decoder_blocks, cfg)
 
-    if use_nnx:
-      if cfg.num_layers_per_pipeline_stage == 1:
-        return base_stage
-      else:
-        return lambda config, mesh, model_mode, rngs, quant=None: SequentialNNXWrapper(
-            decoder_layer_class=base_stage,
-            num_decoder_layers=cfg.num_layers_per_pipeline_stage,
-            config=config,
-            mesh=mesh,
-            model_mode=model_mode,
-            rngs=rngs,
-            quant=quant
-        )
-
     if cfg.set_remat_policy_on_layers_per_stage:
       policy = self.get_remat_policy()
       base_stage = self.set_remat_policy([base_stage], policy)[0]
@@ -730,6 +712,7 @@ def get_layer_to_pipeline(blocks, cfg):
           "layers_per_stage",
           self.mesh,
           in_axes_tuple=(nn.broadcast,) * 4,
+          model_mode=self.model_mode,
       )
     else:
       stage_module = SequentialBlockDecoderLayers(
diff --git a/src/MaxText/layers/embeddings.py b/src/MaxText/layers/embeddings.py
index 76fb1b8a78..5285e530e2 100644
--- a/src/MaxText/layers/embeddings.py
+++ b/src/MaxText/layers/embeddings.py
@@ -342,6 +342,7 @@ def __call__(
           "The embedding dims of the rotary position embedding" "must match the hidden dimension of the inputs."
       )
 
+
     position = position[:, :, jnp.newaxis, jnp.newaxis]
     sinusoid_inp = position / self.timescale
     sin = jnp.sin(sinusoid_inp).astype(inputs.dtype)
diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py
index 20c9dbe51b..2720fa53fd 100644
--- a/src/MaxText/layers/pipeline.py
+++ b/src/MaxText/layers/pipeline.py
@@ -12,628 +12,436 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-""" Pipeline layer wrapping a decoder layer(s). Supports circular pipelining """
-
+"""Pipeline layer wrapping a decoder layer(s). Supports circular pipelining.
+NNX Implementation.
+"""
 import functools
-from typing import Any, Callable
-
-import numpy as np
+from typing import Any, Optional, Callable
 
-from jax import numpy as jnp
-from jax.sharding import Mesh
 import jax
-import jax.ad_checkpoint
+import jax.numpy as jnp
+from jax.sharding import Mesh, NamedSharding, PartitionSpec
 
-from flax.core import meta
-from flax import linen as nn
 from flax import nnx
+from flax import linen as nn
 
 from MaxText.common_types import Config, MODEL_MODE_TRAIN, EP_AS_CONTEXT
-from MaxText.sharding import all_gather_over_fsdp
-from MaxText import max_logging
 from MaxText.layers import nnx_wrappers
-
-
-class Pipeline(nnx.Module):
-  """NNX Module that implements pipelining across stages.
-
-  This module will loop over microbatches and execute the main body with a vmap for both the inputs and weights.
-  This will produce a pipeline pattern if the stage dimension is sharded.
-
-  Supports circular pipelines, and multiple layers per stage are used when a module that executes multiple layers
-  is passed as the layers input.
-
-  Attributes:
-    config: Importantly contains num_pipeline_microbatches, num_pipeline_repeats.
-    layers: A callable (NNX class or Linen class) that each stage can execute. It can either be a single layer such as a
-      LlamaDecoderLayer instance or scanned/looped set of decoder layers to execute multiple layers per stage.
-    mesh:  The device mesh of the system.
-    remat_policy: Remat policy to use for the loop iterations
-  """
-
-  def __init__(
-      self,
-      layers: Callable | type,
-      config: Config,
-      mesh: Mesh,
-      rngs: nnx.Rngs = None,
-      remat_policy: Any = None,
-  ):
-    """Initialize Pipeline with NNX or Linen decoder layers.
-
-    Args:
-      layers: Either an NNX class (type) or Linen class (type) to instantiate for each stage
-      config: Model configuration
-      mesh: Device mesh for sharding
-      rngs: Optional NNX RNG state (passed by ToLinen wrapper)
-      remat_policy: Remat policy for loop iterations
-    """
-    self.config = config
-    self.mesh = mesh
-    self.rngs = rngs
-    self.remat_policy = remat_policy
-
-    self.num_stages = self.config.ici_pipeline_parallelism * self.config.dcn_pipeline_parallelism
-    self.forwarding_delay = 2 if self.config.pipeline_delay_activation_forwarding else 1
-    self.pipeline_microbatch_size = self.config.micro_batch_size_to_train_on // self.config.num_pipeline_microbatches
-    microbatches_per_stage = self.config.num_pipeline_microbatches // self.num_stages
-    self.microbatches_per_stage = microbatches_per_stage
-    self.use_circ_storage = self.need_circ_storage()
-
-    if self.config.expert_shard_attention_option == EP_AS_CONTEXT:
-      self.batch_axis_name = "activation_batch_no_exp"
-      self.seq_len_axis_name = "activation_length"
+from MaxText.layers.initializers import variable_to_logically_partitioned
+
+# ==============================================================================
+# Shared Sharding Logic (Heuristics)
+# ==============================================================================
+
+def _infer_partition_names(value, repeats=1, base_names=None):
+    """Helper to infer logical axis names (PartitionSpec) for pipeline weights."""
+    ndim = value.ndim
+    
+    if repeats > 1:
+        # Structure: [Repeats, Stage, ...]
+        if base_names is not None:
+             new_names = (None, 'stage') + base_names
+        else:
+            # Heuristics for [Repeats, Stage, ...]
+            if ndim == 4: # [Repeats, Stage, Embed, MLP]
+                new_names = (None, 'stage', 'fsdp', 'tensor')
+            elif ndim == 3: # [Repeats, Stage, Bias]
+                new_names = (None, 'stage', 'fsdp')
+            else:
+                new_names = (None, 'stage') + (None,) * (ndim - 2)
     else:
-      self.batch_axis_name = "activation_batch"
-      self.seq_len_axis_name = "activation_length_no_exp"
-
-    # Detect if layers is a Linen class/instance or NNX class
-    self._is_linen = (isinstance(layers, type) and issubclass(layers, nn.Module)) or isinstance(layers, nn.Module)
+        # Structure: [Stage, ...]
+        if base_names is not None:
+            new_names = ('stage',) + base_names
+        else:
+             # Heuristics for [Stage, ...]
+            if ndim == 3: # [Stage, Embed, MLP]
+                new_names = ('stage', 'fsdp', 'tensor')
+            elif ndim == 2: # [Stage, Bias]
+                new_names = ('stage', 'fsdp')
+            else:
+                new_names = ('stage',) + (None,) * (ndim - 1)
+    
+    return new_names
+
+# ==============================================================================
+# Pipeline Class
+# ==============================================================================
 
-    if self._is_linen:
-      if isinstance(layers, nn.Module):
+class Pipeline(nnx.Module):
+    """NNX Implementation of the MaxText Pipeline."""
+
+    def __init__(
+        self,
+        layers: nn.Module,
+        config: Config,
+        mesh: Mesh,
+        rngs: nnx.Rngs,
+        remat_policy: Any = None
+    ):
+        self.config = config
+        self.mesh = mesh
         self.layers = layers
-      else:
-        self.layers = layers(config=config, mesh=mesh, model_mode=MODEL_MODE_TRAIN)
-      self._linen_variables = None
-    else:
-      # Create num_stages independent NNX instances, stored as attributes for
-      # NNX pytree tracking (not as Python lists).
-      for s in range(self.num_stages):
-        stage_rngs = nnx.Rngs(s)
-        instance = layers(
-            config=config,
-            mesh=mesh,
-            model_mode=MODEL_MODE_TRAIN,
-            rngs=stage_rngs,
-            quant=None,
-        )
-        setattr(self, f'stage_{s}', instance)
-
-  def need_circ_storage(self):
-    return (
-        self.config.num_pipeline_repeats > 1
-        and self.config.num_pipeline_microbatches > self.num_stages * self.forwarding_delay
-    )
-
-  def iterations_to_complete_first_microbatch_one_repeat(self):
-    """Returns iterations for microbatch 0 to complete one repeat."""
-    return self.forwarding_delay * (self.num_stages - 1)
-
-  def iterations_to_complete_first_microbatch(self):
-    """Returns iterations for microbatch 0 to complete all repeats."""
-    return (
-        self.config.num_pipeline_microbatches * (self.config.num_pipeline_repeats - 1)
-        + self.iterations_to_complete_first_microbatch_one_repeat()
-    )
-
-  def init_states(self, inputs):
-    """Initialize pipeline loop state buffers.
-
-    Assumes inputs are reshaped to [num_microbatches, micro_batch_size, sequence, embed].
-
-    Returns:
-      Dictionary containing:
-        - shift: Buffer for rotating outputs [num_stages, micro_size, sequence, embed]
-        - prev_outputs: Same shape as shift (only used with pipeline_delay_activation_forwarding)
-        - state_io: Input/output buffer [num_stages, microbatches/stages, micro_size, sequence, embed]
-        - circ_storage: Circular storage buffer (only when num_microbatches > num_stages)
-        - circ_storage_mover: One-iteration delay buffer for circ_storage
-        - loop_iteration: Iteration counter (starts at 0)
-    """
-    shift = jnp.zeros((self.num_stages,) + inputs.shape[1:], dtype=inputs.dtype)
-    shift = self._with_logical_constraint(
-        shift,
-        ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"),
-    )
-
-    if self.config.pipeline_delay_activation_forwarding:
-      prev_outputs = jnp.zeros((self.num_stages,) + inputs.shape[1:], dtype=inputs.dtype)
-      prev_outputs = self._with_logical_constraint(
-          prev_outputs,
-          ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"),
-      )
-    else:
-      prev_outputs = None
+        self.remat_policy = remat_policy
+
+        # --- Dimensions ---
+        self.num_stages = self.config.ici_pipeline_parallelism * self.config.dcn_pipeline_parallelism
+        self.forwarding_delay = 2 if self.config.pipeline_delay_activation_forwarding else 1
+        self.total_microbatches = self.config.num_pipeline_microbatches
+        self.microbatch_size = self.config.micro_batch_size_to_train_on // self.total_microbatches
+        self.microbatches_per_stage = self.total_microbatches // self.num_stages
+        self.use_circ_storage = self._need_circ_storage()
+
+        # Axis Naming
+        if hasattr(self.config, 'expert_shard_attention_option') and self.config.expert_shard_attention_option == EP_AS_CONTEXT:
+            self.batch_axis_name = "activation_batch_no_exp"
+            self.seq_len_axis_name = "activation_length"
+        else:
+            self.batch_axis_name = "activation_batch"
+            self.seq_len_axis_name = "activation_length_no_exp"
+
+        self.input_shape = (self.microbatch_size, self.config.max_target_length, self.config.emb_dim)
+
+        # --- Eager Initialization with JIT + Out Shardings ---
+        init_rng_key = rngs.params()
+        repeats = self.config.num_pipeline_repeats
+
+        # 1. Define the Generator (Pure JAX, no side effects)
+        def generate_vars_fn(key):
+            def init_single_stage(k):
+                # Create dummy inputs locally to ensure shape inference works inside JIT
+                dummy_in = jnp.zeros(self.input_shape, dtype=jnp.float32)
+                dummy_pos_shape = (self.input_shape[0], self.input_shape[1])
+                dummy_positions = jnp.zeros(dummy_pos_shape, dtype=jnp.int32)
+                dummy_segments = jnp.zeros(dummy_pos_shape, dtype=jnp.int32)
+                
+                return self.layers.init(
+                    {'params': k},
+                    dummy_in,
+                    dummy_segments,
+                    dummy_positions,
+                    True,
+                    MODEL_MODE_TRAIN
+                )
+
+            if repeats > 1:
+                repeat_keys = jax.random.split(key, repeats)
+                # Outer vmap over repeats, Inner vmap over stages
+                stage_rng_keys = jax.vmap(lambda k: jax.random.split(k, self.num_stages))(repeat_keys)
+                return jax.vmap(jax.vmap(init_single_stage))(stage_rng_keys)
+            else:
+                stage_rng_keys = jax.random.split(key, self.num_stages)
+                return jax.vmap(init_single_stage)(stage_rng_keys)
+
+        # 2. Abstract Evaluation (Zero Memory)
+        # Determine the shape and structure of the variables without allocating data
+        abstract_variables = jax.eval_shape(generate_vars_fn, init_rng_key)
+
+        # 3. Compute Target Sharding (Heuristics)
+        # We assume the output structure matches what we expect for partitioning
+        def get_target_sharding(abstract_leaf):
+            # abstract_leaf is a ShapeDtypeStruct. We check its ndim.
+            logical_names = _infer_partition_names(abstract_leaf, repeats=repeats)
+            
+            if logical_names:
+                mesh_axes = nn.logical_to_mesh_axes(logical_names, self.config.logical_axis_rules)
+                return NamedSharding(self.mesh, PartitionSpec(*mesh_axes))
+            
+            # Default fallback: Replicated
+            return NamedSharding(self.mesh, PartitionSpec())
+
+        sharding_tree = jax.tree.map(get_target_sharding, abstract_variables)
+
+        # 4. JIT Compile with Out Shardings (Direct Allocation)
+        # This tells XLA to compile a kernel that writes outputs directly to their 
+        # final sharded destination, bypassing the single-device OOM bottleneck.
+        sharded_variables = jax.jit(generate_vars_fn, out_shardings=sharding_tree)(init_rng_key)
+
+        # --- Register with NNX ---
+        self.stage_params = self._to_nnx_structure(sharded_variables)
+        
+        # Handle batch_stats if present
+        if 'batch_stats' in sharded_variables:
+            pass
+    # ==========================================================================
+    # Helper Methods
+    # ==========================================================================
+
+    def _to_nnx_structure(self, node):
+        if isinstance(node, (dict, nn.FrozenDict)):
+            return nnx.Dict({k: self._to_nnx_structure(v) for k, v in node.items()})
+        elif isinstance(node, (list, tuple)):
+            return nnx.List([self._to_nnx_structure(v) for v in node])
+        else:
+            return nnx.Param(node)
+
+    def _to_pure_dict(self, node):
+        if hasattr(node, 'items'): return {k: self._to_pure_dict(v) for k, v in node.items()}
+        elif isinstance(node, (list, tuple)): return [self._to_pure_dict(v) for v in node]
+        elif hasattr(node, 'value'): return node.value
+        return node 
+
+    def _with_logical_constraint(self, x, axis_names):
+        if axis_names is None: return x
+        mesh_axes = nn.logical_to_mesh_axes(axis_names, self.config.logical_axis_rules)
+        sharding = NamedSharding(self.mesh, PartitionSpec(*mesh_axes))
+        return jax.lax.with_sharding_constraint(x, sharding)
+
+    def _need_circ_storage(self):
+        return (self.config.num_pipeline_repeats > 1 and 
+                self.config.num_pipeline_microbatches > self.num_stages * self.forwarding_delay)
+
+    def iterations_to_complete_first_microbatch_one_repeat(self):
+        return self.forwarding_delay * (self.num_stages - 1)
+
+    # ==========================================================================
+    # Buffer Management
+    # ==========================================================================
+
+    def init_states(self, inputs):
+        shift = jnp.zeros((self.num_stages,) + inputs.shape[1:], dtype=inputs.dtype)
+        shift = self._with_logical_constraint(shift, ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"))
+        
+        if self.config.pipeline_delay_activation_forwarding:
+            prev_outputs = jnp.zeros((self.num_stages,) + inputs.shape[1:], dtype=inputs.dtype)
+            prev_outputs = self._with_logical_constraint(prev_outputs, ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"))
+        else:
+            prev_outputs = None
+            
+        state_io = jnp.reshape(inputs, (self.num_stages, self.microbatches_per_stage) + inputs.shape[1:])
+        state_io = self._with_logical_constraint(state_io, ("activation_stage", None, self.batch_axis_name, self.seq_len_axis_name, "activation_embed"))
+        
+        if self.use_circ_storage:
+            circ_storage = jnp.zeros((self.num_stages,) + inputs.shape, dtype=inputs.dtype)
+            circ_storage_mover = shift
+        else:
+            circ_storage = None
+            circ_storage_mover = None
+            
+        return {"state_io": state_io, "shift": shift, "circ_storage": circ_storage, "circ_storage_mover": circ_storage_mover, "loop_iteration": jnp.array(0, dtype=jnp.int32), "prev_outputs": prev_outputs, "rng_stream": jax.random.PRNGKey(0)}
+
+    def get_iteration_inputs(self, loop_iteration, state_io, circ_storage, shift):
+        state_io_batch_idx = loop_iteration % self.microbatches_per_stage
+        state_io_slice = state_io[:, state_io_batch_idx]
+        
+        if self.use_circ_storage:
+            circ_storage_batch_idx = loop_iteration % self.config.num_pipeline_microbatches
+            circular_stage_in = circ_storage[:, circ_storage_batch_idx]
+        else:
+            circular_stage_in = shift
+            
+        first_stage_in = jnp.where(loop_iteration < self.config.num_pipeline_microbatches, state_io_slice, circular_stage_in)
+        
+        stages_in = jnp.where(jax.lax.broadcasted_iota("int32", shift.shape, 0) == 0, first_stage_in, shift)
+        stages_in = self._with_logical_constraint(stages_in, ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"))
+        return stages_in
+
+    def get_new_loop_state(self, output, loop_state):
+        old_state_io = loop_state["state_io"]
+        old_circ_storage = loop_state["circ_storage"]
+        old_circ_storage_mover = loop_state["circ_storage_mover"]
+        loop_iteration = loop_state["loop_iteration"]
+        old_prev_outputs = loop_state["prev_outputs"]
+        
+        def _rotate_right(arr):
+            last = jax.lax.slice_in_dim(arr, self.num_stages - 1, self.num_stages, axis=0)
+            except_last = jax.lax.slice_in_dim(arr, 0, self.num_stages - 1, axis=0)
+            return jnp.concatenate([last, except_last], axis=0)
+        def _shift_right(arr):
+            padding = [[1, 0]] + [[0, 0]] * (arr.ndim - 1)
+            return jax.lax.slice(jnp.pad(arr, padding), [0] * arr.ndim, arr.shape)
+        def _update_shift(output_in):
+            if self.config.num_pipeline_repeats == 1 or self.use_circ_storage:
+                return _shift_right(output_in)
+            return _rotate_right(output_in)
+        
+        if self.config.pipeline_delay_activation_forwarding:
+            new_shift = _update_shift(old_prev_outputs)
+            new_prev_outputs = output
+        else:
+            new_shift = _update_shift(output)
+            new_prev_outputs = None
+            
+        if self.use_circ_storage:
+            rotated = _rotate_right(old_circ_storage_mover)
+            rotated = jnp.expand_dims(rotated, 1)
+            offset = (loop_iteration - self.iterations_to_complete_first_microbatch_one_repeat() - 1) % self.config.num_pipeline_microbatches
+            new_circ_storage = jax.lax.dynamic_update_slice_in_dim(old_circ_storage, rotated, offset, axis=1)
+            new_circ_storage_mover = output
+        else:
+            new_circ_storage = None
+            new_circ_storage_mover = None
+            
+        stream_buf_idx = loop_iteration % self.microbatches_per_stage
+        stream_slice = old_state_io[:, stream_buf_idx]
+        padding = [[0, 1]] + [[0, 0]] * (stream_slice.ndim - 1)
+        stream_slice = jax.lax.slice_in_dim(jnp.pad(stream_slice, padding), 1, stream_slice.shape[0] + 1, axis=0)
+        stream_slice = jnp.where(jax.lax.broadcasted_iota("int32", stream_slice.shape, 0) == self.num_stages - 1, output, stream_slice)
+        stream_slice = jnp.expand_dims(stream_slice, 1)
+        new_state = jax.lax.dynamic_update_slice_in_dim(old_state_io, stream_slice, stream_buf_idx, axis=1)
+        
+        return {"state_io": new_state, "shift": new_shift, "circ_storage": new_circ_storage, "circ_storage_mover": new_circ_storage_mover, "loop_iteration": loop_iteration + 1, "prev_outputs": new_prev_outputs}
+
+    # ==========================================================================
+    # Logic for VMAP and Weight Gathering
+    # ==========================================================================
+
+    def get_microbatch_and_repeat_ids(self, loop_iteration):
+        microbatches_processed = jnp.maximum(loop_iteration - self.forwarding_delay * jnp.arange(self.num_stages), 0)
+        return microbatches_processed % self.total_microbatches, (microbatches_processed // self.total_microbatches).astype(jnp.int32)
+
+    def shard_dim_by_stages(self, x, dim: int):
+        dims_mapping = [PartitionSpec.UNCONSTRAINED] * x.ndim
+        dims_mapping[dim] = "stage"
+        sharding = NamedSharding(self.mesh, PartitionSpec(*dims_mapping))
+        return jax.lax.with_sharding_constraint(x, sharding)
+
+    def vmap_parallel_gather(self, weights, repeat_ids, repeat_dim_in_weights, stages_dim_in_weights):
+        def _gather_one(x, repeat_id):
+            return jnp.squeeze(jax.lax.dynamic_slice_in_dim(x, repeat_id, 1, repeat_dim_in_weights), repeat_dim_in_weights)
+        gathered_weights_stage_dim = 0
+        repeat_ids = self.shard_dim_by_stages(repeat_ids, 0)
+        weights = self.shard_dim_by_stages(weights, stages_dim_in_weights)
+        stage_weights = jax.vmap(_gather_one, in_axes=(stages_dim_in_weights, 0), out_axes=gathered_weights_stage_dim)(weights, repeat_ids)
+        return self.shard_dim_by_stages(stage_weights, gathered_weights_stage_dim)
+
+    def get_current_stage_weights(self, pipeline_weights, loop_iteration):
+        if self.config.num_pipeline_repeats <= 1: return pipeline_weights
+        _, repeat_ids = self.get_microbatch_and_repeat_ids(loop_iteration)
+        return jax.tree.map(lambda w: self.vmap_parallel_gather(w, repeat_ids, 0, 1), pipeline_weights)
+
+    def permute_output_micro_per_stage_dim(self, output):
+        microbatch_0_idx = self.iterations_to_complete_first_microbatch_one_repeat() % self.microbatches_per_stage
+        permutation = (jnp.arange(self.microbatches_per_stage) + microbatch_0_idx) % self.microbatches_per_stage
+        return output[:, permutation]
+
+    # ==========================================================================
+    # Integration Methods
+    # ==========================================================================
+
+    def get_pipeline_remat_policy(self):
+        if self.config.remat_policy == "custom": return self.remat_policy
+        save_input_policy = jax.checkpoint_policies.save_only_these_names("iteration_input", "decoder_layer_input")
+        if self.remat_policy is not None: return jax.checkpoint_policies.save_from_both_policies(self.remat_policy, save_input_policy)
+        return save_input_policy
+
+    def get_weight_sharding(self, *init_args):
+        variables = self.stage_params
+        repeats = self.config.num_pipeline_repeats
+        def _infer_partition_spec(node):
+            if hasattr(node, 'value'):
+                names = _infer_partition_names(node.value, repeats=repeats, base_names=None)
+                return nn.LogicallyPartitioned(node.value, names).get_partition_spec()
+            return None
+        specs = jax.tree.map(_infer_partition_spec, variables, is_leaf=lambda x: hasattr(x, 'value'))
+        return {'params': {'stage_params': specs['params']}}
+
+    def _all_gather_over_fsdp(self, params, partition_spec):
+        def _remove_fsdp_from_spec(spec):
+            if isinstance(spec, PartitionSpec):
+                new_spec = []
+                for axis in spec:
+                    if isinstance(axis, str) and axis in ("fsdp", "fsdp_transpose"): new_spec.append(None)
+                    elif isinstance(axis, (list, tuple)): new_spec.append(tuple(a for a in axis if a not in ("fsdp", "fsdp_transpose")))
+                    else: new_spec.append(axis)
+                return PartitionSpec(*new_spec)
+            return spec
+        def _remove_fsdp_sharding(sharding_tree):
+            return jax.tree.map(lambda x: NamedSharding(self.mesh, _remove_fsdp_from_spec(x.spec)) if isinstance(x, NamedSharding) else x, sharding_tree)
+        physical = nn.logical_to_mesh_sharding(partition_spec, mesh=self.mesh, rules=self.config.logical_axis_rules)
+        physical_no_fsdp = _remove_fsdp_sharding(physical)
+        return jax.lax.with_sharding_constraint(params, physical_no_fsdp)
+
+    # ==========================================================================
+    # Main Call (__call__)
+    # ==========================================================================
+
+    def __call__(self, inputs: jax.Array, segment_ids: Optional[jax.Array] = None, positions: Optional[jax.Array] = None, deterministic: bool = True, model_mode: str = MODEL_MODE_TRAIN, partition_spec: Any = None) -> jax.Array:
+        inputs = inputs.reshape((self.total_microbatches, self.microbatch_size, self.config.max_target_length, self.config.emb_dim))
+        if positions is not None: positions = positions.reshape((self.total_microbatches, self.microbatch_size, self.config.max_target_length))
+        if segment_ids is not None: segment_ids = segment_ids.reshape((self.total_microbatches, self.microbatch_size, self.config.max_target_length))
+        
+        loop_state = self.init_states(inputs)
+        layer_variables = self._to_pure_dict(self.stage_params)
+        
+        compute_dtype = getattr(self.config, 'compute_dtype', jnp.bfloat16)
+        if isinstance(compute_dtype, str):
+            compute_dtype = {'bfloat16': jnp.bfloat16, 'float32': jnp.float32, 'float16': jnp.float16}.get(compute_dtype, jnp.bfloat16)
+        layer_variables = jax.tree.map(lambda x: x.astype(compute_dtype), layer_variables)
+
+        if self.config.pipeline_fsdp_ag_once and partition_spec is not None:
+            try:
+                if "params" in partition_spec and "stage_params" in partition_spec["params"]:
+                     params_only = layer_variables['params']
+                     params_spec = partition_spec['params']['stage_params']['params']
+                     layer_variables['params'] = self._all_gather_over_fsdp(params_only, params_spec)
+            except (KeyError, TypeError): pass
+
+        def scan_body(carry, _):
+            iteration = carry['loop_iteration']
+            current_rng = carry['rng_stream']
+            step_rng, next_rng = jax.random.split(current_rng)
+            stage_rngs = jax.random.split(step_rng, self.num_stages)
+            stages_inputs = self.get_iteration_inputs(iteration, carry['state_io'], carry['circ_storage'], carry['shift'])
+            stages_inputs = jax.ad_checkpoint.checkpoint_name(stages_inputs, "iteration_input")
+            mb_ids, _ = self.get_microbatch_and_repeat_ids(iteration)
+            stages_positions = jnp.take(positions, mb_ids, axis=0) if positions is not None else None
+            stages_segment_ids = jnp.take(segment_ids, mb_ids, axis=0) if segment_ids is not None else None
+            current_vars = self.get_current_stage_weights(layer_variables, iteration)
+
+            def execution_logic(vars_in, inputs, rngs, pos, seg):
+                def stage_fn(v, x, r, po, se):
+                    rngs_dict = {'dropout': r} if not deterministic else {}
+                    mutables = ['aux_loss', 'intermediates']
+                    return self.layers.apply(v, x, se, po, deterministic, model_mode, rngs=rngs_dict, mutable=mutables)
+                vmap_axes = [0, 0, 0]
+                vmap_args = [vars_in, inputs, rngs]
+                vmap_axes.append(0 if pos is not None else None); vmap_args.append(pos)
+                vmap_axes.append(0 if seg is not None else None); vmap_args.append(seg)
+                return jax.vmap(stage_fn, in_axes=tuple(vmap_axes))(*vmap_args)
+
+            if self.config.set_remat_policy_on_pipeline_iterations:
+                policy = self.get_pipeline_remat_policy()
+                execution_logic = jax.checkpoint(execution_logic, policy=policy, prevent_cse=not self.config.scan_pipeline_iterations)
+
+            stages_output, stages_mutables = execution_logic(current_vars, stages_inputs, stage_rngs, stages_positions, stages_segment_ids)
+            if hasattr(self.config, 'scan_layers') and self.config.scan_layers:
+                 if isinstance(stages_output, tuple): stages_output = stages_output[0]
+            
+            new_loop_state = self.get_new_loop_state(stages_output, carry)
+            new_loop_state['rng_stream'] = next_rng
+            return new_loop_state, stages_mutables
+
+        bubble_iterations = self.forwarding_delay * (self.num_stages - 1)
+        real_iterations = self.config.num_pipeline_microbatches * self.config.num_pipeline_repeats
+        total_ticks = real_iterations + bubble_iterations
+        final_state, stacked_mutables = jax.lax.scan(scan_body, loop_state, None, length=total_ticks)
+        output = final_state['state_io']
+        output = self.permute_output_micro_per_stage_dim(output)
+        output = output.reshape((self.config.micro_batch_size_to_train_on, self.config.max_target_length, self.config.emb_dim))
+        return output
 
-    state_io = jnp.reshape(inputs, (self.num_stages, self.microbatches_per_stage) + inputs.shape[1:])
-    state_io = self._with_logical_constraint(
-        state_io,
-        ("activation_stage", None, self.batch_axis_name, self.seq_len_axis_name, "activation_embed"),
-    )
 
-    if self.use_circ_storage:
-      circ_storage = jnp.zeros((self.num_stages,) + inputs.shape, dtype=inputs.dtype)
-      circ_storage_mover = shift
-    else:
-      circ_storage = None
-      circ_storage_mover = None
-
-    init_loop_state = {
-        "state_io": state_io,
-        "shift": shift,
-        "circ_storage": circ_storage,
-        "circ_storage_mover": circ_storage_mover,
-        "loop_iteration": 0,
-        "prev_outputs": prev_outputs,
-    }
-    return init_loop_state
-
-  def _with_logical_constraint(self, tensor, logical_axis_names):
-    """Applies logical sharding constraints to tensor."""
-    return nn.with_logical_constraint(
-        tensor,
-        logical_axis_names,
-        rules=self.config.logical_axis_rules,
-        mesh=self.mesh,
-    )
-
-  def get_iteration_inputs(self, loop_iteration, state_io, circ_storage, shift):
-    """Constructs input array for all stages for this iteration.
-
-    Returns array of shape [stages, micro_size, sequence, embed] with rotated outputs
-    from previous iteration, except stage 0 which gets new input from state_io or circ_storage.
-    """
-    state_io_batch_idx = loop_iteration % self.microbatches_per_stage
-    state_io_slice = state_io[:, state_io_batch_idx]
-
-    if self.use_circ_storage:
-      circ_storage_batch_idx = loop_iteration % self.config.num_pipeline_microbatches
-      circular_stage_in = circ_storage[:, circ_storage_batch_idx]
-    else:
-      circular_stage_in = shift
-
-    first_stage_in = jnp.where(loop_iteration < self.config.num_pipeline_microbatches, state_io_slice, circular_stage_in)
-
-    def select_state_or_input(first_stage_in, shift):
-      return jnp.where(jax.lax.broadcasted_iota("int32", shift.shape, 0) == 0, first_stage_in, shift)
-
-    stages_in = select_state_or_input(first_stage_in, shift)
-    stages_in = self._with_logical_constraint(
-        stages_in,
-        ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"),
-    )
-    return stages_in
-
-  def shard_dim_by_stages(self, x, dim: int):
-    """Shards the specified dimension by stage."""
-    dims_mapping = [jax.sharding.PartitionSpec.UNCONSTRAINED] * x.ndim
-    dims_mapping[dim] = "stage"
-    dims_mapping = tuple(dims_mapping)
-    sharding = jax.sharding.NamedSharding(self.mesh, jax.sharding.PartitionSpec(*dims_mapping))
-    return jax.lax.with_sharding_constraint(x, sharding)
-
-  def get_microbatch_and_repeat_ids(self, loop_iteration):
-    """Gets microbatch and repeat IDs for all stages at this iteration."""
-    microbatches_processed = jnp.maximum(loop_iteration - self.forwarding_delay * jnp.arange(self.num_stages), 0)
-    microbatch_ids = microbatches_processed % self.config.num_pipeline_microbatches
-    repeat_ids = microbatches_processed // self.config.num_pipeline_microbatches
-    return microbatch_ids, repeat_ids
-
-  def vmap_parallel_gather(self, weights, repeat_ids, repeat_dim_in_weights, stages_dim_in_weights):
-    """Sharded parallel gather where each stage has its own weights and gets one slice.
-
-    Args:
-      weights: Per-stage data to gather from.
-      repeat_ids: Integer tensor of shape [num_stages] with repeat indices per stage.
-      repeat_dim_in_weights: Dimension where repeat_ids are applied (removed in output).
-      stages_dim_in_weights: Dimension representing parallel stages.
-
-    Returns:
-      Per-stage gathered values with repeat_dim_in_weights removed.
-    """
-    def _gather_one(x, repeat_id):
-      return jnp.squeeze(jax.lax.dynamic_slice_in_dim(x, repeat_id, 1, repeat_dim_in_weights), repeat_dim_in_weights)
-
-    gathered_weights_stage_dim = 0
-    repeat_ids = self.shard_dim_by_stages(repeat_ids, 0)
-    weights = self.shard_dim_by_stages(weights, stages_dim_in_weights)
-    stage_weights = jax.vmap(_gather_one, in_axes=(stages_dim_in_weights, 0), out_axes=gathered_weights_stage_dim)(
-        weights, repeat_ids
-    )
-    stage_weights = self.shard_dim_by_stages(stage_weights, gathered_weights_stage_dim)
-    return stage_weights
-
-  def vmap_gather(self, xs, ids, ids_dim):
-    """Stage-wise sharded gather with shared input but different offsets per stage.
-
-    Args:
-      xs: Data shared by all stages.
-      ids: Integer tensor of shape [num_stages] with offsets per stage.
-      ids_dim: Dimension where ids are applied (output has [num_stages] size here).
-
-    Returns:
-      Per-stage gathered values with ids_dim size replaced with [num_stages].
-    """
-    def _gather_one(x, i):
-      return jnp.squeeze(jax.lax.dynamic_slice_in_dim(x, i, 1, ids_dim), ids_dim)
-
-    ids = self.shard_dim_by_stages(ids, 0)
-    outs = jax.vmap(_gather_one, in_axes=(None, 0), out_axes=ids_dim)(xs, ids)
-    return self.shard_dim_by_stages(outs, 0)
-
-  def get_new_loop_state(self, output, loop_state):
-    """Updates all pipeline buffers after one iteration.
-
-    Updates shift, state_io, circ_storage, circ_storage_mover, and prev_outputs
-    to advance the pipeline by one step.
-    """
-    old_state_io = loop_state["state_io"]
-    old_circ_storage = loop_state["circ_storage"]
-    old_circ_storage_mover = loop_state["circ_storage_mover"]
-    loop_iteration = loop_state["loop_iteration"]
-    old_prev_outputs = loop_state["prev_outputs"]
-
-    def _rotate_right(arr):
-      last = jax.lax.slice_in_dim(arr, self.num_stages - 1, self.num_stages, axis=0)
-      except_last = jax.lax.slice_in_dim(arr, 0, self.num_stages - 1, axis=0)
-      return jnp.concatenate([last, except_last], axis=0)
-
-    def _shift_right(arr):
-      padding = [[1, 0]] + [[0, 0]] * (arr.ndim - 1)
-      return jax.lax.slice(jnp.pad(arr, padding), [0] * arr.ndim, arr.shape)
-
-    def _update_shift(output_in):
-      if self.config.num_pipeline_repeats == 1 or self.use_circ_storage:
-        return _shift_right(output_in)
-      else:
-        return _rotate_right(output_in)
-
-    if self.config.pipeline_delay_activation_forwarding:
-      new_shift = _update_shift(old_prev_outputs)
-      new_prev_outputs = output
-    else:
-      new_shift = _update_shift(output)
-      new_prev_outputs = None
-
-    if self.use_circ_storage:
-      def _rotate_right_and_update(circ_storage_mover_in, circ_storage_in):
-        rotated = _rotate_right(circ_storage_mover_in)
-        rotated = jnp.expand_dims(rotated, 1)
-        offset = (
-            loop_iteration - self.iterations_to_complete_first_microbatch_one_repeat() - 1
-        ) % self.config.num_pipeline_microbatches
-        return jax.lax.dynamic_update_slice_in_dim(circ_storage_in, rotated, offset, axis=1)
-
-      new_circ_storage = _rotate_right_and_update(old_circ_storage_mover, old_circ_storage)
-      new_circ_storage_mover = output
-    else:
-      new_circ_storage = None
-      new_circ_storage_mover = None
-
-    stream_buf_idx = loop_iteration % self.microbatches_per_stage
-    stream_slice = old_state_io[:, stream_buf_idx]
-
-    def _update_state_io(state_in, stream_slice, output):
-      padding = [[0, 1]] + [[0, 0]] * (stream_slice.ndim - 1)
-      stream_slice = jax.lax.slice_in_dim(jnp.pad(stream_slice, padding), 1, stream_slice.shape[0] + 1, axis=0)
-      stream_slice = jnp.where(
-          jax.lax.broadcasted_iota("int32", stream_slice.shape, 0) == self.num_stages - 1, output, stream_slice
-      )
-      stream_slice = jnp.expand_dims(stream_slice, 1)
-      return jax.lax.dynamic_update_slice_in_dim(state_in, stream_slice, stream_buf_idx, axis=1)
-
-    new_state = _update_state_io(old_state_io, stream_slice, output)
-
-    new_loop_state = {
-        "state_io": new_state,
-        "shift": new_shift,
-        "circ_storage": new_circ_storage,
-        "circ_storage_mover": new_circ_storage_mover,
-        "loop_iteration": loop_iteration + 1,
-        "prev_outputs": new_prev_outputs,
-    }
-    return new_loop_state
-
-  def permute_output_micro_per_stage_dim(self, output):
-    """Permutes output to correct microbatch ordering after pipeline completion."""
-    microbatch_0_idx = self.iterations_to_complete_first_microbatch() % self.microbatches_per_stage
-    permutation = (
-        np.arange(self.microbatches_per_stage) + microbatch_0_idx
-    ) % self.microbatches_per_stage
-    output = output[:, permutation]
-    return output
-
-  def _initialize_linen_parameters(self, sample_input, sample_seg_ids, sample_positions, deterministic, model_mode):
-    """Initialize Linen module parameters for all stages."""
-    if self._linen_variables is not None:
-      return
-
-    linen_rngs = {'params': jax.random.PRNGKey(0), 'dropout': jax.random.PRNGKey(1)}
-
-    base_params = self.layers.init(
-        linen_rngs,
-        sample_input,
-        sample_seg_ids,
-        sample_positions,
-        deterministic,
-        model_mode,
-    )
-
-    stage_params = {}
-    for stage_idx in range(self.num_stages):
-      stage_params[f'stage_{stage_idx}'] = jax.tree_util.tree_map(lambda x: x, base_params)
-
-    self._linen_variables = {'params': stage_params}
-
-  def _run_stages_linen(
-      self,
-      stages_inputs,
-      stages_segment_ids,
-      stages_positions,
-      deterministic,
-      model_mode,
-  ):
-    """Run stages using Linen module with manual vmap."""
-    stage_params_list = [self._linen_variables['params'][f'stage_{i}'] for i in range(self.num_stages)]
-    stacked_params = jax.tree_util.tree_map(
-        lambda *xs: jnp.stack(xs, axis=0),
-        *stage_params_list
-    )
-
-    def apply_stage(stage_params, stage_input, stage_seg_ids, stage_pos):
-      output = self.layers.apply(
-          stage_params,
-          stage_input,
-          stage_seg_ids,
-          stage_pos,
-          deterministic,
-          model_mode,
-      )
-      if isinstance(output, tuple):
-        return output[0]
-      return output
-
-    if stages_segment_ids is None:
-      vmapped_apply = jax.vmap(
-          lambda p, i, pos: apply_stage(p, i, None, pos),
-          in_axes=(0, 0, 0),
-          out_axes=0
-      )
-      stages_outputs = vmapped_apply(stacked_params, stages_inputs, stages_positions)
-    else:
-      vmapped_apply = jax.vmap(apply_stage, in_axes=(0, 0, 0, 0), out_axes=0)
-      stages_outputs = vmapped_apply(stacked_params, stages_inputs, stages_segment_ids, stages_positions)
-
-    return stages_outputs
-
-  def _run_stages_vmapped(
-      self,
-      stages_inputs,
-      stages_segment_ids,
-      stages_positions,
-      deterministic,
-      model_mode,
-  ):
-    """Run all stages in parallel using JAX vmap over NNX instances."""
-    stage_0 = getattr(self, 'stage_0')
-    graphdef, state_0 = nnx.split(stage_0)
-
-    states = [state_0]
-    for s in range(1, self.num_stages):
-      instance = getattr(self, f'stage_{s}')
-      _, state_s = nnx.split(instance)
-      states.append(state_s)
-
-    stacked_state = jax.tree_util.tree_map(
-        lambda *xs: jnp.stack(xs, axis=0),
-        *states
-    )
-
-    def call_stage(state, stage_input, stage_seg_ids, stage_pos):
-      module = nnx.merge(graphdef, state)
-      output = module(stage_input, stage_seg_ids, stage_pos, deterministic, model_mode)
-      if isinstance(output, tuple):
-        return output[0]
-      return output
-
-    if stages_segment_ids is None:
-      def call_stage_no_seg(state, stage_input, stage_pos):
-        module = nnx.merge(graphdef, state)
-        output = module(stage_input, None, stage_pos, deterministic, model_mode)
-        if isinstance(output, tuple):
-          return output[0]
-        return output
 
-      vmapped_call = jax.vmap(call_stage_no_seg, in_axes=(0, 0, 0), out_axes=0)
-      stages_outputs = vmapped_call(stacked_state, stages_inputs, stages_positions)
-    else:
-      vmapped_call = jax.vmap(call_stage, in_axes=(0, 0, 0, 0), out_axes=0)
-      stages_outputs = vmapped_call(stacked_state, stages_inputs, stages_segment_ids, stages_positions)
-
-    return stages_outputs
-
-  def run_one_iteration(
-      self,
-      loop_state,
-      positions,
-      segment_ids,
-      deterministic,
-      model_mode,
-  ):
-    """Run one loop iteration: get inputs, execute stages, update state."""
-    state_io = loop_state["state_io"]
-    shift = loop_state["shift"]
-    circ_storage = loop_state["circ_storage"]
-    loop_iteration = loop_state["loop_iteration"]
-
-    microbatch_ids, _ = self.get_microbatch_and_repeat_ids(loop_iteration)
-
-    stages_inputs = self.get_iteration_inputs(loop_iteration, state_io, circ_storage, shift)
-    stages_inputs = jax.ad_checkpoint.checkpoint_name(stages_inputs, "iteration_input")
-    stages_positions = self.vmap_gather(positions, microbatch_ids, 0) if positions is not None else None
-    stages_segment_ids = self.vmap_gather(segment_ids, microbatch_ids, 0) if segment_ids is not None else None
-
-    if self._is_linen:
-      stages_output = self._run_stages_linen(
-          stages_inputs,
-          stages_segment_ids,
-          stages_positions,
-          deterministic,
-          model_mode,
-      )
-    else:
-      stages_output = self._run_stages_vmapped(
-          stages_inputs,
-          stages_segment_ids,
-          stages_positions,
-          deterministic,
-          model_mode,
-      )
-
-    new_state = self.get_new_loop_state(stages_output, loop_state)
-    return new_state
-
-  def get_pipeline_remat_policy(self):
-    """Returns the remat policy for pipeline iterations."""
-    if self.config.remat_policy == "custom":
-      return self.remat_policy
-
-    save_input_policy = jax.checkpoint_policies.save_only_these_names("iteration_input", "decoder_layer_input")
-    if self.remat_policy is not None:
-      remat_policy = jax.checkpoint_policies.save_from_both_policies(self.remat_policy, save_input_policy)
-    else:
-      remat_policy = save_input_policy
-    return remat_policy
-
-  def __call__(
-      self,
-      inputs: jnp.ndarray,
-      segment_ids: jnp.ndarray,
-      positions: jnp.ndarray,
-      deterministic: bool,
-      model_mode=MODEL_MODE_TRAIN,
-      partition_spec=None,
-  ) -> jnp.ndarray:
-    """Maps decoder layer inputs to outputs using pipeline parallelism.
-
-    Reshapes inputs into microbatches, runs pipeline iterations with bubble
-    handling, and returns outputs reshaped to original batch size.
-    """
-    inputs = inputs.reshape(
-        (
-            self.config.num_pipeline_microbatches,
-            self.pipeline_microbatch_size,
-            self.config.max_target_length,
-            self.config.emb_dim,
-        )
-    )
-
-    if self._is_linen and self._linen_variables is None:
-      example_input = inputs[0]
-      example_seg_ids = segment_ids[0] if segment_ids is not None else None
-      example_pos = positions[0] if positions is not None else None
-      self._initialize_linen_parameters(example_input, example_seg_ids, example_pos, deterministic, model_mode)
-
-    ag_sharding = jax.sharding.NamedSharding(self.mesh, jax.sharding.PartitionSpec(None, None))
-    if positions is not None:
-      positions = jax.lax.with_sharding_constraint(positions, ag_sharding)
-      positions = positions.reshape(
-          (self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length)
-      )
-
-    if segment_ids is not None:
-      segment_ids = jax.lax.with_sharding_constraint(segment_ids, ag_sharding)
-      segment_ids = segment_ids.reshape(
-          (self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length)
-      )
-
-    loop_state = self.init_states(inputs)
-
-    bubble_iterations = self.forwarding_delay * (self.num_stages - 1)
-    real_iterations = self.config.num_pipeline_microbatches * self.config.num_pipeline_repeats
-    total_iterations = real_iterations + bubble_iterations
-
-    if self.config.scan_pipeline_iterations:
-      def run_iteration_scannable(loop_state, xs):
-        return (
-            self.run_one_iteration(
-                loop_state, positions, segment_ids, deterministic, model_mode
-            ),
-            None,
-        )
-
-      if self.config.set_remat_policy_on_pipeline_iterations:
-        run_iteration_scannable = jax.checkpoint(
-            run_iteration_scannable,
-            prevent_cse=False,
-            policy=self.get_pipeline_remat_policy(),
-        )
-
-      loop_state, _ = jax.lax.scan(run_iteration_scannable, loop_state, None, length=total_iterations)
+# ==============================================================================
+# Factory
+# ==============================================================================
+
+def add_stage_axis_to_partitioning(variable, repeats=1):
+    """Metadata helper for ToLinen."""
+    partitioned_obj = variable_to_logically_partitioned(variable)
+    if isinstance(partitioned_obj, nn.LogicallyPartitioned):
+        base_names = partitioned_obj.names
+        value = partitioned_obj.value
     else:
-      for _ in range(total_iterations):
-        loop_state = self.run_one_iteration(
-            loop_state, positions, segment_ids, deterministic, model_mode
-        )
-
-    final_output = self.permute_output_micro_per_stage_dim(loop_state["state_io"])
-
-    final_output = jnp.reshape(
-        final_output, (self.config.micro_batch_size_to_train_on, self.config.max_target_length, self.config.emb_dim)
-    )
-
-    return final_output
-
-
-class PipelineToLinen(nnx_wrappers.ToLinen):
-  """Wrap NNX Pipeline as a Linen module.
-
-  This allows the NNX Pipeline to be used within the Linen Decoder module.
-  """
-  pass
-
-
-def create_pipeline(
-    config: Config,
-    layers: Callable | type,
-    mesh: Mesh,
-    remat_policy: Any = None,
-    use_nnx: bool = True,
-) -> PipelineToLinen:
-  """Factory function to create a Pipeline wrapped as a Linen module.
-
-  Args:
-    config: Model configuration
-    layers: NNX or Linen decoder layer class to use for pipeline stages
-    mesh: Device mesh for sharding
-    remat_policy: Remat policy for loop iterations
-    use_nnx: Whether to use NNX pipeline (True) or Linen (False)
-
-  Returns:
-    PipelineToLinen wrapper around the NNX Pipeline
-  """
-  if not use_nnx:
-    raise ValueError("This implementation only supports NNX pipelines (use_nnx=True)")
-
-  wrapped = PipelineToLinen(
-      Pipeline,
-      kwargs={
-          'layers': layers,
-          'config': config,
-          'mesh': mesh,
-          'remat_policy': remat_policy,
-      }
-  )
-
-  return wrapped
+        value = partitioned_obj
+        if not hasattr(value, 'ndim'): return value
+        base_names = None
+    
+    new_names = _infer_partition_names(value, repeats=repeats, base_names=base_names)
+    return nn.LogicallyPartitioned(value, new_names)
+
+def create_pipeline(config: Config, layers: Callable | type, mesh: Mesh, remat_policy: Any = None) -> nnx_wrappers.ToLinen:
+    repeats = getattr(config, 'num_pipeline_repeats', 1)
+    metadata_fn = functools.partial(add_stage_axis_to_partitioning, repeats=repeats)
+    return nnx_wrappers.to_linen(Pipeline, config=config, mesh=mesh, layers=layers, remat_policy=remat_policy, name="pipeline_module", abstract_init=False, metadata_fn=metadata_fn)
\ No newline at end of file