fix: return error instead of panicking on out-of-bounds categorical index

This commit is contained in:
Manuel Raimann
2026-02-12 13:54:19 +01:00
parent 2d88526d8d
commit 968eb67d07
+29 -2
View File
@@ -467,7 +467,11 @@ impl<T: Clone + Debug> Parameter for CategoricalParam<T> {
fn cast_param_value(&self, param_value: &ParamValue) -> Result<T> {
match param_value {
ParamValue::Categorical(index) => Ok(self.choices[*index].clone()),
ParamValue::Categorical(index) => self
.choices
.get(*index)
.cloned()
.ok_or(Error::Internal("categorical index out of bounds")),
_ => Err(Error::Internal(
"Categorical distribution should return Categorical value",
)),
@@ -708,7 +712,8 @@ impl<T: Categorical + Debug> Parameter for EnumParam<T> {
fn cast_param_value(&self, param_value: &ParamValue) -> Result<T> {
match param_value {
ParamValue::Categorical(index) => Ok(T::from_index(*index)),
ParamValue::Categorical(index) if *index < T::N_CHOICES => Ok(T::from_index(*index)),
ParamValue::Categorical(_) => Err(Error::Internal("categorical index out of bounds")),
_ => Err(Error::Internal(
"Categorical distribution should return Categorical value",
)),
@@ -887,6 +892,17 @@ mod tests {
assert!(param.cast_param_value(&ParamValue::Float(1.0)).is_err());
}
#[test]
fn categorical_param_cast_out_of_bounds() {
let param = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"]);
assert!(param.cast_param_value(&ParamValue::Categorical(3)).is_err());
assert!(
param
.cast_param_value(&ParamValue::Categorical(usize::MAX))
.is_err()
);
}
#[test]
fn bool_param_distribution() {
let param = BoolParam::new();
@@ -955,6 +971,17 @@ mod tests {
assert!(param.cast_param_value(&ParamValue::Float(1.0)).is_err());
}
#[test]
fn enum_param_cast_out_of_bounds() {
let param = EnumParam::<TestEnum>::new();
assert!(param.cast_param_value(&ParamValue::Categorical(3)).is_err());
assert!(
param
.cast_param_value(&ParamValue::Categorical(usize::MAX))
.is_err()
);
}
#[test]
fn float_param_suggest_via_trial() {
let param = FloatParam::new(0.0, 1.0);