[terminal_command] TERMINAL ()
   190          z = nn.LayerNorm(   191              param_dtype=self.param_dtype,   192              dtype=self.dtype,   193          )(x)   194          # FIXME (f.srambical): Here, the attention hidden dimension is the same as the FFN's. Usually, FFN hidden dimension is 4x model_dim   195          z = nn.Dense(   196              self.dim,   197              param_dtype=self.param_dtype,   198              dtype=self.dtype,   199          )(z)   200          z = nn.gelu(z)   201          x = x + z   202   203          return x</stdout>```bashcat -n models/lam.py```<stdout>     1  from typing import Dict, Any     2     3  import jax.numpy as jnp     4  import flax.linen as nn     5     6  from utils.preprocess import patchify, unpatchify     7  from utils.nn import STTransformer, VectorQuantizer     8     9    10  class LatentActionModel(nn.Module):    11      """Latent Action ST-ViVit VQ-VAE"""    12    13      in_dim: int    14      model_dim: int    15      latent_dim: int    16      num_latents: int    17      patch_size: int    18      num_blocks: int    19      num_heads: int    20      dropout: float    21      codebook_dropout: float    22      param_dtype: jnp.dtype    23      dtype: jnp.dtype    24      use_flash_attention: bool    25    26      def setup(self):    27          self.patch_token_dim = self.in_dim * self.patch_size**2    28          self.encoder = STTransformer(    29              self.model_dim,    30              self.latent_dim,    31              self.num_blocks,    32              self.num_heads,    33              self.dropout,    34              self.param_dtype,    35              self.dtype,    36              use_flash_attention=self.use_flash_attention,    37          )    38          self.action_in = self.param(    39              "action_in",    40              nn.initializers.lecun_uniform(),    41              (1, 1, 1, self.patch_token_dim),    42          )    43          self.vq = VectorQuantizer(    44              self.latent_dim,    45              self.num_latents,    46              self.codebook_dropout,    47          )    48          self.patch_up = nn.Dense(    49              self.model_dim,    50              param_dtype=self.param_dtype,    51              dtype=self.dtype,    52          )    53          self.action_up = nn.Dense(    54              self.model_dim,    55              param_dtype=self.param_dtype,    56              dtype=self.dtype,    57          )    58          self.decoder = STTransformer(    59              self.model_dim,    60              self.patch_token_dim,    61              self.num_blocks,    62              self.num_heads,    63              self.dropout,    64              self.param_dtype,    65              self.dtype,    66              use_flash_attention=self.use_flash_attention,    67          )    68    69      def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:    70          # --- Encode + VQ ---    71          H, W = batch["videos"].shape[2:4]    72          outputs = self.vq_encode(batch["videos"], training)    73          video_action_patches = self.action_up(outputs["z_q"]) + self.patch_up(    74              outputs["patches"][:, :-1]    75          )    76          del outputs["patches"]    77    78          # --- Decode ---    79          video_recon = self.decoder(video_action_patches)    80          video_recon = video_recon.astype(jnp.float32)    81          video_recon = nn.sigmoid(video_recon)    82          video_recon = video_recon.astype(self.dtype)    83          outputs["recon"] = unpatchify(video_recon, self.patch_size, H, W)    84          return outputs    85    86      def vq_encode(self, videos: Any, training: bool = True) -> Dict[str, Any]:    87          # --- Preprocess videos ---    88          B, T = videos.shape[:2]    89          patches = patchify(videos, self.patch_size)    90          action_pad = jnp.broadcast_to(self.action_in, (B, T, 1, self.patch_token_dim))    91          # FIXME mihir do this the other way around    92          padded_patches = jnp.concatenate((action_pad, patches), axis=2)    93    94          # --- Encode ---    95          z = self.encoder(padded_patches)  # (B, T, N, E)    96          # Get latent action for all future frames    97          z = z[:, 1:, 0]  # (B, T-1, E)    98    99          # --- Vector quantize ---   100          z = z.reshape(B * (T - 1), self.latent_dim)   101          z_q, z, emb, indices = self.vq(z, training)   102          z_q = z_q.reshape(B, T - 1, 1, self.latent_dim)   103          return dict(patches=patches, z_q=z_q, z=z, emb=emb, indices=indices)</stdout>```bashcat -n models/lam.py | sed -n '18,38p'```<stdout>    18      num_blocks: int    19      num_heads: int    20      dropout: float    21      codebook_dropout: float    22      param_dtype: jnp.dtype(Pdb) l358             row = df.iloc[i]359             file_path: str = row["File"]360             event_type = row["Type"]361  362             if i % 50 == 0:363  ->             breakpoint()364  365             match event_type:366                 case "tab":367                     _flush_all_pending_edits()368                     _flush_terminal_output_buffer()(Pdb) quit()*** SystemExit: None(Pdb) Traceback (most recent call last):  File "/fast/home/franz.srambical/crowd-pilot/crowd-pilot/serialize_dataset_array_record.py", line 143, in <module>    main()    ~~~~^^  File "/fast/home/franz.srambical/crowd-pilot/crowd-pilot/serialize_dataset_array_record.py", line 139, in main    to_array_record(cfg)    ~~~~~~~~~~~~~~~^^^^^  File "/fast/home/franz.srambical/crowd-pilot/crowd-pilot/serialize_dataset_array_record.py", line 83, in to_array_record    transcript = session_to_bash_formatted_transcript(        session_df,    )  File "/fast/home/franz.srambical/crowd-pilot/crowd-pilot/serialization_utils.py", line 363, in session_to_bash_formatted_transcript    file_path: str = row["File"]        ^^^^^^^^^^^^  File "/home/franz.srambical/.local/share/uv/python/cpython-3.13.5-linux-x86_64-gnu/lib/python3.13/bdb.py", line 116, in trace_dispatch    return self.dispatch_opcode(frame, arg)           ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^  File "/home/franz.srambical/.local/share/uv/python/cpython-3.13.5-linux-x86_64-gnu/lib/python3.13/bdb.py", line 215, in dispatch_opcode    self.user_opcode(frame)    ~~~~~~~~~~~~~~~~^^^^^^^  File "/home/franz.srambical/.local/share/uv/python/cpython-3.13.5-linux-x86_64-gnu/lib/python3.13/pdb.py", line 448, in user_line    self.interaction(frame, None)    ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^  File "/home/franz.srambical/.local/share/uv/python/cpython-3.13.5-linux-x86_64-gnu/lib/python3.13/pdb.py", line 625, in interaction    self._cmdloop()    ~~~~~~~~~~~~~^^  File "/home/franz.srambical/.local/share/uv/python/cpython-3.13.5-linux-x86_64-gnu/lib/python3.13/pdb.py", line 511, in _cmdloop    self.cmdloop()    ~~~~~~~~~~~~^^  File "/home/franz.srambical/.local/share/uv/python/cpython-3.13.5-linux-x86_64-gnu/lib/python3.13/cmd.py", line 134, in cmdloop    line = input(self.prompt)ValueError: I/O operation on closed file.