p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] genie.py (python)
                # assert block.spatial_attention.cached_key is not None
                # assert block.spatial_attention.cached_value is not None
                # assert block.spatial_attention.cache_index is not None
                # block.spatial_attention.cached_key = jnp.zeros((B * T, 1, self.dyna_dim), dtype=self.dtype)
                # block.spatial_attention.cached_value[...] = jnp.zeros((B * T, 1, self.dyna_dim), dtype=self.dtype)
                # block.temporal_attention.cache_index[...] = jnp.array(0, dtype=jnp.int32)
[content] genie.py (python)
        for block in dynamics_causal.transformer.blocks:
            block.spatial_attention.init_cache((B * T, 1, self.dyna_dim), dtype=self.dtype)
            block.temporal_attention.init_cache((B * N, 1, self.dyna_dim), dtype=self.dtype)
[tab] experiments/sample.sh (shellscript)
source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/causal_dynamics_openai_grain_tok_restore"

export XLA_FLAGS='--xla_gpu_autotune_level=0'
export PYTHONUNBUFFERED=1
srun python sample.py \
    --dyna_type "causal" \
    --batch_size 1 \
    --seq_len 3 \
    --start_frame 1 \
    --checkpoint $ckpt_dir \
    --data_dir $data_dir
[content] sample.py (python)
    # Measure autoregressive generation time once and report it.
    elapsed = timeit.timeit(
        lambda: _autoreg_sample(genie, rng, video_batch_BSHWC, action_batch_E).block_until_ready(),
        number=1,
    )
    print(f"Autoregressive generation time: {elapsed:.4f}s")
[content] sample.py (python)

        for block in dynamics_causal.transformer.blocks:
            block.spatial_attention.init_cache((B * seq_len, 1, self.dyna_dim), dtype=self.dtype)
            block.temporal_attention.init_cache((B * (N + 1), 1, self.dyna_dim), dtype=self.dtype)
[content] sample.py (python)
# _autoreg_sample(genie, rng, video_batch_BSHWC, action_batch_E).block_until_ready()
    # # Measure autoregressive generation time once and report it.
    # elapsed = timeit.timeit(
    #     lambda: _autoreg_sample(genie, rng, video_batch_BSHWC, action_batch_E).block_until_ready(),
    #     number=1,
    # )
    # print(f"Autoregressive generation time: {elapsed:.4f}s")
[content] genie.py (python)

    _autoreg_sample(genie, rng, video_batch_BSHWC, action_batch_E).block_until_ready()
    # Measure autoregressive generation time once and report it.
    elapsed = timeit.timeit(
        lambda: _autoreg_sample(genie, rng, video_batch_BSHWC, action_batch_E).block_until_ready(),
        number=1,
    )
    print(f"Autoregressive generation time: {elapsed:.4f}s")
[content] genie.py (python)

            # --- Reset spatial KV caches before each frame ---
            for block in dynamics_causal.transformer.blocks:
                block.spatial_attention.init_cache((B * seq_len, 1, self.dyna_dim), dtype=self.dtype)

            # --- Initialize and run causal loop ---
            init_carry_causal = (
                step_rng,
                current_token_idxs_BSN,
                action_tokens_EL,
                jnp.array(step_t, dtype=jnp.int32),
            )

            # Measure autoregressive generation time once and report it.
            elapsed = timeit.timeit(
                lambda: _autoreg_sample(genie, rng, video_batch_BSHWC, action_batch_E).block_until_ready(),
                number=1,
            )
            print(f"Autoregressive generation time: {elapsed:.4f}s")
            final_carry_causal, _ = jax.lax.scan(
                causal_step_fn, init_carry_causal, jnp.arange(N)
            )
[content] genie.py (python)


            # Measure autoregressive generation time once and report it.
            # elapsed = timeit.timeit(
            #     lambda: _autoreg_sample(genie, rng, video_batch_BSHWC, action_batch_E).block_until_ready(),
            #     number=1,
            # )
            # print(f"Autoregressive generation time: {elapsed:.4f}s")
[content] genie.py (python)

            # --- Reset spatial KV caches before each frame ---
            for block in genie.dynamics.transformer.blocks:
                block.spatial_attention.init_cache((B * seq_len, 1, self.dyna_dim), dtype=self.dtype)
[content] genie.py (python)
Reset spatial KV caches before each frame ---
            for block in genie.dynamics.transformer.blocks:
                block.spatial_attention.init_cache((B * seq_len, 1, self.dyna_dim), dtype=self.dtype)

            # ---
[content] genie.py (python)

            # Measure autoregressive generation time once and report it.
            # elapsed = timeit.timeit(
            #     lambda: _autoreg_sample(genie, rng, video_batch_BSHWC, action_batch_E).block_until_ready(),
            #     number=1,
            # )
            # print(f"Autoregressive generation time: {elapsed:.4f}s")
[content] genie.py (python)
# --- Reset spatial KV caches before each frame ---
            for block in dynamics_causal.transformer.blocks:
                block.spatial_attention.init_cache((B * seq_len, 1, self.dyna_dim), dtype=self.dtype)

            # --- Initialize and run causal loop ---
            init_carry_causal = (
                step_rng,
                current_token_idxs_BSN,
                action_tokens_EL,
                jnp.array(step_t, dtype=jnp.int32),
            )

            # Measure autoregressive generation time once and report it.
            elapsed = timeit.timeit(
                lambda: _autoreg_sample(genie, rng, video_batch_BSHWC, action_batch_E).block_until_ready(),
                number=1,
            )
            print(f"Autoregressive generation time: {elapsed:.4f}s")
            final_carry_causal, _ = jax.lax.scan(
                causal_step_fn, init_carry_causal, jnp.arange(N)
            )
