diff --git a/rust/src/engine-core-client/src/protocol/logprobs.rs b/rust/src/engine-core-client/src/protocol/logprobs.rs index a01186d0dcb9..6a52cab82510 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs.rs @@ -269,6 +269,19 @@ impl WireLogprobs { ); } + // Empty position lists may be encoded as either [0, 0] or [0, k + 1]. + if token_ids.rows == 0 { + return Ok(Logprobs { + positions: Vec::new(), + }); + } + if token_ids.cols == 0 { + bail_ext_value_decode!( + "{field_prefix}: zero-column logprobs payload with {} rows", + token_ids.rows + ); + } + let mut positions = Vec::with_capacity(token_ids.rows); for ((token_ids_row, logprobs_row), sampled_rank) in token_ids .data diff --git a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs index b105ea8d4a08..42f7ec78b599 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs @@ -303,3 +303,49 @@ fn rejects_non_none_cu_num_generated_tokens() { "messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]" ); } + +#[test] +fn decodes_zero_row_logprobs_as_empty() { + for shape in [[0usize, 0], [0, 3]] { + let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields( + None, + Some(Value::Array(vec![ + ndarray_value("