#!/bin/bash
# Ablations that test the mechanism rather than decorate the result.
#
#  0,1  recurrent WITHOUT deep supervision  -> is per-step supervision the
#       ingredient, or is recurrence alone enough?
#  2,3  32 DISTINCT layers, d=128, no weight sharing -> is it the depth, or the
#       weight sharing? Same depth, ~32x the parameters.
#  4,5  R=8 instead of 32 -> how much recurrence does the task actually need?
#
# All use the schedule-matched objective, which is the arm that reached 99.70%.
#SBATCH -J abl
#SBATCH -p gpu-h100
#SBATCH --gpus=h100:1
#SBATCH -c 8
#SBATCH --mem=48g
#SBATCH -t 5:00:00
#SBATCH -o logs/abl-%A_%a.out
#SBATCH --array=0-5%6
source /etc/profile
SCR=/scratch/zt1/project/msml612/user/$USER
module load python/3.10.10
export PYTHONPATH=
cd $SCR/pcmd/src
I=$SLURM_ARRAY_TASK_ID
S=$((I % 2))
COMMON="--mode full --seed $S --pe ape --hard $SCR/pcmd/data/sudoku_hard --augment \
  --heads 4 --bs 256 --steps 120000 --lr 3e-4 --warmup 3000 --K 8 --lam 1.0 \
  --eval_bs 250 --ckpt_every 10000"
case $((I / 2)) in
  0) NAME="abl-nodeepsup-s$S"
     ARGS="--recurrent --R 32 --R_infer 64 --no_deep_sup --d 128 --layers 1" ;;
  1) NAME="abl-deep32L-s$S"
     ARGS="--d 128 --layers 32" ;;
  2) NAME="abl-R8-s$S"
     ARGS="--recurrent --R 8 --R_infer 16 --d 128 --layers 1" ;;
esac
echo "=== $NAME ==="
$SCR/hf_env/bin/python -u train_sudoku.py $ARGS $COMMON --out $SCR/pcmd/runs/$NAME