p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] jasmine/test_dit.py (python)
from models.dynamics import DynamicsDiffusion
import jax.numpy as jnp
import jax
from flax import nnx

if __name__ == "__main__":
    model = DynamicsDiffusion(
        model_dim=128,
        ffn_dim=128,
        num_latents=128,
        latent_action_dim=128,
        num_blocks=128,
        num_heads=128,
        dropout=0.0,
        param_dtype=jnp.float32,
        dtype=jnp.float32,
        use_flash_attention=False,
        decode=False,
        denoise_steps=64,
        rngs=nnx.Rngs(params=0, dropout=jax.random.key(1))
    )
    batch = {
        "videos": jnp.zeros((4, 16, 64, 64, 3)),
        "rng": jax.random.key(0),
    }
    x_t, x_0 = model(batch)
    print(x_t.shape)
    print(x_0.shape)
[tab] jasmine/test_dit copy.py (python)
from models.dynamics import DynamicsDiffusion
import jax.numpy as jnp
import jax
from flax import nnx

if __name__ == "__main__":
    model = DynamicsDiffusion(
        model_dim=128,
        ffn_dim=128,
        num_latents=128,
        latent_action_dim=128,
        num_blocks=128,
        num_heads=128,
        dropout=0.0,
        param_dtype=jnp.float32,
        dtype=jnp.float32,
        use_flash_attention=False,
        decode=False,
        denoise_steps=64,
        rngs=nnx.Rngs(params=0, dropout=jax.random.key(1))
    )
    batch = {
        "videos": jnp.zeros((4, 16, 64, 64, 3)),
        "rng": jax.random.key(0),
    }
    x_t, x_0 = model(batch)
    print(x_t.shape)
    print(x_0.shape)
[tab] jasmine/test_masking.py (python)
from models.dynamics import DynamicsDiffusion
import jax.numpy as jnp
import jax
from flax import nnx

if __name__ == "__main__":
    model = DynamicsDiffusion(
        model_dim=128,
        ffn_dim=128,
        num_latents=128,
        latent_action_dim=128,
        num_blocks=128,
        num_heads=128,
        dropout=0.0,
        param_dtype=jnp.float32,
        dtype=jnp.float32,
        use_flash_attention=False,
        decode=False,
        denoise_steps=64,
        rngs=nnx.Rngs(params=0, dropout=jax.random.key(1))
    )
    batch = {
        "videos": jnp.zeros((4, 16, 64, 64, 3)),
        "rng": jax.random.key(0),
    }
    x_t, x_0 = model(batch)
    print(x_t.shape)
    print(x_0.shape)
[content] jasmine/test_masking.py (python)
            _rng_prob, _rng_mask = jax.random.split(rng, 2)
            mask_prob = jax.random.uniform(
                _rng_prob, shape=(B, T), minval=0, maxval=self.max_mask_ratio
            )
            mask_BTN = jax.vmap(
                lambda rng, prob: jax.random.bernoulli(rng, prob[..., None], (N,)),
                in_axes=(0, 0),
            )(jax.random.split(_rng_mask, B * T).reshape(B, T, -1), mask_prob)