[content] /fast/home/franz.srambical/jafar/genie.py (python)
        for block in dynamics_causal.transformer.blocks:
            block.spatial_attention.init_cache((B * seq_len, 1, self.dyna_dim), dtype=self.dtype)
            block.temporal_attention.init_cache((B * (N + 1), 1, self.dyna_dim), dtype=self.dtype)
[content] sample.py (python)
    B, S, H, W, _ = video_batch_BSHWC.shape
    N = math.ceil(H / args.patch_size) * math.ceil(W / args.patch_size)
    for block in genie.dynamics.transformer.blocks:
        block.spatial_attention.init_cache((B * S, 1, args.dyna_dim), dtype=args.dtype)
        block.temporal_attention.init_cache((B * (N + 1), 1, args.dyna_dim), dtype=args.dtype)
[content] /fast/home/franz.srambical/jafar/genie.py (python)
            block.spatial_attention.init_cache((B * seq_len, (N + 1), self.dyna_dim), dtype=self.dtype)
            block.temporal_attention.init_cache((B * (N + 1), seq_len, self.dyna_dim), dtype=self.dtype)

        @nnx.jit
        def causal_step_fn(
            carry: tuple[jax.Array, jax.Array, jax.Array, jax.Array], step_n: jax.Array
        ) -> tuple[tuple[jax.Array, jax.Array, jax.Array, jax.Array], None]:
            rng, token_idxs_BSN, action_tokens_EL, step_t = carry
            S, N = token_idxs_BSN.shape[1:]
            L = action_tokens_EL.shape[-1]

            # --- Construct + encode video ---
            vid_embed_BSNM = dynamics_causal.patch_embed(token_idxs_BSN)

            # --- Predict transition ---
            action_tokens_BSm1L = jnp.reshape(action_tokens_EL, (B, S - 1, L))
            act_embed_BSm1M = dynamics_causal.action_up(action_tokens_BSm1L)
            act_embed_BSM = jnp.pad(act_embed_BSm1M, ((0, 0), (1, 0), (0, 0)))
            act_embed_BS1M = jnp.reshape(act_embed_BSM, (B, S, 1, act_embed_BSM.shape[-1]))
            vid_embed_BSNp1M = jnp.concatenate([act_embed_BS1M, vid_embed_BSNM], axis=2)
            final_logits_BTNp1V = dynamics_causal.transformer(vid_embed_BSNp1M, (step_t, step_n)) / temperature
            final_logits_BV = final_logits_BTNp1V[:, step_t, step_n, :]

            # --- Sample new tokens for final frame ---
            if sample_argmax:
                sampled_token_idxs_B = jnp.argmax(final_logits_BV, axis=-1)
            else:
                rng, _rng = jax.random.split(rng)
                sampled_token_idxs_B = jax.random.categorical(_rng, final_logits_BV)
            # Update next tokens only
            token_idxs_BSN = token_idxs_BSN.at[:, step_t, step_n].set(sampled_token_idxs_B)
            step_t += 1

            new_carry = (rng, token_idxs_BSN, action_tokens_EL, step_t)
            return new_carry, None

        # --- Run the autoregressive generation using a for loop ---
        rng = batch["rng"]
        current_token_idxs_BSN = token_idxs_BSN
        
        for step_t in range(T, seq_len):
            rng, step_rng = jax.random.split(rng)

            # --- Reset spatial KV caches before each frame ---
            # for block in dynamics_causal.transformer.blocks:
            #     block.spatial_attention.init_cache((B * seq_len, 1, self.dyna_dim), dtype=self.dtype)
            #breakpoint()
[tab] test.ipynb (python)
from flax import nnx

class MyModel(nnx.Module):
  """A simple model containing a Linear layer and dropout."""

  def __init__(self, use_bias: bool, dropout: float, *, rngs: nnx.Rngs):
    self.linear = nnx.Linear(10, 20, use_bias=use_bias, rngs=rngs)
    self.dropout = nnx.Dropout(dropout, rngs=rngs)

  def __call__(self, x):
    x = self.linear(x)
    x = self.dropout(x)
    return x
[tab] test.ipynb (python)
# nnx.state(nnx.eval_shape(lambda: optimizer.model.tokenizer.vq)) # <- works
# nnx.state(nnx.eval_shape(lambda: optimizer.model.tokenizer)).vq) # <- works
# nnx.state(nnx.eval_shape(lambda: optimizer.model)).tokenizer.vq)
# nnx.state(nnx.eval_shape(lambda: optimizer)).model.tokenizer.vq)

# jax.tree.map(ocp.utils.to_shape_dtype_struct, nnx.state(optimizer.model.tokenizer.vq)) # <- works
# jax.tree.map(ocp.utils.to_shape_dtype_struct, nnx.state(optimizer.model.tokenizer)).vq # <- works
# jax.tree.map(ocp.utils.to_shape_dtype_struct, nnx.state(optimizer.model)).tokenizer.vq
# jax.tree.map(ocp.utils.to_shape_dtype_struct, nnx.state(optimizer)).model.tokenizer.vq