candding 0.1.0

use super::*;

/// The configs differ only in `auto_map`, which candding never reads, so the probe alone tells the shapes apart.
#[test]
fn every_block_shape_loads_from_the_header() {
  for (qk, split) in [(false, false), (true, false), (true, true)] {
    let varmap = VarMap::new();
    let model = tiny_model_with(&varmap, qk, split);
    let out = run(&model, &[vec![3, 5, 7, 9]]);
    assert_eq!(
      (out.len(), out[0].len(), out[0][0].len()),
      (1, 4, 16),
      "qk_norm={qk} split_norms={split}"
    );
  }
}

/// Zeroing a norm's weight collapses its output onto its bias, so only a norm that is reached moves the result.
#[test]
fn both_split_norms_reach_the_output() {
  for name in ["layer_norm_1", "layer_norm_2"] {
    let mut varmap = VarMap::new();
    let model = tiny_model_with(&varmap, true, true);
    let rows = [vec![3u32, 5, 7, 9]];
    let before = run(&model, &rows);
    varmap
      .set_one(
        format!("encoder.layer.0.{name}.weight"),
        Tensor::zeros(16, DType::F32, &Device::Cpu).unwrap(),
      )
      .unwrap();
    assert!(
      max_shift(&run(&model, &rows), &before) > 1e-3,
      "{name} is loaded and never applied"
    );
  }
}

#[test]
fn the_qk_norm_reaches_the_scores_only_where_the_header_has_it() {
  let mut varmap = VarMap::new();
  let model = tiny_model_with(&varmap, true, false);
  let rows = [vec![3u32, 5, 7, 9]];
  let before = run(&model, &rows);
  varmap
    .set_one(
      "encoder.layer.0.attention.self.layer_norm_q.weight",
      Tensor::zeros(16, DType::F32, &Device::Cpu).unwrap(),
    )
    .unwrap();
  assert!(
    max_shift(&run(&model, &rows), &before) > 1e-3,
    "layer_norm_q is not applied"
  );

  let base = VarMap::new();
  tiny_model_with(&base, false, false);
  assert!(
    !base
      .data()
      .lock()
      .unwrap()
      .contains_key("encoder.layer.0.attention.self.layer_norm_q.weight"),
    "the base shape must never request a QK norm"
  );
}

#[test]
fn an_unsupported_config_is_refused_at_load() {
  for (field, value) in [
    ("position_embedding_type", "absolute"),
    ("feed_forward_type", "reglu"),
  ] {
    let varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
    let mut config = tiny_config();
    match field {
      "position_embedding_type" => config.position_embedding_type = value.to_string(),
      _ => config.feed_forward_type = value.to_string(),
    }
    assert!(
      matches!(
        JinaBertModel::load(config, vb, AlibiSlopeVariant::Canonical).unwrap_err(),
        crate::error::CanddingError::UnsupportedArchitecture(_)
      ),
      "{field} = {value}"
    );
  }
}

/// The reference feeds the feed-forward `layer_norm_1(residual + attention_output)`, not the attention output.
/// A constant `layer_norm_1` makes the correct layer row-uniform; the second assertion rules out a collapse to zero.
#[test]
fn the_split_feed_forward_consumes_the_normalized_residual() {
  let mut varmap = VarMap::new();
  let model = tiny_model_with(&varmap, true, true);
  let config = tiny_config();
  for layer in 0..config.num_hidden_layers {
    for (suffix, value) in [("weight", 0.0f32), ("bias", 1.0)] {
      varmap
        .set_one(
          format!("encoder.layer.{layer}.layer_norm_1.{suffix}"),
          Tensor::full(value, config.hidden_size, &Device::Cpu).unwrap(),
        )
        .unwrap();
    }
  }
  let out = run(&model, &[vec![3u32, 5, 7, 9]]);
  let first = &out[0][0];
  let spread = out[0]
    .iter()
    .flat_map(|row| row.iter().zip(first).map(|(a, b)| (a - b).abs()))
    .fold(0f32, f32::max);
  assert!(
    spread < 1e-5,
    "every position feeds the same normalized residual, so the rows must agree: spread {spread}"
  );
  let magnitude = first.iter().fold(0f32, |m, v| m.max(v.abs()));
  assert!(
    magnitude > 1e-3,
    "a layer collapsed to zero would be row-uniform too: magnitude {magnitude}"
  );
}

/// `VarMap` creates whatever the load asks for, so the keys left behind are exactly the names a shape requests.
#[test]
fn each_shape_reads_its_own_feed_forward_names() {
  for (split, wanted, unwanted) in [
    (
      false,
      ["mlp.gated_layers.weight", "mlp.wo.weight"],
      ["mlp.up_gated_layer.weight", "mlp.down_layer.weight"],
    ),
    (
      true,
      ["mlp.up_gated_layer.weight", "mlp.down_layer.weight"],
      ["mlp.gated_layers.weight", "mlp.wo.weight"],
    ),
  ] {
    let varmap = VarMap::new();
    tiny_model_with(&varmap, split, split);
    let data = varmap.data().lock().unwrap();
    for name in wanted {
      let key = format!("encoder.layer.0.{name}");
      assert!(data.contains_key(&key), "split_norms={split} wants {key}");
    }
    for name in unwanted {
      let key = format!("encoder.layer.0.{name}");
      assert!(!data.contains_key(&key), "split_norms={split} read {key}");
    }
  }
}