[terminal_command] TERMINAL ()
ects/jasmine/jasmine/sample.py", line 185, in _autoreg_sample    generated_vid_BSHWC = _sampling_fn(genie, batch)                          ^^^^^^^^^^^^^^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/jasmine/sample.py", line 171, in _sampling_fn    frames, _ = model.sample(                ^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/jasmine/genie.py", line 253, in sample    return self.sample_diffusion(batch, seq_len, diffusion_steps)           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/jasmine/genie.py", line 682, in sample_diffusion    final_carry = autoregressive_step_fn(initial_carry, jnp.arange(T, seq_len))                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/graph.py", line 2060, in update_context_manager_wrapper    return f(*args, **kwargs)           ^^^^^^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/transforms/iteration.py", line 1344, in scan_wrapper    carry_out, scan_out = jax.lax.scan(                          ^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/transforms/iteration.py", line 1040, in __call__    out = self.f(*args)          ^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/jasmine/genie.py", line 674, in autoregressive_step_fn    final_carry_denoise = denoise_step_fn(carry_denoise, jnp.arange(diffusion_steps))                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/graph.py", line 2060, in update_context_manager_wrapper    return f(*args, **kwargs)           ^^^^^^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/transforms/iteration.py", line 1344, in scan_wrapper    carry_out, scan_out = jax.lax.scan(                          ^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/transforms/iteration.py", line 1040, in __call__    out = self.f(*args)          ^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/jasmine/genie.py", line 654, in denoise_step_fn    v_pred_BSHWC = self.dynamics.diffusion_transformer(vid_BSHWC, t_vector, dt_base, act_embed_window_BSM)                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/jasmine/utils/nn.py", line 740, in __call__    c = te + dte + act_embed_BTM # (B, T, hidden_size)        ~~~~~~~~~^~~~~~~~~~~~~~~  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py", line 1141, in op    return getattr(self.aval, f"_{name}")(self, *args)           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py", line 604, in deferring_binary_op    return binary_op(*args)           ^^^^^^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/jax/_src/numpy/ufunc_api.py", line 183, in __call__    return call(*args)           ^^^^^^^^^^^  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/jax/_src/numpy/ufuncs.py", line 1237, in add    out = lax.add(x, y)          ^^^^^^^^^^^^^TypeError: add got incompatible shapes for broadcasting: (12, 20, 512), (12, 16, 512).--------------------For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.W1022 14:38:48.639839 3559973 pjrt_client.cc:1469] WatchJobStateAsync failed for task goo.gle/debugonly  job_name: "jax_worker": UNAVAILABLE: Cancelling all callsAdditional GRPC error information from remote target coordination_service while calling /tensorflow.CoordinationService/WatchJobState::UNKNOWN:Error received from peer  {grpc_message:"Cancelling all calls", grpc_status:14}srun: error: hai005: task 0: Exited with exit code 1(jasmine) [mihir.mahajan@hai005.haico(jasmine) [mihir.mahajan@hai005.haico(jasmine) [mihir.mahajan@hai005.haicore.berlin:~/Projects/jasmine] $ smi(jasmine) [mihir.mahajan@hai005.haicore.berlin:~/Projects/jasmine] $ smi(jasmine) [mihir.mahajan@hai005.haicore.berlin:~/Projects/jasmine] $ sh slurm/jobs/mihir/berlin/diffusion/breakout/sample_breakout.shSampling from checkpoint: /fast/project/HFMI_ProFOUND/mihir/checkpoints/dyn/train_dyn_dit/32420/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:1269: UserWarning: Sharding info not provided when restoring. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.  warnings.warn(Per-frame SSIM: [0.9998514  0.99978817 0.9997904  0.9995727  0.99978346 0.99954486 0.99977845 0.9997005  0.9994912  0.99961317 0.9997052  0.99972427]Per-frame PSNR: [66.070694 65.63579  66.90076  58.81781  60.54789  61.108715 58.53357 53.421627 55.367622 54.68206  57.980724 59.64276 ]SSIM: 0.999695360660553PSNR: 59.89249801635742W1022 15:38:17.178686 3585757 pjrt_client.cc:1469] WatchJobStateAsync failed for task goo.gle/debugproto    job_name: "jax_worker": CANCELLED: CANCELLEDAdditional GRPC error information from remote target coordination_service while calling /tensorflow.CoordinationService/WatchJobState::UNKNOWN:Error received from peer  {grpc_status:1, grpc_message:"CANCELLED"} [type.googleapis.com/tensorflow.DerivedStatus=''](jasmine) [mihir.mahajan@hai005.haicore.berlin:~/Projects/jasmine] $ sh slurm/jobs/mihir/berlin/diffusion/coinrun/dyn_dit_coinrun.sh          dyn_dit_coinrun_lr_1e-4.sh  sample_coinrun.sh(jasmine) [mihir.mahajan@hai005.haicore.berlin:~/Projects/jasmine] $ sh slurm/jobs/mihir/berlin/diffusion/coinrun/sample_coinrun.sh Sampling from checkpoint: /fast/project/HFMI_ProFOUND/mihir/checkpoints/dyn/train_dyn_dit/32419/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:1269: UserWarning: Sharding info not provided when restoring. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.  warnings.warn(Per-frame SSIM: [0.750481   0.70428133 0.6765215  0.6440681  0.64454424 0.6443093 0.60871184 0.5925966  0.5756893  0.5915309  0.57996917 0.5621155 0.3758793  0.3986288  0.45873427 0.40834853]Per-frame PSNR: [26.71417  24.57674  23.092754 21.809902 21.242256 21.011755 20.4632 20.189533 19.655071 19.36661  19.010542 18.608376 12.160603 12.028679 16.642351 11.588016]SSIM: 0.5760256052017212PSNR: 19.260034561157227W1022 15:40:30.988200 3587719 pjrt_client.cc:1469] WatchJobStateAsync failed for task goo.gle/debugproto   job_name: "jax_worker": CANCELLED: CANCELLEDAdditional GRPC error information from remote target coordination_service while calling /tensorflow.CoordinationService/WatchJobState::UNKNOWN:Error received from peer  {grpc_status:1, grpc_message:"CANCELLED"} [type.googleapis.com/tensorflow.DerivedStatus=''](jasmine) [mihir.mahajan@hai005.haicore.berlin:~/Projects/jasmine] $ salloc: Job 32441 has exceeded its time limit and its allocation has been revoked.                                            srun: Job step aborted: Waiting up to 32 seconds for job step to finish.[2025-10-22T15:41:27.003] error: *** STEP 32441.interactive ON hai005 CANCELLED AT 2025-10-22T15:41:27 DUE TO TIME LIMIT ***^C(jasmine) [mihir.mahajan@hai005.haicore.berlin:~/Projects/jasmine] $ exitsrun: error: hai005: task 0: Exited with exit code 130