p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] train_lam.py (python)

    # extract and manually clip grad norm for logging (actual clipping is done in the optax.chain)
    raw_grad_norm = optax.global_norm(grads)
    g_norm_clipped = jnp.minimum(raw_grad_norm, args.grad_clip_threshold)
    metrics["grad_norm"] = g_norm_clipped
[content] train_dynamics.py (python)
# extract and manually clip grad norm for logging (actual clipping is done in the optax.chain)
    raw_grad_norm = optax.global_norm(grads)
    g_norm_clipped = jnp.minimum(raw_grad_norm, args.grad_clip_threshold)
    metrics["grad_norm"] = g_norm_clipped
[content] train_tokenizer.py (python)

    # extract and manually clip grad norm for logging (actual clipping is done in the optax.chain)
    raw_grad_norm = optax.global_norm(grads)
    g_norm_clipped = jnp.minimum(raw_grad_norm, args.grad_clip_threshold)
    metrics["grad_norm"] = g_norm_clipped
[content] train_dynamics.py (python)

    # extract and manually clip grad norm for logging (actual clipping is done in the optax.chain)
    raw_grad_norm = optax.global_norm(grads)
    metrics["grad_norm"] = jnp.minimum(raw_grad_norm, args.grad_clip_threshold) if args.grad_clip_threshold else raw_grad_norm
[content] train_lam.py (python)

    # extract and manually clip grad norm for logging (actual clipping is done in the optax.chain)
    raw_grad_norm = optax.global_norm(grads)
    metrics["grad_norm"] = jnp.minimum(raw_grad_norm, args.grad_clip_threshold) if args.grad_clip_threshold else raw_grad_norm
[content] train_lam.py (python)

    if args.grad_clip_threshold:
        tx = optax.chain(
            optax.clip_by_global_norm(args.grad_clip_threshold),
            optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4)
        )
    else:
        tx = optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4)
[content] train_dynamics.py (python)

    if args.grad_clip_threshold:
        tx = optax.chain(
            optax.clip_by_global_norm(args.grad_clip_threshold),
            optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4)
        )
    else:
        tx = optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4)
[content] train_dynamics.py (python)
    lr_schedule = get_lr_schedule(args.lr_schedule, 
                                  args.init_lr, 
                                  args.max_lr, 
                                  args.decay_end, 
                                  args.num_steps, 
                                  args.warmup_steps, 
                                  args.wsd_decay_steps)
    if args.grad_clip_threshold:
        tx = optax.chain(
            optax.clip_by_global_norm(args.grad_clip_threshold),
            optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4)
        )
    else:
        tx = optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4, mu_dtype=args.dtype)
[content] train_lam.py (python)
    lr_schedule = get_lr_schedule(args.lr_schedule, 
                                  args.init_lr, 
                                  args.max_lr, 
                                  args.decay_end, 
                                  args.num_steps, 
                                  args.warmup_steps, 
                                  args.wsd_decay_steps)
    if args.grad_clip_threshold:
        tx = optax.chain(
            optax.clip_by_global_norm(args.grad_clip_threshold),
            optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4)
        )
    else:
        tx = optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4, mu_dtype=args.dtype)
[content] train_tokenizer.py (python)
    lr_schedule = get_lr_schedule(args.lr_schedule, 
                                  args.init_lr, 
                                  args.max_lr, 
                                  args.decay_end, 
                                  args.num_steps, 
                                  args.warmup_steps, 
                                  args.wsd_decay_steps)
    if args.grad_clip_threshold:
        tx = optax.chain(
            optax.clip_by_global_norm(args.grad_clip_threshold),
            optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4)
        )
    else:
        tx = optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4, mu_dtype=args.dtype)
[content] utils/nn.py (python)
class CausalTransformerBlock(nn.Module):
    model_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # LayerNorm + Causal Self-Attention
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        seq_len = z.shape[1]
        # Causal mask: (1, 1, seq_len, seq_len)
        causal_mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=bool))
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.model_dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z, mask=causal_mask)
        x = x + z

        # Feedforward
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        x = x + z

        return x

class CausalTransformer(nn.Module):
    model_dim: int
    out_dim: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # Input projection and normalization
        x = nn.Sequential(
            [
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
                nn.Dense(self.model_dim,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                ),
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
            ]
        )(x)
        # Causal transformer blocks
        for _ in range(self.num_blocks):
            x = CausalTransformerBlock(
                model_dim=self.model_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
            )(x)
        # Output projection
        x = nn.Dense(
            self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        return x  # (B, T, E)
[content] models/dynamics.py (python)
class DynamicsAutoregressive(nn.Module):
    """Autoregressive (causal) dynamics model"""

    model_dim: int
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    def setup(self):
        # Use a causal transformer instead of STTransformer
        from utils.nn import CausalTransformer  # Make sure this exists or implement it
        self.dynamics = CausalTransformer(
            self.model_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
        )
        self.patch_embed = nn.Embed(self.num_latents, self.model_dim)
        self.action_up = nn.Dense(
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        # No random masking, just embed the tokens
        vid_embed = self.patch_embed(batch["video_tokens"])
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        logits = self.dynamics(vid_embed, causal=True)  # Pass causal=True if needed
        return dict(token_logits=logits, mask=None)