From 3c46bbb1985d79cbe14435c90194ff927b06d2fa Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 9 Jul 2026 17:57:33 +0000 Subject: [PATCH 1/5] Add gemma-4-12B unified (encoder-free) preprocessing support The gemma-4-12B "unified" model is encoder-free: it consumes raw merged pixel patches and raw waveform frames directly, not the SigLIP/log-mel contract of the standard gemma-4 processor. Add the two ort-extensions pieces needed so onnxruntime-genai can preprocess natively (no external HuggingFace processor): * Audio: new Gemma4UnifiedAudioFrames feature op that chunks a decoded 16 kHz waveform into fixed audio_samples_per_token (640) frames, reproducing Gemma4UnifiedAudioFeatureExtractor exactly (pad to a whole number of frames, reshape to (num_tokens, 640)). * Image: no new op required. The unified 48px merged patch (patch_dim 6912) produced by HuggingFace via (16px patchify -> 3x3 patches_merge) is provably identical to a direct 48px patchify, so the existing Gemma4ImageTransform yields the unified contract when configured with patch_size=48, pooling_kernel_size=1. A test locks this equivalence in. Adds test configs under test/data/models/gemma-4-unified/ and gtests that verify the 640-sample audio framing and the 6912-dim image contract (cross-checked against the standard config's teacher patches). Signed-off-by: Justin Chu --- shared/api/gemma4_audio_features.hpp | 75 ++++++++++++ shared/api/speech_extractor.cc | 3 +- .../audio_feature_extraction.json | 23 ++++ .../gemma-4-unified/image_processor.json | 27 +++++ test/pp_api_test/test_feature_extraction.cc | 38 +++++- test/pp_api_test/test_processor.cc | 114 ++++++++++++++++++ 6 files changed, 278 insertions(+), 2 deletions(-) create mode 100644 test/data/models/gemma-4-unified/audio_feature_extraction.json create mode 100644 test/data/models/gemma-4-unified/image_processor.json diff --git a/shared/api/gemma4_audio_features.hpp b/shared/api/gemma4_audio_features.hpp index 1f90ad57f..d26374ce4 100644 --- a/shared/api/gemma4_audio_features.hpp +++ b/shared/api/gemma4_audio_features.hpp @@ -344,4 +344,79 @@ class Gemma4LogMel { std::vector mel_filters_; // (n_freq x feature_size), row-major }; +// Gemma 4 *unified* (encoder-free, gemma-4-12B) audio feature extraction. +// +// Unlike Gemma4LogMel (128-dim USM log-mel), the unified model has no audio +// encoder: each audio soft token is simply a fixed-length chunk of the raw +// 16 kHz waveform. This op reproduces HuggingFace +// ``Gemma4UnifiedAudioFeatureExtractor._extract_waveform_features`` exactly: +// zero-pad the waveform to a multiple of ``audio_samples_per_token`` and +// reshape it into ``(num_tokens, audio_samples_per_token)`` frames. +// +// Pipeline: AudioDecoder -> Gemma4UnifiedAudioFrames +// +// Inputs: float (1, num_samples) — mono PCM at `sampling_rate` Hz +// Outputs: float (num_tokens, audio_samples_per_token) — raw waveform frames +// +// The frame-level mask is intentionally not emitted here: for the single-clip +// inference path the downstream processor fills an all-true mask, and the +// batch framework zero-pads ragged clips when stacking. +class Gemma4UnifiedAudioFrames { + public: + Gemma4UnifiedAudioFrames() = default; + + OrtxStatus Compute(const ortc::Tensor& pcm_input, + ortc::Tensor& frames_out) { + const auto& pcm_shape = pcm_input.Shape(); + if (pcm_shape.size() != 2 || pcm_shape[0] != 1) { + return {kOrtxErrorInvalidArgument, + "[Gemma4UnifiedAudioFrames]: expected (1, num_samples) float input"}; + } + + const int64_t num_samples = pcm_shape[1]; + const int64_t spt = audio_samples_per_token_; + // Zero-pad to a whole number of frames (ceil division), matching HF's + // ``pad_len = (-len(waveform)) % audio_samples_per_token``. + const int64_t num_tokens = (num_samples + spt - 1) / spt; + + float* out = frames_out.Allocate({num_tokens, spt}); + if (num_tokens == 0) { + return {}; + } + // Fill the (possibly padded) tail of the last frame with the padding value, + // then copy the real samples over the front. + std::fill(out, out + static_cast(num_tokens) * spt, padding_value_); + std::copy(pcm_input.Data(), pcm_input.Data() + num_samples, out); + return {}; + } + + template + OrtxStatus Init(const DictT& attrs) { + for (const auto& [key, value] : attrs) { + if (key == "audio_samples_per_token" || key == "feature_size") { + // ``feature_size`` is accepted as an alias: HF sets feature_size == + // audio_samples_per_token (both default to 640). + audio_samples_per_token_ = std::get(value); + } else if (key == "sampling_rate") { + sampling_rate_ = std::get(value); + } else if (key == "padding_value") { + padding_value_ = static_cast(std::get(value)); + } else { + return {kOrtxErrorInvalidArgument, + "[Gemma4UnifiedAudioFrames]: unknown attribute '" + key + "'"}; + } + } + if (audio_samples_per_token_ <= 0) { + return {kOrtxErrorInvalidArgument, + "[Gemma4UnifiedAudioFrames]: audio_samples_per_token must be positive"}; + } + return {}; + } + + private: + int64_t audio_samples_per_token_ = 640; // 640 samples = 40 ms @ 16 kHz + int64_t sampling_rate_ = 16000; + float padding_value_ = 0.0f; +}; + } // namespace ort_extensions diff --git a/shared/api/speech_extractor.cc b/shared/api/speech_extractor.cc index c32100ddd..97679f172 100644 --- a/shared/api/speech_extractor.cc +++ b/shared/api/speech_extractor.cc @@ -17,7 +17,8 @@ Operation::KernelRegistry SpeechFeatureExtractor::kernel_registry_ = { {"NemoLogMel", []() { return CreateKernelInstance(&NemoLogMel::Compute); }}, {"PerFeatureNormalize", []() { return CreateKernelInstance(&PerFeatureNormalize::Compute); }}, {"Phi4AudioEmbed", []() { return CreateKernelInstance(&Phi4AudioEmbed::Compute); }}, - {"Gemma4LogMel", []() { return CreateKernelInstance(&Gemma4LogMel::Compute); }}}; + {"Gemma4LogMel", []() { return CreateKernelInstance(&Gemma4LogMel::Compute); }}, + {"Gemma4UnifiedAudioFrames", []() { return CreateKernelInstance(&Gemma4UnifiedAudioFrames::Compute); }}}; SpeechFeatureExtractor::SpeechFeatureExtractor() : OrtxObjectImpl(extObjectKind_t::kOrtxKindFeatureExtractor) {} diff --git a/test/data/models/gemma-4-unified/audio_feature_extraction.json b/test/data/models/gemma-4-unified/audio_feature_extraction.json new file mode 100644 index 000000000..abf92ed7d --- /dev/null +++ b/test/data/models/gemma-4-unified/audio_feature_extraction.json @@ -0,0 +1,23 @@ +{ + "feature_extraction": { + "sequence": [ + { + "operation": { + "name": "audio_decoder", + "type": "AudioDecoder" + } + }, + { + "operation": { + "name": "gemma4_unified_audio_frames", + "type": "Gemma4UnifiedAudioFrames", + "attrs": { + "audio_samples_per_token": 640, + "sampling_rate": 16000, + "padding_value": 0.0 + } + } + } + ] + } +} diff --git a/test/data/models/gemma-4-unified/image_processor.json b/test/data/models/gemma-4-unified/image_processor.json new file mode 100644 index 000000000..ac9252e7b --- /dev/null +++ b/test/data/models/gemma-4-unified/image_processor.json @@ -0,0 +1,27 @@ +{ + "processor": { + "name": "gemma_4_unified_image_processing", + "transforms": [ + { + "operation": { + "name": "decode_image", + "type": "DecodeImage", + "attrs": { + "color_space": "RGB" + } + } + }, + { + "operation": { + "name": "gemma4_image_transform", + "type": "Gemma4ImageTransform", + "attrs": { + "patch_size": 48, + "max_soft_tokens": 280, + "pooling_kernel_size": 1 + } + } + } + ] + } +} diff --git a/test/pp_api_test/test_feature_extraction.cc b/test/pp_api_test/test_feature_extraction.cc index 7acb1497e..e4ec78696 100644 --- a/test/pp_api_test/test_feature_extraction.cc +++ b/test/pp_api_test/test_feature_extraction.cc @@ -446,4 +446,40 @@ TEST(ExtractorTest, TestGemma4AudioFeatureExtractionMultiFile) { ASSERT_EQ(err, kOrtxOK); ASSERT_EQ(mask_dims, 2ULL); ASSERT_EQ(mask_shape[0], 2); -} \ No newline at end of file +} + +TEST(ExtractorTest, TestGemma4UnifiedAudioFrames) { + // gemma-4-12B "unified" (encoder-free) audio: raw 16 kHz waveform chunked + // into fixed 640-sample frames. Pipeline: AudioDecoder -> Gemma4UnifiedAudioFrames + const char* audio_path[] = {"data/jfk.flac"}; + OrtxObjectPtr raw_audios; + extError_t err = OrtxLoadAudios(raw_audios.ToBeAssigned(), audio_path, 1); + ASSERT_EQ(err, kOrtxOK); + + OrtxObjectPtr feature_extractor( + OrtxCreateSpeechFeatureExtractor, "data/models/gemma-4-unified/audio_feature_extraction.json"); + OrtxObjectPtr result; + err = OrtxFeatureExtraction(feature_extractor.get(), raw_audios.get(), result.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + + // Output 0: raw waveform frames — float (batch, num_tokens, 640) + OrtxObjectPtr tensor; + err = OrtxTensorResultGetAt(result.get(), 0, tensor.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + + const float* data{}; + const int64_t* shape{}; + size_t num_dims; + err = OrtxGetTensorData(tensor.get(), reinterpret_cast(&data), &shape, &num_dims); + ASSERT_EQ(err, kOrtxOK); + ASSERT_EQ(num_dims, 3ULL); // (batch, num_tokens, samples_per_token) + ASSERT_EQ(shape[0], 1); // single audio + ASSERT_EQ(shape[2], 640); // 640 raw samples per token + EXPECT_GT(shape[1], 0); // at least one frame + + // All values finite and within the normalized PCM range. + for (int64_t i = 0; i < std::min(shape[1] * 640, 5000); ++i) { + ASSERT_TRUE(std::isfinite(data[i])) << "frame value at index " << i << " is not finite"; + ASSERT_LE(std::abs(data[i]), 4.0f) << "frame value at index " << i << " out of range"; + } +} diff --git a/test/pp_api_test/test_processor.cc b/test/pp_api_test/test_processor.cc index fea0d19f9..b48f89d03 100644 --- a/test/pp_api_test/test_processor.cc +++ b/test/pp_api_test/test_processor.cc @@ -407,6 +407,120 @@ TEST(ProcessorTest, TestGemma4ImageProcessing) { EXPECT_EQ(nst_peek[0], 260) << "num_soft_tokens should be 260 for australia.jpg (HF reference)"; } +TEST(ProcessorTest, TestGemma4UnifiedImageProcessing) { + // gemma-4-12B "unified" (encoder-free) vision preprocessing. + // + // The unified model consumes 48px MERGED patches (patch_dim = 48*48*3 = 6912) + // directly, with no SigLIP encoder to pool 3x3 teacher patches. HuggingFace + // produces these via (16px patchify -> 3x3 patches_merge). That is provably + // identical to a direct 48px patchify, so the unified contract is generated by + // reusing Gemma4ImageTransform with patch_size=48, pooling_kernel_size=1. + // + // This test verifies (a) the 6912-dim contract and (b) that the reused op's + // top-left 48px patch matches the 16px teacher patch from the standard config. + const char* image_path[] = {"data/processor/australia.jpg"}; + + OrtxObjectPtr raw_images{}; + extError_t err = OrtxLoadImages(raw_images.ToBeAssigned(), image_path, 1, nullptr); + ASSERT_EQ(err, kOrtxOK); + + // --- unified config: 48px merged patches --- + OrtxObjectPtr processor; + err = OrtxCreateProcessor(processor.ToBeAssigned(), "data/models/gemma-4-unified/image_processor.json"); + if (err != kOrtxOK) { + std::cout << "Error: " << OrtxGetLastErrorMessage() << std::endl; + } + ASSERT_EQ(err, kOrtxOK); + + OrtxObjectPtr result; + err = OrtxImagePreProcess(processor.get(), raw_images.get(), result.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + + OrtxObjectPtr tensor; + err = OrtxTensorResultGetAt(result.get(), 0, tensor.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + + const float* pv_data{}; + const int64_t* shape{}; + size_t num_dims; + err = OrtxGetTensorData(tensor.get(), reinterpret_cast(&pv_data), &shape, &num_dims); + ASSERT_EQ(err, kOrtxOK); + ASSERT_EQ(num_dims, 3ULL); + ASSERT_EQ(shape[0], 1); + constexpr int64_t kMaxSoftTokens = 280; + constexpr int64_t kMergedPatchDim = 48 * 48 * 3; // 6912 + ASSERT_EQ(shape[1], kMaxSoftTokens); + ASSERT_EQ(shape[2], kMergedPatchDim); + + // position_ids — merged 48-grid coordinates. + err = OrtxTensorResultGetAt(result.get(), 1, tensor.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + const int64_t* pos_data{}; + err = OrtxGetTensorData(tensor.get(), reinterpret_cast(&pos_data), &shape, &num_dims); + ASSERT_EQ(err, kOrtxOK); + EXPECT_EQ(pos_data[0], 0); // patch 0 x + EXPECT_EQ(pos_data[1], 0); // patch 0 y + + // num_soft_tokens: merged count = teacher-grid / 9. For australia.jpg the + // teacher grid is 60x39 -> merged 20x13 = 260 (same value as the standard + // config, which reports it before pooling). + OrtxObjectPtr nst_tensor; + err = OrtxTensorResultGetAt(result.get(), 2, nst_tensor.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + const int64_t* nst{}; + const int64_t* nst_shape{}; + size_t nst_dims{}; + err = OrtxGetTensorData(nst_tensor.get(), reinterpret_cast(&nst), &nst_shape, &nst_dims); + ASSERT_EQ(err, kOrtxOK); + EXPECT_EQ(nst[0], 260) << "merged soft-token count for australia.jpg (HF reference)"; + const int64_t num_merged = nst[0]; + // Last real merged patch position: (20-1, 13-1) = (19, 12). + EXPECT_EQ(pos_data[(num_merged - 1) * 2], 19); + EXPECT_EQ(pos_data[(num_merged - 1) * 2 + 1], 12); + // Padding beyond real patches is (-1, -1). + for (int64_t i = num_merged; i < kMaxSoftTokens; ++i) { + EXPECT_EQ(pos_data[i * 2], -1); + EXPECT_EQ(pos_data[i * 2 + 1], -1); + } + + // Copy merged patch 0 before running the second (teacher) config, which + // invalidates pv_data. + std::vector merged_patch0(pv_data, pv_data + kMergedPatchDim); + + // --- standard config: 16px teacher patches --- + OrtxObjectPtr teacher_proc; + err = OrtxCreateProcessor(teacher_proc.ToBeAssigned(), "data/models/gemma-4/image_processor.json"); + ASSERT_EQ(err, kOrtxOK); + OrtxObjectPtr teacher_result; + err = OrtxImagePreProcess(teacher_proc.get(), raw_images.get(), teacher_result.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + OrtxObjectPtr teacher_tensor; + err = OrtxTensorResultGetAt(teacher_result.get(), 0, teacher_tensor.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + const float* teacher_pv{}; + const int64_t* teacher_shape{}; + size_t teacher_dims{}; + err = OrtxGetTensorData(teacher_tensor.get(), reinterpret_cast(&teacher_pv), &teacher_shape, + &teacher_dims); + ASSERT_EQ(err, kOrtxOK); + + // Merged patch 0 is the top-left 48x48 image block, HWC. Its top-left 16x16 + // sub-block must equal teacher patch 0 (also the top-left 16x16 block, HWC). + // merged flat index: ((r*48 + col)*3 + c) for r,col in [0,16) + // teacher flat index: ((r*16 + col)*3 + c) + constexpr int64_t kTeacherPatchDim = 16 * 16 * 3; // 768 + for (int64_t r = 0; r < 16; ++r) { + for (int64_t col = 0; col < 16; ++col) { + for (int64_t c = 0; c < 3; ++c) { + float m = merged_patch0[(r * 48 + col) * 3 + c]; + float t = teacher_pv[(r * 16 + col) * 3 + c]; + ASSERT_NEAR(m, t, 1e-6f) << "merged vs teacher mismatch at r=" << r << " col=" << col << " c=" << c; + } + } + } + (void)kTeacherPatchDim; +} + TEST(ProcessorTest, TestGemma4ImageProcessingMultiImage) { // Verify batched processing works with multiple images of different sizes. const char* image_paths[] = {"data/processor/standard_s.jpg", "data/processor/australia.jpg"}; From ba71805cbfbf6e2a228992513fce8cb0456a1457 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 9 Jul 2026 22:34:32 +0000 Subject: [PATCH 2/5] Address review: sampling_rate validation and test hardening * Gemma4UnifiedAudioFrames: document that the op frames by sample count and does not resample (so sampling_rate is informational, matching the decoder's expected output rate), and reject non-positive sampling_rate. * Unified image test: assert the teacher tensor's rank/shape (using kTeacherPatchDim) before indexing to avoid a potential OOB read, and reword the pv_data snapshot comment to describe tensor-buffer rebinding rather than invalidation. Signed-off-by: Justin Chu --- shared/api/gemma4_audio_features.hpp | 12 ++++++++++++ test/pp_api_test/test_processor.cc | 10 +++++++--- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/shared/api/gemma4_audio_features.hpp b/shared/api/gemma4_audio_features.hpp index d26374ce4..2b380eba3 100644 --- a/shared/api/gemma4_audio_features.hpp +++ b/shared/api/gemma4_audio_features.hpp @@ -410,11 +410,23 @@ class Gemma4UnifiedAudioFrames { return {kOrtxErrorInvalidArgument, "[Gemma4UnifiedAudioFrames]: audio_samples_per_token must be positive"}; } + // The op frames whatever PCM the upstream AudioDecoder produces; the frame + // size is defined in samples, so this op is intrinsically sample-rate + // agnostic. ``sampling_rate`` therefore only documents the rate the decoder + // is expected to output (16 kHz for the gemma-4 contract, where 640 samples + // == 40 ms). Reject non-positive values so a misconfiguration is loud + // rather than silently producing frames at an unintended rate. + if (sampling_rate_ <= 0) { + return {kOrtxErrorInvalidArgument, + "[Gemma4UnifiedAudioFrames]: sampling_rate must be positive"}; + } return {}; } private: int64_t audio_samples_per_token_ = 640; // 640 samples = 40 ms @ 16 kHz + // Expected decoder output rate. Informational only: the op frames by sample + // count and does not resample (see the note in Init()). int64_t sampling_rate_ = 16000; float padding_value_ = 0.0f; }; diff --git a/test/pp_api_test/test_processor.cc b/test/pp_api_test/test_processor.cc index b48f89d03..1859f79f5 100644 --- a/test/pp_api_test/test_processor.cc +++ b/test/pp_api_test/test_processor.cc @@ -483,8 +483,9 @@ TEST(ProcessorTest, TestGemma4UnifiedImageProcessing) { EXPECT_EQ(pos_data[i * 2 + 1], -1); } - // Copy merged patch 0 before running the second (teacher) config, which - // invalidates pv_data. + // Copy merged patch 0 out now: the tensor's backing buffer is owned by the + // first result, and reusing `tensor`/running the teacher config below would + // rebind pv_data, so snapshot it into an independent buffer first. std::vector merged_patch0(pv_data, pv_data + kMergedPatchDim); // --- standard config: 16px teacher patches --- @@ -509,6 +510,10 @@ TEST(ProcessorTest, TestGemma4UnifiedImageProcessing) { // merged flat index: ((r*48 + col)*3 + c) for r,col in [0,16) // teacher flat index: ((r*16 + col)*3 + c) constexpr int64_t kTeacherPatchDim = 16 * 16 * 3; // 768 + // Guard the raw indexing below against an unexpected teacher tensor layout. + ASSERT_EQ(teacher_dims, 3ULL); // (batch, num_patches, patch_dim) + ASSERT_GE(teacher_shape[1], 1); // at least patch 0 + ASSERT_EQ(teacher_shape[2], kTeacherPatchDim); for (int64_t r = 0; r < 16; ++r) { for (int64_t col = 0; col < 16; ++col) { for (int64_t c = 0; c < 3; ++c) { @@ -518,7 +523,6 @@ TEST(ProcessorTest, TestGemma4UnifiedImageProcessing) { } } } - (void)kTeacherPatchDim; } TEST(ProcessorTest, TestGemma4ImageProcessingMultiImage) { From 2f4280015f79638665bd19f7362cf5dafbfb2870 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 10 Jul 2026 04:10:04 +0000 Subject: [PATCH 3/5] Consolidate gemma-4 audio into a single Gemma4Audio op (review) Per review feedback, instead of shipping a second kernel (Gemma4UnifiedAudioFrames) alongside Gemma4LogMel, expose one generic Gemma4Audio op that selects the front-end via a 'type' attribute: type="log_mel" (default) -> 128-dim USM log-mel spectrogram (E2B/E4B) type="raw_frames" -> raw 640-sample waveform frames (12B unified) Both branches now share the (features: float, mask: bool) output signature; the raw-frames path emits an all-true frame mask, matching HuggingFace Gemma4UnifiedAudioFeatureExtractor (input_features + input_features_mask). Gemma4Audio internally delegates to the unchanged Gemma4LogMel DSP and the raw-frames implementation. Gemma4LogMel stays registered as a backward-compat alias so already-deployed gemma-4 configs keep loading. The repo's own gemma-4 and gemma-4-unified configs now use Gemma4Audio. All 9 Gemma4 tests pass (log-mel via Gemma4Audio type=log_mel, raw frames via type=raw_frames incl. mask assertions). Signed-off-by: Justin Chu --- shared/api/gemma4_audio_features.hpp | 78 +++++++++++++++++-- shared/api/speech_extractor.cc | 2 +- .../audio_feature_extraction.json | 5 +- .../gemma-4/audio_feature_extraction.json | 5 +- test/pp_api_test/test_feature_extraction.cc | 24 +++++- 5 files changed, 101 insertions(+), 13 deletions(-) diff --git a/shared/api/gemma4_audio_features.hpp b/shared/api/gemma4_audio_features.hpp index 2b380eba3..a0685e7b0 100644 --- a/shared/api/gemma4_audio_features.hpp +++ b/shared/api/gemma4_audio_features.hpp @@ -294,6 +294,9 @@ class Gemma4LogMel { } else if (key == "per_bin_stddev") { auto& v = std::get>(value); per_bin_stddev_.assign(v.begin(), v.end()); + } else if (key == "type") { + // Consumed by the Gemma4Audio dispatcher (selects this log-mel path); + // ignored here so a forwarded attribute dict does not error. } else { return {kOrtxErrorInvalidArgument, "[Gemma4LogMel]: unknown attribute '" + key + "'"}; @@ -353,20 +356,24 @@ class Gemma4LogMel { // zero-pad the waveform to a multiple of ``audio_samples_per_token`` and // reshape it into ``(num_tokens, audio_samples_per_token)`` frames. // -// Pipeline: AudioDecoder -> Gemma4UnifiedAudioFrames +// Pipeline: AudioDecoder -> Gemma4Audio (type="raw_frames") // // Inputs: float (1, num_samples) — mono PCM at `sampling_rate` Hz // Outputs: float (num_tokens, audio_samples_per_token) — raw waveform frames +// bool (num_tokens,) — frame-level mask (all true) // -// The frame-level mask is intentionally not emitted here: for the single-clip -// inference path the downstream processor fills an all-true mask, and the -// batch framework zero-pads ragged clips when stacking. +// The mask is emitted (all-true) so that this path shares the (features, mask) +// output signature of Gemma4LogMel, letting a single Gemma4Audio op cover both. +// It matches HuggingFace ``Gemma4UnifiedAudioFeatureExtractor``, which returns +// ``input_features`` and ``input_features_mask``; ragged clips are zero-padded +// by the batch framework when stacking, so per-clip frames are all valid. class Gemma4UnifiedAudioFrames { public: Gemma4UnifiedAudioFrames() = default; OrtxStatus Compute(const ortc::Tensor& pcm_input, - ortc::Tensor& frames_out) { + ortc::Tensor& frames_out, + ortc::Tensor& mask_out) { const auto& pcm_shape = pcm_input.Shape(); if (pcm_shape.size() != 2 || pcm_shape[0] != 1) { return {kOrtxErrorInvalidArgument, @@ -380,6 +387,7 @@ class Gemma4UnifiedAudioFrames { const int64_t num_tokens = (num_samples + spt - 1) / spt; float* out = frames_out.Allocate({num_tokens, spt}); + bool* mask = mask_out.Allocate({num_tokens}); if (num_tokens == 0) { return {}; } @@ -387,6 +395,8 @@ class Gemma4UnifiedAudioFrames { // then copy the real samples over the front. std::fill(out, out + static_cast(num_tokens) * spt, padding_value_); std::copy(pcm_input.Data(), pcm_input.Data() + num_samples, out); + // Every frame of a single clip is valid (padding lives within the last frame). + std::fill(mask, mask + num_tokens, true); return {}; } @@ -401,6 +411,8 @@ class Gemma4UnifiedAudioFrames { sampling_rate_ = std::get(value); } else if (key == "padding_value") { padding_value_ = static_cast(std::get(value)); + } else if (key == "type") { + // Consumed by the Gemma4Audio dispatcher (selects this raw-frames path). } else { return {kOrtxErrorInvalidArgument, "[Gemma4UnifiedAudioFrames]: unknown attribute '" + key + "'"}; @@ -431,4 +443,60 @@ class Gemma4UnifiedAudioFrames { float padding_value_ = 0.0f; }; +// Unified Gemma 4 audio feature extraction op. +// +// A single registered op that dispatches, via the ``type`` attribute, to one of +// the gemma-4 audio front-ends rather than exposing a separate kernel per model +// variant: +// +// type = "log_mel" (default) -> 128-dim USM log-mel spectrogram (E2B/E4B) +// type = "raw_frames" -> raw 640-sample waveform frames (12B unified) +// +// Both branches share the (features: float, mask: bool) output signature. +// +// Pipeline: AudioDecoder -> Gemma4Audio +class Gemma4Audio { + public: + Gemma4Audio() = default; + + OrtxStatus Compute(const ortc::Tensor& pcm_input, + ortc::Tensor& features_out, + ortc::Tensor& mask_out) { + if (mode_ == Mode::kRawFrames) { + return raw_frames_.Compute(pcm_input, features_out, mask_out); + } + return log_mel_.Compute(pcm_input, features_out, mask_out); + } + + template + OrtxStatus Init(const DictT& attrs) { + // Select the front-end from the ``type`` attribute, then forward the full + // attribute dict to the chosen implementation (each ignores the ``type`` + // key). Defaults to log-mel for backward compatibility. + for (const auto& [key, value] : attrs) { + if (key == "type") { + const std::string& type = std::get(value); + if (type == "raw_frames") { + mode_ = Mode::kRawFrames; + } else if (type == "log_mel") { + mode_ = Mode::kLogMel; + } else { + return {kOrtxErrorInvalidArgument, + "[Gemma4Audio]: unknown type '" + type + "' (expected 'log_mel' or 'raw_frames')"}; + } + } + } + if (mode_ == Mode::kRawFrames) { + return raw_frames_.Init(attrs); + } + return log_mel_.Init(attrs); + } + + private: + enum class Mode { kLogMel, kRawFrames }; + Mode mode_ = Mode::kLogMel; + Gemma4LogMel log_mel_; + Gemma4UnifiedAudioFrames raw_frames_; +}; + } // namespace ort_extensions diff --git a/shared/api/speech_extractor.cc b/shared/api/speech_extractor.cc index 97679f172..ffd3c6960 100644 --- a/shared/api/speech_extractor.cc +++ b/shared/api/speech_extractor.cc @@ -18,7 +18,7 @@ Operation::KernelRegistry SpeechFeatureExtractor::kernel_registry_ = { {"PerFeatureNormalize", []() { return CreateKernelInstance(&PerFeatureNormalize::Compute); }}, {"Phi4AudioEmbed", []() { return CreateKernelInstance(&Phi4AudioEmbed::Compute); }}, {"Gemma4LogMel", []() { return CreateKernelInstance(&Gemma4LogMel::Compute); }}, - {"Gemma4UnifiedAudioFrames", []() { return CreateKernelInstance(&Gemma4UnifiedAudioFrames::Compute); }}}; + {"Gemma4Audio", []() { return CreateKernelInstance(&Gemma4Audio::Compute); }}}; SpeechFeatureExtractor::SpeechFeatureExtractor() : OrtxObjectImpl(extObjectKind_t::kOrtxKindFeatureExtractor) {} diff --git a/test/data/models/gemma-4-unified/audio_feature_extraction.json b/test/data/models/gemma-4-unified/audio_feature_extraction.json index abf92ed7d..db1de383e 100644 --- a/test/data/models/gemma-4-unified/audio_feature_extraction.json +++ b/test/data/models/gemma-4-unified/audio_feature_extraction.json @@ -9,9 +9,10 @@ }, { "operation": { - "name": "gemma4_unified_audio_frames", - "type": "Gemma4UnifiedAudioFrames", + "name": "gemma4_audio", + "type": "Gemma4Audio", "attrs": { + "type": "raw_frames", "audio_samples_per_token": 640, "sampling_rate": 16000, "padding_value": 0.0 diff --git a/test/data/models/gemma-4/audio_feature_extraction.json b/test/data/models/gemma-4/audio_feature_extraction.json index 699156dc4..b070f01cc 100644 --- a/test/data/models/gemma-4/audio_feature_extraction.json +++ b/test/data/models/gemma-4/audio_feature_extraction.json @@ -9,9 +9,10 @@ }, { "operation": { - "name": "gemma4_log_mel", - "type": "Gemma4LogMel", + "name": "gemma4_audio", + "type": "Gemma4Audio", "attrs": { + "type": "log_mel", "feature_size": 128, "sampling_rate": 16000, "frame_length_ms": 20.0, diff --git a/test/pp_api_test/test_feature_extraction.cc b/test/pp_api_test/test_feature_extraction.cc index e4ec78696..f155a453f 100644 --- a/test/pp_api_test/test_feature_extraction.cc +++ b/test/pp_api_test/test_feature_extraction.cc @@ -344,7 +344,7 @@ TEST(ExtractorTest, TestSplitSignalSegments) { TEST(ExtractorTest, TestGemma4AudioFeatureExtraction) { // Use existing test audio files to verify the Gemma 4 USM-style log-mel pipeline: - // AudioDecoder -> Gemma4LogMel + // AudioDecoder -> Gemma4Audio (type="log_mel") const char* audio_path[] = {"data/jfk.flac"}; OrtxObjectPtr raw_audios; extError_t err = OrtxLoadAudios(raw_audios.ToBeAssigned(), audio_path, 1); @@ -450,7 +450,8 @@ TEST(ExtractorTest, TestGemma4AudioFeatureExtractionMultiFile) { TEST(ExtractorTest, TestGemma4UnifiedAudioFrames) { // gemma-4-12B "unified" (encoder-free) audio: raw 16 kHz waveform chunked - // into fixed 640-sample frames. Pipeline: AudioDecoder -> Gemma4UnifiedAudioFrames + // into fixed 640-sample frames via the generic Gemma4Audio op with + // type="raw_frames". Pipeline: AudioDecoder -> Gemma4Audio const char* audio_path[] = {"data/jfk.flac"}; OrtxObjectPtr raw_audios; extError_t err = OrtxLoadAudios(raw_audios.ToBeAssigned(), audio_path, 1); @@ -476,10 +477,27 @@ TEST(ExtractorTest, TestGemma4UnifiedAudioFrames) { ASSERT_EQ(shape[0], 1); // single audio ASSERT_EQ(shape[2], 640); // 640 raw samples per token EXPECT_GT(shape[1], 0); // at least one frame + const int64_t num_tokens = shape[1]; // All values finite and within the normalized PCM range. - for (int64_t i = 0; i < std::min(shape[1] * 640, 5000); ++i) { + for (int64_t i = 0; i < std::min(num_tokens * 640, 5000); ++i) { ASSERT_TRUE(std::isfinite(data[i])) << "frame value at index " << i << " is not finite"; ASSERT_LE(std::abs(data[i]), 4.0f) << "frame value at index " << i << " out of range"; } + + // Output 1: frame mask — bool (batch, num_tokens), all true for a single clip. + err = OrtxTensorResultGetAt(result.get(), 1, tensor.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + const bool* mask_data{}; + const int64_t* mask_shape{}; + size_t mask_dims; + err = OrtxGetTensorData(tensor.get(), reinterpret_cast(&mask_data), &mask_shape, &mask_dims); + ASSERT_EQ(err, kOrtxOK); + ASSERT_EQ(mask_dims, 2ULL); // (batch, num_tokens) + ASSERT_EQ(mask_shape[0], 1); + ASSERT_EQ(mask_shape[1], num_tokens); // same frame count as features + for (int64_t i = 0; i < num_tokens; ++i) { + EXPECT_TRUE(mask_data[i]) << "single-clip frame " << i << " should be valid"; + } } + From 920f54b9cf8142577ba3da48c471d89bbbd4a59d Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 10 Jul 2026 23:00:08 +0000 Subject: [PATCH 4/5] Address review: checked attr extraction, alias conflict, multi-file test Follow-up to the Gemma4Audio consolidation: * Use checked std::get_if via a GetTypedAttr helper across Gemma4LogMel, Gemma4UnifiedAudioFrames and Gemma4Audio Init so a malformed config value (wrong JSON type) returns kOrtxErrorInvalidArgument instead of throwing std::bad_variant_access past the OrtxStatus boundary. * Reject conflicting 'audio_samples_per_token' vs 'feature_size' aliases when both are provided with different values, rather than silently taking the last. * Add TestGemma4UnifiedAudioFramesMultiFile: two clips of different lengths, asserting batch stacking pads the shorter clip and the frame mask marks the padded tail invalid (true-prefix / false-suffix per row). All Gemma4 tests pass. Signed-off-by: Justin Chu --- shared/api/gemma4_audio_features.hpp | 96 ++++++++++++++++----- test/pp_api_test/test_feature_extraction.cc | 67 ++++++++++++++ 2 files changed, 143 insertions(+), 20 deletions(-) diff --git a/shared/api/gemma4_audio_features.hpp b/shared/api/gemma4_audio_features.hpp index a0685e7b0..a3c3743bc 100644 --- a/shared/api/gemma4_audio_features.hpp +++ b/shared/api/gemma4_audio_features.hpp @@ -18,10 +18,28 @@ namespace ort_extensions { +namespace gemma4_audio_detail { + +// Checked attribute extraction: returns kOrtxErrorInvalidArgument instead of +// throwing std::bad_variant_access when a config value has an unexpected type +// (e.g. a string where a number is required, or an int where a float is). +template +OrtxStatus GetTypedAttr(const VariantT& value, const char* op_name, const std::string& key, T& out) { + const T* ptr = std::get_if(&value); + if (ptr == nullptr) { + return {kOrtxErrorInvalidArgument, + std::string("[") + op_name + "]: attribute '" + key + "' has an unexpected value type"}; + } + out = *ptr; + return {}; +} + +} // namespace gemma4_audio_detail + // Gemma 4 audio feature extraction: USM-style log-mel spectrogram that matches // the HuggingFace Gemma4AudioFeatureExtractor exactly. // -// Pipeline: AudioDecoder -> Gemma4LogMel +// Pipeline: AudioDecoder -> Gemma4Audio (type="log_mel") // // Inputs: float (1, num_samples) — mono PCM at `sampling_rate` Hz // Outputs: float (num_frames, feature_size) — log-mel features @@ -267,33 +285,45 @@ class Gemma4LogMel { template OrtxStatus Init(const DictT& attrs) { + using gemma4_audio_detail::GetTypedAttr; + constexpr const char* kOp = "Gemma4LogMel"; for (const auto& [key, value] : attrs) { if (key == "feature_size") { - feature_size_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, feature_size_); !st.IsOk()) return st; } else if (key == "sampling_rate") { - sampling_rate_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, sampling_rate_); !st.IsOk()) return st; } else if (key == "frame_length_ms") { - frame_length_ms_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, frame_length_ms_); !st.IsOk()) return st; } else if (key == "hop_length_ms") { - hop_length_ms_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, hop_length_ms_); !st.IsOk()) return st; } else if (key == "min_frequency") { - min_frequency_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, min_frequency_); !st.IsOk()) return st; } else if (key == "max_frequency") { - max_frequency_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, max_frequency_); !st.IsOk()) return st; } else if (key == "preemphasis") { - preemphasis_ = static_cast(std::get(value)); + double tmp = 0.0; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + preemphasis_ = static_cast(tmp); } else if (key == "preemphasis_htk_flavor") { - preemphasis_htk_flavor_ = std::get(value) != 0; + int64_t tmp = 0; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + preemphasis_htk_flavor_ = tmp != 0; } else if (key == "fft_overdrive") { - fft_overdrive_ = std::get(value) != 0; + int64_t tmp = 0; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + fft_overdrive_ = tmp != 0; } else if (key == "mel_floor") { - mel_floor_ = static_cast(std::get(value)); + double tmp = 0.0; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + mel_floor_ = static_cast(tmp); } else if (key == "per_bin_mean") { - auto& v = std::get>(value); - per_bin_mean_.assign(v.begin(), v.end()); + std::vector tmp; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + per_bin_mean_.assign(tmp.begin(), tmp.end()); } else if (key == "per_bin_stddev") { - auto& v = std::get>(value); - per_bin_stddev_.assign(v.begin(), v.end()); + std::vector tmp; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + per_bin_stddev_.assign(tmp.begin(), tmp.end()); } else if (key == "type") { // Consumed by the Gemma4Audio dispatcher (selects this log-mel path); // ignored here so a forwarded attribute dict does not error. @@ -402,15 +432,27 @@ class Gemma4UnifiedAudioFrames { template OrtxStatus Init(const DictT& attrs) { + using gemma4_audio_detail::GetTypedAttr; + constexpr const char* kOp = "Gemma4UnifiedAudioFrames"; + // Track the two aliases separately so conflicting values are rejected rather + // than silently taking whichever key appears last. + bool samples_set = false, feature_size_set = false; + int64_t samples_val = 0, feature_size_val = 0; for (const auto& [key, value] : attrs) { - if (key == "audio_samples_per_token" || key == "feature_size") { + if (key == "audio_samples_per_token") { + if (auto st = GetTypedAttr(value, kOp, key, samples_val); !st.IsOk()) return st; + samples_set = true; + } else if (key == "feature_size") { // ``feature_size`` is accepted as an alias: HF sets feature_size == // audio_samples_per_token (both default to 640). - audio_samples_per_token_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, feature_size_val); !st.IsOk()) return st; + feature_size_set = true; } else if (key == "sampling_rate") { - sampling_rate_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, sampling_rate_); !st.IsOk()) return st; } else if (key == "padding_value") { - padding_value_ = static_cast(std::get(value)); + double tmp = 0.0; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + padding_value_ = static_cast(tmp); } else if (key == "type") { // Consumed by the Gemma4Audio dispatcher (selects this raw-frames path). } else { @@ -418,6 +460,17 @@ class Gemma4UnifiedAudioFrames { "[Gemma4UnifiedAudioFrames]: unknown attribute '" + key + "'"}; } } + if (samples_set && feature_size_set && samples_val != feature_size_val) { + return {kOrtxErrorInvalidArgument, + "[Gemma4UnifiedAudioFrames]: conflicting 'audio_samples_per_token' (" + + std::to_string(samples_val) + ") and 'feature_size' (" + std::to_string(feature_size_val) + + "); they are aliases and must match"}; + } + if (samples_set) { + audio_samples_per_token_ = samples_val; + } else if (feature_size_set) { + audio_samples_per_token_ = feature_size_val; + } if (audio_samples_per_token_ <= 0) { return {kOrtxErrorInvalidArgument, "[Gemma4UnifiedAudioFrames]: audio_samples_per_token must be positive"}; @@ -475,7 +528,10 @@ class Gemma4Audio { // key). Defaults to log-mel for backward compatibility. for (const auto& [key, value] : attrs) { if (key == "type") { - const std::string& type = std::get(value); + std::string type; + if (auto st = gemma4_audio_detail::GetTypedAttr(value, "Gemma4Audio", key, type); !st.IsOk()) { + return st; + } if (type == "raw_frames") { mode_ = Mode::kRawFrames; } else if (type == "log_mel") { diff --git a/test/pp_api_test/test_feature_extraction.cc b/test/pp_api_test/test_feature_extraction.cc index f155a453f..5cdcb7f61 100644 --- a/test/pp_api_test/test_feature_extraction.cc +++ b/test/pp_api_test/test_feature_extraction.cc @@ -501,3 +501,70 @@ TEST(ExtractorTest, TestGemma4UnifiedAudioFrames) { } } +TEST(ExtractorTest, TestGemma4UnifiedAudioFramesMultiFile) { + // Two clips of different lengths: verify batch stacking pads the shorter clip's + // frames and that the frame mask marks the padded tail invalid (false), while + // the real frames of each clip are valid (true). Locks in the unified batch + + // mask behavior, mirroring the log-mel multi-file coverage. + const char* audio_path[] = {"data/jfk.flac", "data/1272-141231-0002.wav"}; + OrtxObjectPtr raw_audios; + extError_t err = OrtxLoadAudios(raw_audios.ToBeAssigned(), audio_path, 2); + ASSERT_EQ(err, kOrtxOK); + + OrtxObjectPtr feature_extractor( + OrtxCreateSpeechFeatureExtractor, "data/models/gemma-4-unified/audio_feature_extraction.json"); + OrtxObjectPtr result; + err = OrtxFeatureExtraction(feature_extractor.get(), raw_audios.get(), result.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + + // Output 0: frames — float (2, max_tokens, 640) + OrtxObjectPtr tensor; + err = OrtxTensorResultGetAt(result.get(), 0, tensor.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + const float* data{}; + const int64_t* shape{}; + size_t num_dims; + err = OrtxGetTensorData(tensor.get(), reinterpret_cast(&data), &shape, &num_dims); + ASSERT_EQ(err, kOrtxOK); + ASSERT_EQ(num_dims, 3ULL); + ASSERT_EQ(shape[0], 2); // batch of 2 clips + ASSERT_EQ(shape[2], 640); // raw samples per token + const int64_t max_tokens = shape[1]; + + // Output 1: mask — bool (2, max_tokens) + err = OrtxTensorResultGetAt(result.get(), 1, tensor.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + const bool* mask_data{}; + const int64_t* mask_shape{}; + size_t mask_dims; + err = OrtxGetTensorData(tensor.get(), reinterpret_cast(&mask_data), &mask_shape, &mask_dims); + ASSERT_EQ(err, kOrtxOK); + ASSERT_EQ(mask_dims, 2ULL); + ASSERT_EQ(mask_shape[0], 2); + ASSERT_EQ(mask_shape[1], max_tokens); + + // Each row's mask must be a contiguous true-prefix (real frames) followed by a + // false-suffix (batch padding). Count valid frames per clip. + int64_t valid_counts[2] = {0, 0}; + for (int64_t b = 0; b < 2; ++b) { + const bool* row = mask_data + b * max_tokens; + bool seen_false = false; + for (int64_t i = 0; i < max_tokens; ++i) { + if (row[i]) { + ASSERT_FALSE(seen_false) << "clip " << b << " mask must not have a true frame after padding"; + ++valid_counts[b]; + } else { + seen_false = true; + } + } + EXPECT_GT(valid_counts[b], 0) << "clip " << b << " should have at least one valid frame"; + } + + // The two clips have different lengths, so exactly one clip fills all max_tokens + // and the shorter clip has a padded (false) tail. + EXPECT_NE(valid_counts[0], valid_counts[1]) << "test clips should differ in length"; + EXPECT_EQ(std::max(valid_counts[0], valid_counts[1]), max_tokens); + const int64_t shorter = std::min(valid_counts[0], valid_counts[1]); + EXPECT_LT(shorter, max_tokens) << "shorter clip should be zero-padded in the batch"; +} + From ff3d01b10b370a894d45a8a5091117cbf426a588 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 10 Jul 2026 23:20:23 +0000 Subject: [PATCH 5/5] Address Copilot review: test shape guards, PCM bound, doc clarity * Unified image test: assert position_ids rank/shape and that num_soft_tokens is within [1, max_soft_tokens] before indexing pos_data, so a changed output layout fails cleanly instead of reading out of bounds. * Unified raw-frames test: tighten the amplitude bound to the normalized PCM range (<= 1.0001, epsilon for full-scale rounding) so it matches the comment. * Gemma4LogMel doc block: clarify it is the log-mel implementation, registered under its own name for back-compat and reused by Gemma4Audio type="log_mel" (was ambiguously labeled as the Gemma4Audio pipeline). All 10 Gemma4 tests pass. Signed-off-by: Justin Chu --- shared/api/gemma4_audio_features.hpp | 6 +++++- test/pp_api_test/test_feature_extraction.cc | 5 +++-- test/pp_api_test/test_processor.cc | 10 ++++++++++ 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/shared/api/gemma4_audio_features.hpp b/shared/api/gemma4_audio_features.hpp index a3c3743bc..cf72770cc 100644 --- a/shared/api/gemma4_audio_features.hpp +++ b/shared/api/gemma4_audio_features.hpp @@ -39,7 +39,11 @@ OrtxStatus GetTypedAttr(const VariantT& value, const char* op_name, const std::s // Gemma 4 audio feature extraction: USM-style log-mel spectrogram that matches // the HuggingFace Gemma4AudioFeatureExtractor exactly. // -// Pipeline: AudioDecoder -> Gemma4Audio (type="log_mel") +// This is the log-mel implementation. It stays registered under its own +// "Gemma4LogMel" name (backward compatibility) and is also used internally by +// the generic Gemma4Audio op with type="log_mel". +// +// Pipeline: AudioDecoder -> Gemma4LogMel (or Gemma4Audio type="log_mel") // // Inputs: float (1, num_samples) — mono PCM at `sampling_rate` Hz // Outputs: float (num_frames, feature_size) — log-mel features diff --git a/test/pp_api_test/test_feature_extraction.cc b/test/pp_api_test/test_feature_extraction.cc index 5cdcb7f61..fb694b847 100644 --- a/test/pp_api_test/test_feature_extraction.cc +++ b/test/pp_api_test/test_feature_extraction.cc @@ -479,10 +479,11 @@ TEST(ExtractorTest, TestGemma4UnifiedAudioFrames) { EXPECT_GT(shape[1], 0); // at least one frame const int64_t num_tokens = shape[1]; - // All values finite and within the normalized PCM range. + // Raw waveform frames are the decoded PCM samples, which the AudioDecoder + // normalizes to [-1, 1]; a small epsilon covers float rounding at full scale. for (int64_t i = 0; i < std::min(num_tokens * 640, 5000); ++i) { ASSERT_TRUE(std::isfinite(data[i])) << "frame value at index " << i << " is not finite"; - ASSERT_LE(std::abs(data[i]), 4.0f) << "frame value at index " << i << " out of range"; + ASSERT_LE(std::abs(data[i]), 1.0001f) << "frame value at index " << i << " out of normalized PCM range"; } // Output 1: frame mask — bool (batch, num_tokens), all true for a single clip. diff --git a/test/pp_api_test/test_processor.cc b/test/pp_api_test/test_processor.cc index 1859f79f5..52c4dd390 100644 --- a/test/pp_api_test/test_processor.cc +++ b/test/pp_api_test/test_processor.cc @@ -458,6 +458,11 @@ TEST(ProcessorTest, TestGemma4UnifiedImageProcessing) { const int64_t* pos_data{}; err = OrtxGetTensorData(tensor.get(), reinterpret_cast(&pos_data), &shape, &num_dims); ASSERT_EQ(err, kOrtxOK); + // Guard the raw pos_data indexing below against an unexpected layout. + ASSERT_EQ(num_dims, 3ULL); // (batch, max_soft_tokens, 2) + ASSERT_EQ(shape[0], 1); + ASSERT_EQ(shape[1], kMaxSoftTokens); + ASSERT_EQ(shape[2], 2); EXPECT_EQ(pos_data[0], 0); // patch 0 x EXPECT_EQ(pos_data[1], 0); // patch 0 y @@ -472,8 +477,13 @@ TEST(ProcessorTest, TestGemma4UnifiedImageProcessing) { size_t nst_dims{}; err = OrtxGetTensorData(nst_tensor.get(), reinterpret_cast(&nst), &nst_shape, &nst_dims); ASSERT_EQ(err, kOrtxOK); + ASSERT_EQ(nst_dims, 2ULL); // (batch, 1) + ASSERT_EQ(nst_shape[0], 1); EXPECT_EQ(nst[0], 260) << "merged soft-token count for australia.jpg (HF reference)"; const int64_t num_merged = nst[0]; + // Value must be within the padded grid before it indexes pos_data below. + ASSERT_GT(num_merged, 0); + ASSERT_LE(num_merged, kMaxSoftTokens); // Last real merged patch position: (20-1, 13-1) = (19, 12). EXPECT_EQ(pos_data[(num_merged - 1) * 2], 19); EXPECT_EQ(pos_data[(num_merged - 1) * 2 + 1], 12);