chore: import upstream snapshot with attribution

tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/check_copyright_header.py

@@ -0,0 +1,120 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A script to check that copyright headers exists"""
+
+import argparse
+import re
+import sys
+from datetime import datetime
+from pathlib import Path
+
+EXCLUSIONS = ["scripts/get_commonvoice_data.py"]
+
+
+def get_top_comments(_data):
+    # Get all lines where comments should exist
+    lines_to_extract = []
+    for i, line in enumerate(_data):
+        # If empty line, skip
+        if line in ["", "\n", "", "\r", "\r\n"]:
+            continue
+        # If it is a comment line, we should get it
+        if line.startswith("#"):
+            lines_to_extract.append(i)
+        # Assume all copyright headers occur before any import statements not enclosed in a comment block
+        elif "import" in line:
+            break
+
+    comments = []
+    for line in lines_to_extract:
+        comments.append(_data[line])
+
+    return comments
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Usage for copyright header insertion script")
+    parser.add_argument(
+        '--dir',
+        help='Path to source files to add copyright header to. Will recurse through subdirectories',
+        required=True,
+        type=str,
+    )
+    args = parser.parse_args()
+
+    current_year = int(datetime.today().year)
+    starting_year = 2020
+    python_header_path = "tests/py_cprheader.txt"
+    with open(python_header_path, 'r', encoding='utf-8') as original:
+        pyheader = original.read().split("\n")
+        pyheader_lines = len(pyheader)
+
+    problematic_files = []
+
+    for filename in Path(args.dir).rglob('*.py'):
+        if str(filename) in EXCLUSIONS:
+            continue
+        with open(str(filename), 'r', encoding='utf-8') as original:
+            data = original.readlines()
+
+        data = get_top_comments(data)
+        if len(data) < pyheader_lines:
+            print(f"{filename} has less header lines than the copyright template")
+            problematic_files.append(filename)
+            continue
+
+        found = False
+        for i, line in enumerate(data):
+            if re.search(re.compile("Copyright.*NVIDIA.*", re.IGNORECASE), line):
+                # if re.search(re.compile("Copyright.*", re.IGNORECASE), line):
+                found = True
+                # Check 1st line manually
+                year_good = False
+                for year in range(starting_year, current_year + 1):
+                    year_line = pyheader[0].format(CURRENT_YEAR=year)
+                    if year_line in data[i]:
+                        year_good = True
+                        break
+                    year_line_aff = year_line.split(".")
+                    year_line_aff = year_line_aff[0] + " & AFFILIATES." + year_line_aff[1]
+                    if year_line_aff in data[i]:
+                        year_good = True
+                        break
+                if not year_good:
+                    problematic_files.append(filename)
+                    print(f"{filename} had an error with the year")
+                    break
+                while "opyright" in data[i]:
+                    i += 1
+                for j in range(1, pyheader_lines):
+                    if pyheader[j] not in data[i + j - 1]:
+                        problematic_files.append(filename)
+                        print(f"{filename} missed the line: {pyheader[j]}")
+                        break
+            if found:
+                break
+        if not found:
+            print(f"{filename} did not match the regex: `Copyright.*NVIDIA.*`")
+            problematic_files.append(filename)
+
+    if len(problematic_files) > 0:
+        print("check_copyright_headers.py found the following files that might not have a copyright header:")
+        for _file in problematic_files:
+            print(_file)
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+    main()

tests/collections/asr/confidence/test_asr_confidence.py

@@ -0,0 +1,170 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import math
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import pytest
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.models import ASRModel, EncDecMultiTaskModel
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig
+from nemo.collections.asr.parts.submodules.ctc_greedy_decoding import GreedyCTCInferConfig
+from nemo.collections.asr.parts.submodules.multitask_decoding import MultiTaskDecodingConfig
+from nemo.collections.asr.parts.submodules.multitask_greedy_decoding import AEDGreedyInferConfig
+from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig
+from nemo.collections.asr.parts.submodules.rnnt_greedy_decoding import GreedyBatchedRNNTInferConfig
+from nemo.collections.asr.parts.utils.asr_confidence_benchmarking_utils import run_confidence_benchmark
+from nemo.collections.asr.parts.utils.asr_confidence_utils import ConfidenceConfig
+
+# both models recognize the test data without errors, thus every metric except ece return default values
+# ECE values for fast conformer models (stt_en_fastconformer_ctc_large and stt_en_fastconformer_transducer_large)
+ECE_VALUES = {("token", "ctc"): 0.86, ("token", "rnnt"): 0.75, ("word", "ctc"): 0.89, ("word", "rnnt"): 0.80}
+
+TOL_DEGREE = 2
+TOL = 2 / math.pow(10, TOL_DEGREE)
+
+
+@pytest.fixture(scope="module")
+def audio_and_texts(test_data_dir):
+    # get filenames and reference texts from manifest
+    filepaths = []
+    reference_texts = []
+    manifest = Path(test_data_dir) / Path("asr/an4_val.json")
+    with open(manifest, 'r') as f:
+        for line in f:
+            item = json.loads(line)
+            # alaptev: maybe fix those paths in the manifest?
+            audio_file = Path(item['audio_filepath'].replace("/data/", "/.data/"))
+            filepaths.append(str(audio_file.absolute()))
+            reference_texts.append(item['text'])
+    return filepaths, reference_texts
+
+
+class TestASRConfidenceBenchmark:
+    @pytest.mark.integration
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize('model_name', ("ctc", "rnnt"))
+    @pytest.mark.parametrize('target_level', ("token", "word"))
+    def test_run_confidence_benchmark(
+        self, model_name, target_level, audio_and_texts, fast_conformer_ctc_model, fast_conformer_transducer_model
+    ):
+        model = fast_conformer_ctc_model if model_name == "ctc" else fast_conformer_transducer_model
+        assert isinstance(model, ASRModel)
+        filepaths, reference_texts = audio_and_texts
+        confidence_cfg = (
+            ConfidenceConfig(preserve_frame_confidence=True, preserve_token_confidence=True)
+            if target_level == "token"
+            else ConfidenceConfig(preserve_frame_confidence=True, preserve_word_confidence=True)
+        )
+        model.change_decoding_strategy(
+            RNNTDecodingConfig(
+                fused_batch_size=-1,
+                strategy="greedy_batch",
+                confidence_cfg=confidence_cfg,
+                greedy=GreedyBatchedRNNTInferConfig(loop_labels=False),
+            )
+            if model_name == "rnnt"
+            else CTCDecodingConfig(confidence_cfg=confidence_cfg)
+        )
+        with tempfile.TemporaryDirectory() as tmpdir:
+            assert np.allclose(
+                np.array(
+                    run_confidence_benchmark(model, target_level, filepaths, reference_texts, plot_dir=tmpdir)[
+                        target_level
+                    ]
+                ),
+                np.array([0.5, 1.0, 0.0, -math.inf, ECE_VALUES[(target_level, model_name)], 0.0, 0.0, 0.0]),
+                atol=TOL,
+            )
+
+    @pytest.mark.integration
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize('model_name', ("ctc", "rnnt"))
+    def test_deprecated_config_args(self, model_name, fast_conformer_ctc_model, fast_conformer_transducer_model):
+        assert ConfidenceConfig().method_cfg.alpha == 0.33, "default `alpha` is supposed to be 0.33"
+        model = fast_conformer_ctc_model if model_name == "ctc" else fast_conformer_transducer_model
+        assert isinstance(model, ASRModel)
+
+        conf = OmegaConf.create({"temperature": 0.5})
+        test_args_main = {"method_cfg": conf}
+        test_args_greedy = {"confidence_method_cfg": conf}
+        confidence_cfg = ConfidenceConfig(preserve_word_confidence=True, **test_args_main)
+        model.change_decoding_strategy(
+            RNNTDecodingConfig(fused_batch_size=-1, strategy="greedy", confidence_cfg=confidence_cfg)
+            if model_name == "rnnt"
+            else CTCDecodingConfig(confidence_cfg=confidence_cfg)
+        )
+        assert model.cfg.decoding.confidence_cfg.method_cfg.alpha == 0.5
+        model.change_decoding_strategy(
+            RNNTDecodingConfig(
+                fused_batch_size=-1,
+                strategy="greedy",
+                greedy=GreedyBatchedRNNTInferConfig(preserve_frame_confidence=True, **test_args_greedy),
+            )
+            if model_name == "rnnt"
+            else CTCDecodingConfig(greedy=GreedyCTCInferConfig(preserve_frame_confidence=True, **test_args_greedy))
+        )
+        assert model.cfg.decoding.greedy.confidence_method_cfg.alpha == 0.5
+
+    @pytest.mark.integration
+    @pytest.mark.with_downloads
+    def test_aed_multitask_model_confidence(self, canary_1b_v2, test_data_dir):
+        """Test token and word confidence for AED multitask models (Canary)."""
+        model = canary_1b_v2
+        assert isinstance(model, EncDecMultiTaskModel)
+
+        audio_file = Path(test_data_dir) / "asr" / "train" / "an4" / "wav" / "an46-mmap-b.wav"
+
+        # Configure decoding with confidence
+        decode_cfg = MultiTaskDecodingConfig(
+            strategy='greedy',
+            greedy=AEDGreedyInferConfig(preserve_token_confidence=True),
+            confidence_cfg=ConfidenceConfig(preserve_token_confidence=True, preserve_word_confidence=True),
+        )
+        model.change_decoding_strategy(decode_cfg)
+
+        hypotheses = model.transcribe(
+            audio=str(audio_file),
+            batch_size=1,
+            return_hypotheses=True,
+        )
+
+        assert len(hypotheses) == 1
+        hyp = hypotheses[0]
+
+        # Verify text is present
+        assert isinstance(hyp.text, str)
+        assert len(hyp.text) > 0
+
+        # Verify y_sequence is present
+        assert hyp.y_sequence is not None
+        assert len(hyp.y_sequence) > 0
+
+        # Verify token confidence is present and has correct length
+        assert hyp.token_confidence is not None
+        assert len(hyp.token_confidence) == len(hyp.y_sequence)
+
+        # Verify word confidence is present
+        assert hyp.word_confidence is not None
+        assert len(hyp.word_confidence) > 0
+
+        # Verify confidence values are in valid range [0, 1]
+        for conf in hyp.token_confidence:
+            assert 0.0 <= conf <= 1.0, f"Token confidence {conf} not in valid range [0, 1]"
+        for conf in hyp.word_confidence:
+            assert 0.0 <= conf <= 1.0, f"Word confidence {conf} not in valid range [0, 1]"

tests/collections/asr/confidence/test_asr_confidence_metrics.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import tempfile
+
+import numpy as np
+import pytest
+from scipy.stats import uniform
+
+from nemo.collections.asr.parts.utils.confidence_metrics import (
+    auc_nt,
+    auc_pr,
+    auc_roc,
+    auc_yc,
+    ece,
+    nce,
+    save_confidence_hist,
+    save_custom_confidence_curve,
+    save_nt_curve,
+    save_pr_curve,
+    save_roc_curve,
+)
+
+# set convenient name2metric mapping
+name2metric = {
+    f.__name__: (f, ans)
+    for f, ans in zip((auc_roc, auc_pr, auc_nt, auc_yc, ece, nce), (0.833, 0.917, 0.833, 0.421, 0.232, 0.403))
+}
+# ece does not have a default value
+name2metric_all_correct = {
+    f.__name__: (f, ans) for f, ans in zip((auc_roc, auc_pr, auc_nt, auc_yc, nce), (0.5, 1.0, 0.0, 0.0, -math.inf))
+}
+name2metric_all_incorrect = {
+    f.__name__: (f, ans) for f, ans in zip((auc_roc, auc_pr, auc_nt, auc_yc, nce), (0.5, 0.0, 1.0, 0.0, -math.inf))
+}
+
+# Initialize data
+Y_TRUE = [1, 0, 0, 1, 1]
+Y_TRUE_ALL_CORRECT = [1, 1, 1, 1, 1]
+Y_TRUE_ALL_INCORRECT = [0, 0, 0, 0, 0]
+Y_SCORE = [0.6, 0.7, 0.02, 0.95, 0.8]
+Y_TRUE_RANDOM = np.random.choice(2, 1000, p=[0.2, 0.8])
+# probability distribution with mean ~= 0.65 and std ~= 0.25
+Y_SCORE_RANDOM = uniform.rvs(size=1000, loc=0.5, scale=0.5) - 0.5 * np.random.choice(2, 1000, p=[0.8, 0.2])
+
+TOL_DEGREE = 3
+TOL = 1 / math.pow(10, TOL_DEGREE)
+
+
+class TestConfidenceMetrics:
+    @pytest.mark.unit
+    @pytest.mark.parametrize('metric_name', name2metric.keys())
+    def test_metric_main(self, metric_name):
+        metric, ans = name2metric[metric_name]
+
+        assert round(metric(Y_TRUE, Y_SCORE), TOL_DEGREE) == ans
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('metric_name', name2metric_all_correct.keys())
+    def test_metric_all_correct(self, metric_name):
+        metric, ans = name2metric_all_correct[metric_name]
+
+        assert round(metric(Y_TRUE_ALL_CORRECT, Y_SCORE), TOL_DEGREE) == ans
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('metric_name', name2metric_all_incorrect.keys())
+    def test_metric_all_incorrect(self, metric_name):
+        metric, ans = name2metric_all_incorrect[metric_name]
+
+        assert round(metric(Y_TRUE_ALL_INCORRECT, Y_SCORE), TOL_DEGREE) == ans
+
+    @pytest.mark.unit
+    def test_metric_auc_yc_aux(self):
+        n_bins = 10
+        result, result_std, result_max, (thresholds, yc_curve) = auc_yc(
+            Y_TRUE, Y_SCORE, n_bins=n_bins, return_std_maximum=True, return_curve=True
+        )
+
+        assert round(result_std, TOL_DEGREE) == 0.228
+        assert round(result_max, TOL_DEGREE) == 0.667
+        assert np.allclose(np.array(thresholds), np.array([i / n_bins for i in range(0, n_bins + 1)]), atol=TOL)
+        assert np.allclose(
+            np.array(yc_curve), np.array([0.0, 0.5, 0.5, 0.5, 0.5, 0.5, 0.167, 0.667, 0.667, 0.333, 0.0]), atol=TOL
+        )
+
+
+class TestSaveConfidencePlot:
+    @pytest.mark.unit
+    def test_save_confidence_hist(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            save_confidence_hist(Y_SCORE_RANDOM, tmpdir)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('plot_func', (save_roc_curve, save_pr_curve, save_nt_curve))
+    def test_save_simple_confidence_curve(self, plot_func):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            plot_func(Y_TRUE_RANDOM, Y_SCORE_RANDOM, tmpdir)
+
+    @pytest.mark.unit
+    def test_save_custom_confidence_curve(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            ranges = np.arange(0, 1, 0.01)
+            save_custom_confidence_curve(ranges, ranges, tmpdir)

tests/collections/asr/confidence/test_asr_confidence_primitives.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import pytest
+import torch
+
+from nemo.collections.asr.parts.utils.asr_confidence_utils import (
+    get_confidence_aggregation_bank,
+    get_confidence_measure_bank,
+)
+
+# Initialize probability vectors
+VOCAB_SIZES = (100, 1000, 10000)
+ONE_VEC_SET, ZERO_VEC_SET, RAND_VEC_SET, OVERFIT_RAND_VEC_SET = {}, {}, {}, {}
+for vocab_size in VOCAB_SIZES:
+    # batch size 2 to test different positions of probability one
+    ONE_VEC_SET[vocab_size] = torch.nan_to_num(
+        torch.cat(
+            [
+                torch.tensor([[0] + [float('-inf')] * (vocab_size - 1)]),
+                torch.tensor([[float('-inf')] * (vocab_size - 3) + [0] + [float('-inf')] * 2]),
+            ]
+        )
+    )
+    ZERO_VEC_SET[vocab_size] = torch.nan_to_num(torch.tensor([[math.log(1 / vocab_size)] * vocab_size] * 2))
+    # batch size 1
+    rand_logit = torch.rand((1, vocab_size))
+    rand_logit_overfit = rand_logit.clone()
+    rand_logit_overfit[0, 0] += vocab_size
+    RAND_VEC_SET[vocab_size] = torch.nan_to_num(torch.nn.functional.log_softmax(rand_logit, -1))
+    OVERFIT_RAND_VEC_SET[vocab_size] = torch.nan_to_num(torch.nn.functional.log_softmax(rand_logit_overfit, -1))
+AGGREGATION_VEC_SIMPLE = [0.0, 0.5, 1]
+
+TOL_DEGREE = 6
+TOL = 1 / math.pow(10, TOL_DEGREE)
+
+
+def get_measure_parametrize_ranges():
+    confidence_measure_bank = {}
+    alpha_range = (0.25, 0.5, 1.0)
+    bank_exception = None
+    try:
+        confidence_measure_bank = get_confidence_measure_bank()
+    except Exception as e:
+        alpha_range = ()
+        bank_exception = e
+    return confidence_measure_bank, alpha_range, bank_exception
+
+
+def get_aggregation_parametrize_ranges():
+    confidence_aggregation_bank = {}
+    bank_exception = None
+    try:
+        confidence_aggregation_bank = get_confidence_aggregation_bank()
+    except Exception as e:
+        bank_exception = e
+    return confidence_aggregation_bank, bank_exception
+
+
+class TestConfidenceMeasureBank:
+    measure_bank, alphas, bank_build_exception = get_measure_parametrize_ranges()
+
+    @pytest.mark.unit
+    def test_measure_bank(self):
+        if self.bank_build_exception is not None:
+            raise self.bank_build_exception
+
+        assert isinstance(self.measure_bank, dict)
+        assert len(self.measure_bank) > 0
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('measure_name', measure_bank.keys())
+    @pytest.mark.parametrize('alpha', alphas)
+    @pytest.mark.parametrize('vocab_size', VOCAB_SIZES)
+    def test_confidence_measures_one(self, measure_name, alpha, vocab_size):
+        measure = self.measure_bank[measure_name]
+
+        assert torch.allclose(measure(ONE_VEC_SET[vocab_size], vocab_size, alpha), torch.tensor([1.0, 1.0]), atol=TOL)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('measure_name', measure_bank.keys())
+    @pytest.mark.parametrize('alpha', alphas)
+    @pytest.mark.parametrize('vocab_size', VOCAB_SIZES)
+    def test_confidence_measures_zero(self, measure_name, alpha, vocab_size):
+        measure = self.measure_bank[measure_name]
+
+        assert torch.allclose(measure(ZERO_VEC_SET[vocab_size], vocab_size, alpha), torch.tensor([0.0, 0.0]), atol=TOL)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('measure_name', measure_bank.keys())
+    @pytest.mark.parametrize('alpha', alphas)
+    @pytest.mark.parametrize('vocab_size', VOCAB_SIZES)
+    def test_confidence_measures_partial_order(self, measure_name, alpha, vocab_size):
+        measure = self.measure_bank[measure_name]
+        value_normal = round(float(measure(RAND_VEC_SET[vocab_size], vocab_size, alpha)[0]), TOL_DEGREE)
+        value_overfit = round(float(measure(OVERFIT_RAND_VEC_SET[vocab_size], vocab_size, alpha)[0]), TOL_DEGREE)
+
+        assert 0 <= value_normal < value_overfit <= 1, (
+            measure(RAND_VEC_SET[vocab_size], vocab_size, alpha),
+            measure(OVERFIT_RAND_VEC_SET[vocab_size], vocab_size, alpha),
+        )
+
+
+class TestConfidenceAggregationBank:
+    aggregation_bank, bank_build_exception = get_aggregation_parametrize_ranges()
+
+    @pytest.mark.unit
+    def test_aggregation_bank(self):
+        if self.bank_build_exception is not None:
+            raise self.bank_build_exception
+
+        assert isinstance(self.aggregation_bank, dict)
+        assert len(self.aggregation_bank) > 0
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('aggregation_name', aggregation_bank.keys())
+    def test_confidence_agregation_simple(self, aggregation_name):
+        # alaptev: would skipif work with parametrize arguments?
+        if aggregation_name not in ("mean", "min", "max", "prod"):
+            pytest.skip(f"{aggregation_name} is not a simple aggregation")
+        aggregation = self.aggregation_bank[aggregation_name]
+        if aggregation_name == "mean":
+            assert aggregation(AGGREGATION_VEC_SIMPLE) == 0.5
+        elif aggregation_name == "min":
+            assert aggregation(AGGREGATION_VEC_SIMPLE) == 0.0
+        if aggregation_name == "max":
+            assert aggregation(AGGREGATION_VEC_SIMPLE) == 1.0
+        if aggregation_name == "prod":
+            assert aggregation(AGGREGATION_VEC_SIMPLE) == 0.0

tests/collections/asr/conftest.py

@@ -0,0 +1,387 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Optional, Type
+
+import numpy as np
+import pytest
+import torch
+
+from nemo.collections.asr.models import ASRModel
+
+
+class RNNTTestHelper:
+    @staticmethod
+    def wrap_and_call(fn, acts, labels, device, input_lengths=None, target_lengths=None):
+        if not torch.is_tensor(acts):
+            acts = torch.FloatTensor(acts)
+
+        if 'cuda' in device:
+            acts = acts.cuda()
+
+        if not acts.requires_grad:
+            acts.requires_grad = True
+
+        labels = torch.LongTensor(labels)
+
+        if input_lengths is None:
+            lengths = [acts.shape[1]] * acts.shape[0]
+            lengths = torch.LongTensor(lengths)
+        else:
+            lengths = input_lengths
+
+        if target_lengths is None:
+            label_lengths = [len(l) for l in labels]
+            label_lengths = torch.LongTensor(label_lengths)
+        else:
+            label_lengths = target_lengths
+
+        if 'cuda' in device:
+            labels = labels.cuda()
+            lengths = lengths.cuda()
+            label_lengths = label_lengths.cuda()
+
+        costs = fn(acts, labels, lengths, label_lengths)
+        cost = torch.sum(costs)
+        cost.backward()
+
+        if 'cuda' in device:
+            torch.cuda.synchronize()
+
+        if acts.grad is not None:
+            grad = acts.grad.data.cpu().numpy()
+        else:
+            grad = None
+
+        return costs.data.cpu().numpy(), grad
+
+
+@dataclass
+class RnntLossSampleData:
+    vocab_size: int
+    blank_id: int
+
+    logits: torch.Tensor
+    targets: torch.Tensor
+    input_lengths: torch.Tensor
+    target_lengths: torch.Tensor
+
+    expected_cost: Optional[torch.Tensor] = None
+    expected_grads: Optional[torch.Tensor] = None
+
+    @classmethod
+    def get_sample_small(cls) -> "RnntLossSampleData":
+        activations = np.array(
+            [
+                [
+                    [[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.2, 0.8, 0.1]],
+                    [[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.2, 0.1, 0.1], [0.7, 0.1, 0.2, 0.1, 0.1]],
+                ]
+            ]
+        )
+        labels = np.asarray([[1, 2]])
+
+        expected_cost = [4.495666]
+        expected_grads = np.array(
+            [
+                [
+                    [
+                        [-0.13116688, -0.3999269, 0.17703125, 0.17703125, 0.17703125],
+                        [-0.18572757, 0.12247056, -0.18168412, 0.12247056, 0.12247056],
+                        [-0.32091254, 0.06269141, 0.06928472, 0.12624499, 0.06269141],
+                    ],
+                    [
+                        [0.05456069, -0.21824276, 0.05456069, 0.05456069, 0.05456069],
+                        [0.12073959, 0.12073959, -0.48295835, 0.12073959, 0.12073959],
+                        [-0.6925882, 0.16871116, 0.18645467, 0.16871116, 0.16871116],
+                    ],
+                ]
+            ]
+        )
+        return RnntLossSampleData(
+            vocab_size=3,
+            blank_id=0,
+            logits=torch.from_numpy(activations).to(torch.float32),
+            targets=torch.from_numpy(labels),
+            input_lengths=torch.tensor([2]),
+            target_lengths=torch.tensor([2]),
+            expected_cost=torch.tensor(expected_cost).to(torch.float32),
+            expected_grads=torch.from_numpy(expected_grads),
+        )
+
+    @classmethod
+    def get_sample_small_blank_last(cls) -> "RnntLossSampleData":
+        activations = np.array(
+            [
+                [
+                    [[0.0, 1.0, 3.0], [0.0, 2.0, 3.0], [1.0, 1.0, 3.0], [2.0, 3.0, 2.0]],
+                    [[0.0, 0.0, 1.0], [0.0, 1.0, 1.0], [1.0, 0.0, 1.0], [2.0, 2.0, 0.0]],
+                    [[0.0, 2.0, 5.0], [0.0, 3.0, 5.0], [1.0, 2.0, 5.0], [2.0, 4.0, 4.0]],
+                    [[0.0, 3.0, 4.0], [0.0, 4.0, 4.0], [1.0, 3.0, 4.0], [2.0, 5.0, 3.0]],
+                    [[2.0, 2.0, 1.0], [2.0, 3.0, 1.0], [3.0, 2.0, 1.0], [4.0, 4.0, 0.0]],
+                ]
+            ]
+        )
+        labels = np.array([[0, 1, 0]])
+
+        expected_cost = [6.789285182952881]
+        expected_grads = np.array(
+            [
+                [
+                    [
+                        [-0.03551076725125313, 0.11419519782066345, -0.07868456840515137],
+                        [0.0027224558871239424, 0.00704305712133646, -0.009765520691871643],
+                        [0.0013856772566214204, 0.0013924005907028913, -0.0027780719101428986],
+                        [1.4249643527364242e-06, 3.873454716085689e-06, -5.298420546751004e-06],
+                    ],
+                    [
+                        [-0.1934257447719574, 0.19551163911819458, -0.0020859241485595703],
+                        [0.07043898105621338, 0.05738453567028046, -0.12782356142997742],
+                        [0.061031512916088104, 0.02286236733198166, -0.08389391005039215],
+                        [0.0005252412520349026, 0.0005252412520349026, -0.0010504829697310925],
+                    ],
+                    [
+                        [-0.007841046899557114, 0.025142310187220573, -0.017301201820373535],
+                        [0.0019501042552292347, 0.0005148053169250488, -0.0024650096893310547],
+                        [0.0027856370434165, 0.008609085343778133, -0.01139475405216217],
+                        [9.526080975774676e-05, 0.0007038871408440173, -0.000799147819634527],
+                    ],
+                    [
+                        [-0.01533521432429552, 0.1386115401983261, -0.12327653169631958],
+                        [0.002850571647286415, -0.006693005561828613, 0.003842458128929138],
+                        [0.009236274287104607, 0.08995233476161957, -0.0991886705160141],
+                        [0.0001865450612967834, 0.0037468576338142157, -0.003933403175324202],
+                    ],
+                    [
+                        [-0.2888762652873993, 0.211185485124588, 0.07769080251455307],
+                        [0.15952755510807037, -0.2182144820690155, 0.05868690833449364],
+                        [-0.3332723379135132, 0.2436419129371643, 0.0896308496594429],
+                        [0.4954628646373749, 0.4954628646373749, -0.9909257292747498],
+                    ],
+                ]
+            ]
+        )
+        return RnntLossSampleData(
+            vocab_size=3,
+            blank_id=2,
+            logits=torch.from_numpy(activations).to(torch.float32),
+            targets=torch.from_numpy(labels),
+            input_lengths=torch.tensor([5]),
+            target_lengths=torch.tensor([3]),
+            expected_cost=torch.tensor(expected_cost).to(torch.float32),
+            expected_grads=torch.from_numpy(expected_grads),
+        )
+
+    @classmethod
+    def get_sample_medium(cls) -> "RnntLossSampleData":
+        # minibatch x T x U x alphabet_size
+        activations = [
+            [
+                [
+                    [0.06535690384862791, 0.7875301411923206, 0.08159176605666074],
+                    [0.5297155426466327, 0.7506749639230854, 0.7541348379087998],
+                    [0.6097641124736383, 0.8681404965673826, 0.6225318186056529],
+                ],
+                [
+                    [0.6685222872103057, 0.8580392805336061, 0.16453892311765583],
+                    [0.989779515236694, 0.944298460961015, 0.6031678586829663],
+                    [0.9467833543605416, 0.666202507295747, 0.28688179752461884],
+                ],
+                [
+                    [0.09418426230195986, 0.3666735970751962, 0.736168049462793],
+                    [0.1666804425271342, 0.7141542198635192, 0.3993997272216727],
+                    [0.5359823524146038, 0.29182076440286386, 0.6126422611507932],
+                ],
+                [
+                    [0.3242405528768486, 0.8007644367291621, 0.5241057606558068],
+                    [0.779194617063042, 0.18331417220174862, 0.113745182072432],
+                    [0.24022162381327106, 0.3394695622533106, 0.1341595066017014],
+                ],
+            ],
+            [
+                [
+                    [0.5055615569388828, 0.051597282072282646, 0.6402903936686337],
+                    [0.43073311517251, 0.8294731834714112, 0.1774668847323424],
+                    [0.3207001991262245, 0.04288308912457006, 0.30280282975568984],
+                ],
+                [
+                    [0.6751777088333762, 0.569537369330242, 0.5584738347504452],
+                    [0.08313242153985256, 0.06016544344162322, 0.10795752845152584],
+                    [0.7486153608562472, 0.943918041459349, 0.4863558118797222],
+                ],
+                [
+                    [0.4181986264486809, 0.6524078485043804, 0.024242983423721887],
+                    [0.13458171554507403, 0.3663418070512402, 0.2958297395361563],
+                    [0.9236695822497084, 0.6899291482654177, 0.7418981733448822],
+                ],
+                [
+                    [0.25000547599982104, 0.6034295486281007, 0.9872887878887768],
+                    [0.5926057265215715, 0.8846724004467684, 0.5434495396894328],
+                    [0.6607698886038497, 0.3771277082495921, 0.3580209022231813],
+                ],
+            ],
+        ]
+
+        expected_cost = [4.2806528590890736, 3.9384369822503591]
+        expected_grads = [
+            [
+                [
+                    [-1.86843902e-01, -6.25548810e-02, 2.49398798e-01],
+                    [-2.03376666e-01, 2.02399328e-01, 9.77333169e-04],
+                    [-1.41016081e-01, 7.91234672e-02, 6.18926100e-02],
+                ],
+                [
+                    [-1.15517676e-02, -8.12802389e-02, 9.28319991e-02],
+                    [-1.54257029e-01, 2.29432687e-01, -7.51756504e-02],
+                    [-2.46593088e-01, 1.46404594e-01, 1.00188486e-01],
+                ],
+                [
+                    [-1.29182907e-02, -6.15932420e-02, 7.45115355e-02],
+                    [-5.59857301e-02, 2.19830811e-01, -1.63845062e-01],
+                    [-4.97626871e-01, 2.09239945e-01, 2.88386941e-01],
+                ],
+                [
+                    [1.36048580e-02, -3.02196294e-02, 1.66147724e-02],
+                    [1.13924511e-01, 6.27811998e-02, -1.76705718e-01],
+                    [-6.67078257e-01, 3.67658824e-01, 2.99419403e-01],
+                ],
+            ],
+            [
+                [
+                    [-3.56343776e-01, -5.53474613e-02, 4.11691219e-01],
+                    [-9.69219357e-02, 2.94591039e-02, 6.74628317e-02],
+                    [-6.35175705e-02, 2.76544970e-02, 3.58630717e-02],
+                ],
+                [
+                    [-1.54499024e-01, -7.39420280e-02, 2.28441030e-01],
+                    [-1.66789949e-01, -8.78955179e-05, 1.66877866e-01],
+                    [-1.72369644e-01, 1.05565332e-01, 6.68043196e-02],
+                ],
+                [
+                    [2.38748826e-02, -1.18255816e-01, 9.43809375e-02],
+                    [-1.04707085e-01, -1.08934477e-01, 2.13641584e-01],
+                    [-3.69844258e-01, 1.80118099e-01, 1.89726159e-01],
+                ],
+                [
+                    [2.57137045e-02, -7.94617534e-02, 5.37480488e-02],
+                    [1.22328237e-01, -2.38788679e-01, 1.16460443e-01],
+                    [-5.98686993e-01, 3.02203178e-01, 2.96483815e-01],
+                ],
+            ],
+        ]
+        activations = np.array(activations)
+        labels = np.array([[1, 2], [1, 1]])
+        expected_grads = np.array(expected_grads)
+
+        return RnntLossSampleData(
+            vocab_size=3,
+            blank_id=0,
+            logits=torch.from_numpy(activations).to(torch.float32),
+            targets=torch.from_numpy(labels),
+            input_lengths=torch.tensor([4, 4]),
+            target_lengths=torch.tensor([2, 2]),
+            expected_cost=torch.tensor(expected_cost).to(torch.float32),
+            expected_grads=torch.from_numpy(expected_grads),
+        )
+
+    @classmethod
+    def get_sample_small_random(cls, blank_first: bool, device=torch.device("cpu")) -> "RnntLossSampleData":
+        vocab_size = 4
+        blank_id = 0 if blank_first else vocab_size - 1
+        num_frames = 4
+        text_len = 2
+        if blank_first:
+            text = np.asarray([1, 3])
+        else:
+            text = np.asarray([0, 2])
+
+        targets = torch.from_numpy(text).unsqueeze(0).to(device)
+        logits = torch.rand([1, num_frames, text_len + 1, vocab_size], requires_grad=True, device=device)
+        input_lengths = torch.tensor([num_frames], device=device)
+        target_lengths = torch.tensor([text_len], device=device)
+        return RnntLossSampleData(
+            vocab_size=vocab_size,
+            blank_id=blank_id,
+            logits=logits,
+            targets=targets,
+            input_lengths=input_lengths,
+            target_lengths=target_lengths,
+        )
+
+    @classmethod
+    def get_sample_medium_random_var_size(cls, blank_first: bool, device=torch.device("cpu")) -> "RnntLossSampleData":
+        vocab_size = 32
+        blank_id = 0 if blank_first else vocab_size - 1
+        num_frames = 32
+        text_len = 27
+        min_symbol = 1 if blank_first else 0
+        max_symbol = vocab_size if blank_first else vocab_size - 1
+        batch_size = 4
+
+        rs = np.random.RandomState(2021)
+        text = rs.randint(min_symbol, max_symbol, size=(batch_size, text_len))
+        targets = torch.from_numpy(text).to(device)
+
+        logits = torch.rand([batch_size, num_frames, text_len + 1, vocab_size], requires_grad=True, device=device)
+        input_lengths = torch.tensor([num_frames, num_frames // 2, text_len, text_len // 2], device=device).long()
+        target_lengths = torch.tensor([text_len, text_len - 1, text_len - 3, text_len - 10], device=device)
+        return RnntLossSampleData(
+            vocab_size=vocab_size,
+            blank_id=blank_id,
+            logits=logits,
+            targets=targets,
+            input_lengths=input_lengths,
+            target_lengths=target_lengths,
+        )
+
+
+@pytest.fixture(scope="session")
+def rnnt_test_helper() -> Type[RNNTTestHelper]:
+    return RNNTTestHelper
+
+
+@pytest.fixture(scope="session")
+def rnn_loss_sample_data() -> Type[RnntLossSampleData]:
+    return RnntLossSampleData
+
+
+@pytest.fixture(scope='session')
+def fast_conformer_transducer_model():
+    return ASRModel.from_pretrained("stt_en_fastconformer_transducer_large")
+
+
+@pytest.fixture(scope='session')
+def fast_conformer_ctc_model():
+    return ASRModel.from_pretrained("stt_en_fastconformer_ctc_large")
+
+
+@pytest.fixture(scope='session')
+def fast_conformer_hybrid_model():
+    return ASRModel.from_pretrained("parakeet-tdt_ctc-110m")
+
+
+@pytest.fixture(scope='session')
+def canary_1b_flash():
+    return ASRModel.restore_from("/home/TestData/asr/canary/models/canary-1b-flash_HF_20250318.nemo")
+
+
+@pytest.fixture(scope='session')
+def canary_1b_v2():
+    return ASRModel.restore_from("/home/TestData/asr/canary/models/canary-1b-v2_20250809.nemo")
+
+
+@pytest.fixture(scope='session')
+def hybrid_rnnt_ctc_bpe_model_with_prompt():
+    return ASRModel.restore_from("/home/TestData/asr/hybrid_rnnt_ctc_bpe_model_with_prompt.nemo")

tests/collections/asr/decoding/conftest.py

@@ -0,0 +1,135 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from pathlib import Path
+
+import pytest
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
+from tests.collections.asr.decoding.utils import make_preprocessor_deterministic, preserve_decoding_cfg_and_cpu_device
+
+CHECKPOINTS_PATH = Path("/home/TestData/asr")
+
+
+@pytest.fixture(scope="session")
+def an4_val_manifest_corrected(tmp_path_factory, test_data_dir):
+    """
+    Correct an4_val manifest audio filepaths, e.g.,
+    "tests/data/asr/test/an4/wav/an440-mjgm-b.wav" -> test_data_dir / "test/an4/wav/an440-mjgm-b.wav"
+    """
+    an4_val_manifest_orig_path = Path(test_data_dir) / "asr/an4_val.json"
+    an4_val_manifest_corrected_path = tmp_path_factory.mktemp("manifests") / "an4_val_corrected.json"
+    an4_val_records = read_manifest(an4_val_manifest_orig_path)
+    for record in an4_val_records:
+        record["audio_filepath"] = record["audio_filepath"].replace(
+            "tests/data/asr", str(an4_val_manifest_orig_path.resolve().parent)
+        )
+    write_manifest(an4_val_manifest_corrected_path, an4_val_records)
+    return an4_val_manifest_corrected_path
+
+
+@pytest.fixture(scope="session")
+def an4_train_manifest_corrected(tmp_path_factory, test_data_dir):
+    """
+    Correct an4_train manifest audio filepaths, e.g.,
+    "tests/data/asr/test/an4/wav/an440-mjgm-b.wav" -> test_data_dir / "test/an4/wav/an440-mjgm-b.wav"
+    """
+    an4_train_manifest_orig_path = Path(test_data_dir) / "asr/an4_train.json"
+    an4_train_manifest_corrected_path = tmp_path_factory.mktemp("manifests") / "an4_train_corrected.json"
+    an4_train_records = read_manifest(an4_train_manifest_orig_path)
+    for record in an4_train_records:
+        record["audio_filepath"] = record["audio_filepath"].replace(
+            "tests/data/asr", str(an4_train_manifest_orig_path.resolve().parent)
+        )
+    write_manifest(an4_train_manifest_corrected_path, an4_train_records)
+    return an4_train_manifest_corrected_path
+
+
+@pytest.fixture(scope="package")
+def _stt_en_conformer_transducer_small_raw():
+    if CHECKPOINTS_PATH.exists():
+        model = ASRModel.restore_from(
+            str(CHECKPOINTS_PATH / "stt_en_conformer_transducer_small.nemo"), map_location="cpu"
+        )
+    else:
+        model_name = "stt_en_conformer_transducer_small"
+        model = ASRModel.from_pretrained(model_name, map_location="cpu")
+    make_preprocessor_deterministic(model)
+    return model
+
+
+@pytest.fixture(scope="package")
+def _stt_en_fastconformer_transducer_large_raw():
+    if CHECKPOINTS_PATH.exists():
+        model = ASRModel.restore_from(
+            str(CHECKPOINTS_PATH / "stt_en_fastconformer_transducer_large.nemo"), map_location="cpu"
+        )
+    else:
+        model_name = "stt_en_fastconformer_transducer_large"
+        model = ASRModel.from_pretrained(model_name, map_location="cpu")
+    make_preprocessor_deterministic(model)
+    return model
+
+
+@pytest.fixture(scope="package")
+def _stt_en_fastconformer_tdt_large_raw():
+    if CHECKPOINTS_PATH.exists():
+        model = ASRModel.restore_from(
+            str(CHECKPOINTS_PATH / "stt_en_fastconformer_tdt_large.nemo"), map_location="cpu"
+        )
+    else:
+        model_name = "nvidia/stt_en_fastconformer_tdt_large"
+        model = ASRModel.from_pretrained(model_name, map_location="cpu")
+    make_preprocessor_deterministic(model)
+    return model
+
+
+@pytest.fixture(scope="package")
+def _canary_180m_flash_raw():
+    model_name = "nvidia/canary-180m-flash"
+    model = ASRModel.from_pretrained(model_name, map_location="cpu")
+    make_preprocessor_deterministic(model)
+    return model
+
+
+@pytest.fixture
+def stt_en_conformer_transducer_small(_stt_en_conformer_transducer_small_raw):
+    """Function-level fixture for model. Guarantees to preserve decoding config and device"""
+    model = _stt_en_conformer_transducer_small_raw
+    with preserve_decoding_cfg_and_cpu_device(model):
+        yield model
+
+
+@pytest.fixture
+def stt_en_fastconformer_transducer_large(_stt_en_fastconformer_transducer_large_raw):
+    """Function-level fixture for model. Guarantees to preserve decoding config and device"""
+    model = _stt_en_fastconformer_transducer_large_raw
+    with preserve_decoding_cfg_and_cpu_device(model):
+        yield model
+
+
+@pytest.fixture
+def stt_en_fastconformer_tdt_large(_stt_en_fastconformer_tdt_large_raw):
+    """Function-level fixture for model. Guarantees to preserve decoding config and device"""
+    model = _stt_en_fastconformer_tdt_large_raw
+    with preserve_decoding_cfg_and_cpu_device(model):
+        yield model
+
+
+@pytest.fixture
+def canary_180m_flash(_canary_180m_flash_raw):
+    """Function-level fixture for model. Guarantees to preserve decoding config and device"""
+    model = _canary_180m_flash_raw
+    with preserve_decoding_cfg_and_cpu_device(model):
+        yield model

tests/collections/asr/decoding/test_batched_beam_decoding.py

@@ -0,0 +1,883 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import glob
+import os
+from pathlib import Path
+
+import pytest
+import torch
+from kaldialign import edit_distance
+from omegaconf import open_dict
+from tqdm import tqdm
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.models.ctc_models import EncDecCTCModel
+from nemo.collections.asr.parts.submodules.ctc_beam_decoding import BeamBatchedCTCInfer
+from nemo.collections.asr.parts.submodules.ngram_lm import NGramGPULanguageModel
+from nemo.collections.asr.parts.submodules.rnnt_beam_decoding import BeamBatchedRNNTInfer
+from nemo.collections.asr.parts.submodules.tdt_beam_decoding import BeamBatchedTDTInfer
+from nemo.collections.asr.parts.utils import rnnt_utils
+from nemo.core.utils import numba_utils
+from nemo.core.utils.cuda_python_utils import skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+from tests.collections.asr.decoding.utils import load_audio
+
+RNNT_MODEL = "stt_en_conformer_transducer_small"
+CTC_MODEL = "nvidia/stt_en_conformer_ctc_small"
+TDT_MODEL = "nvidia/stt_en_fastconformer_tdt_large"
+MAX_SAMPLES = 10
+
+DEVICES = [torch.device("cpu")]
+
+if torch.cuda.is_available():
+    DEVICES.append(torch.device('cuda'))
+
+NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cpu_is_supported(
+    __NUMBA_MINIMUM_VERSION__
+) or numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+
+
+# available audio filename fixtures
+@pytest.fixture(scope="module")
+def test_audio_filenames(test_data_dir):
+    return tuple(glob.glob(os.path.join(test_data_dir, "asr", "test", "an4", "wav", "*.wav")))
+
+
+# model fixtures
+@pytest.fixture(scope="module")
+def rnnt_model():
+    model = ASRModel.from_pretrained(model_name=RNNT_MODEL, map_location="cpu")
+    model.eval()
+    return model
+
+
+@pytest.fixture(scope="module")
+def tdt_model():
+    model = ASRModel.from_pretrained(model_name=TDT_MODEL, map_location="cpu")
+    model.eval()
+    return model
+
+
+@pytest.fixture(scope="module")
+def ctc_model():
+    model = ASRModel.from_pretrained(model_name=CTC_MODEL, map_location="cpu")
+    model.eval()
+    return model
+
+
+# encoder output fixtures
+@pytest.fixture(scope="module")
+def get_rnnt_encoder_output(rnnt_model, test_audio_filenames):
+    encoder_output, encoded_lengths = get_transducer_model_encoder_output(
+        test_audio_filenames, MAX_SAMPLES, rnnt_model
+    )
+    return encoder_output, encoded_lengths
+
+
+@pytest.fixture(scope="module")
+def get_tdt_encoder_output(tdt_model, test_audio_filenames):
+    encoder_output, encoded_lengths = get_transducer_model_encoder_output(test_audio_filenames, MAX_SAMPLES, tdt_model)
+    return encoder_output, encoded_lengths
+
+
+@pytest.fixture(scope="module")
+def get_ctc_output(ctc_model, test_audio_filenames):
+    encoder_output, encoded_lengths = get_ctc_model_output(test_audio_filenames, MAX_SAMPLES, ctc_model)
+    return encoder_output, encoded_lengths
+
+
+@pytest.fixture(scope="module")
+def kenlm_model_path(tmp_path_factory, test_data_dir):
+    lm_path = Path(test_data_dir) / "asr/kenlm_ngram_lm/parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+    assert os.path.exists(lm_path), f"LM file not found: {lm_path}"
+    lm_nemo_path = tmp_path_factory.mktemp("lm") / f"{lm_path.name}.nemo"
+    NGramGPULanguageModel.from_file(lm_path, vocab_size=1024).save_to(f"{lm_nemo_path}")
+    return f"{lm_nemo_path}"
+
+
+def get_transducer_model_encoder_output(
+    test_audio_filenames,
+    num_samples: int,
+    model: ASRModel,
+    device: torch.device = torch.device("cpu"),
+    dtype: torch.dtype = torch.float32,
+):
+    audio_filepaths = test_audio_filenames[:num_samples]
+
+    with torch.no_grad():
+        model.preprocessor.featurizer.dither = 0.0
+        model.preprocessor.featurizer.pad_to = 0
+        model.eval()
+
+        all_inputs, all_lengths = [], []
+        for audio_file in tqdm(audio_filepaths, desc="Loading audio files"):
+            audio_tensor, _ = load_audio(audio_file)
+            all_inputs.append(audio_tensor)
+            all_lengths.append(torch.tensor(audio_tensor.shape[0], dtype=torch.int64))
+
+        input_batch = torch.nn.utils.rnn.pad_sequence(all_inputs, batch_first=True).to(device=device, dtype=dtype)
+        length_batch = torch.tensor(all_lengths, dtype=torch.int64).to(device)
+
+        encoded_outputs, encoded_length = model(input_signal=input_batch, input_signal_length=length_batch)
+
+    return encoded_outputs, encoded_length
+
+
+def get_ctc_model_output(
+    test_audio_filenames,
+    num_samples: int,
+    model: ASRModel,
+    device: torch.device = torch.device("cpu"),
+    dtype: torch.dtype = torch.float32,
+):
+    audio_filepaths = test_audio_filenames[:num_samples]
+
+    with torch.no_grad():
+        model.preprocessor.featurizer.dither = 0.0
+        model.preprocessor.featurizer.pad_to = 0
+        model.eval()
+
+        all_inputs, all_lengths = [], []
+        for audio_file in tqdm(audio_filepaths, desc="Loading audio files"):
+            audio_tensor, _ = load_audio(audio_file)
+            all_inputs.append(audio_tensor)
+            all_lengths.append(torch.tensor(audio_tensor.shape[0], dtype=torch.int64))
+
+        input_batch = torch.nn.utils.rnn.pad_sequence(all_inputs, batch_first=True).to(device=device, dtype=dtype)
+        length_batch = torch.tensor(all_lengths, dtype=torch.int64).to(device)
+
+        log_probs, encoded_length, _ = model(input_signal=input_batch, input_signal_length=length_batch)
+
+    return log_probs, encoded_length
+
+
+def print_unit_test_info(strategy, batch_size, beam_size, allow_cuda_graphs, device):
+    print(
+        f"""Beam search algorithm: {strategy},
+                Batch size: {batch_size},
+                Beam size: {beam_size},
+                Cuda Graphs: {allow_cuda_graphs},
+                Decoding device: {device}
+            """
+    )
+
+
+def check_res_best_hyps(num_samples, hyps):
+    assert type(hyps) == list
+    assert type(hyps[0]) == rnnt_utils.Hypothesis
+
+    assert len(hyps) == num_samples
+
+    assert all(
+        [
+            hasattr(hyps[hyp_idx], "y_sequence")
+            and hasattr(hyps[hyp_idx], "score")
+            and hasattr(hyps[hyp_idx], "timestamp")
+            for hyp_idx in range(num_samples)
+        ]
+    )
+
+
+def print_res_best_hyps(hyps):
+    for hyp_idx, hyp in enumerate(hyps):
+        print("Sample: ", hyp_idx)
+        print("Decoded text: ", hyp.text)
+        print("Score: ", hyp.score)
+        print("Transcript", hyp.y_sequence)
+        print("Timesteps", hyp.timestamp)
+        print()
+
+
+def check_res_nbest_hyps(num_samples, batch_nbest_hyps):
+    assert type(batch_nbest_hyps) == list
+    assert type(batch_nbest_hyps[0]) == rnnt_utils.NBestHypotheses
+
+    assert len(batch_nbest_hyps) == num_samples
+
+    for idx in range(num_samples):
+        assert all(
+            [
+                hasattr(batch_nbest_hyps[idx].n_best_hypotheses[hyp_idx], "y_sequence")
+                and hasattr(batch_nbest_hyps[idx].n_best_hypotheses[hyp_idx], "score")
+                and hasattr(batch_nbest_hyps[idx].n_best_hypotheses[hyp_idx], "timestamp")
+                for hyp_idx in range(len(batch_nbest_hyps[idx].n_best_hypotheses))
+            ]
+        )
+
+        # Empty transcript (blank-only beam) is valid; y_sequence and timestamp must stay aligned.
+        assert all(
+            len(batch_nbest_hyps[idx].n_best_hypotheses[hyp_idx].y_sequence)
+            == len(batch_nbest_hyps[idx].n_best_hypotheses[hyp_idx].timestamp)
+            for hyp_idx in range(len(batch_nbest_hyps[idx].n_best_hypotheses))
+        )
+
+
+def print_res_nbest_hyps(batch_nbest_hyps):
+    for batch_idx, nbest_hyps in enumerate(batch_nbest_hyps):
+        print(f"Batch idx: {batch_idx}")
+        for idx, hyp in enumerate(nbest_hyps):
+            print(f"Hyp index: {idx + 1}")
+            print("Text: ", hyp.text)
+            print("Score: ", hyp.score)
+            print("Transcripts: ", hyp.y_sequence)
+            print("Timesteps: ", hyp.timestamp)
+            print()
+
+
+def decode_text_from_hypotheses(hyps, model):
+    if isinstance(model, EncDecCTCModel):
+        return model.decoding.decode_hypothesis(hyps, fold_consecutive=False)
+    else:
+        return model.decoding.decode_hypothesis(hyps)
+
+
+def decode_text_from_nbest_hypotheses(hyps, model):
+    if isinstance(model, EncDecCTCModel):
+        return [
+            model.decoding.decode_hypothesis(nbest_hyp.n_best_hypotheses, fold_consecutive=False) for nbest_hyp in hyps
+        ]
+    else:
+        return [model.decoding.decode_hypothesis(nbest_hyp.n_best_hypotheses) for nbest_hyp in hyps]
+
+
+class TestRNNTDecoding:
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "malsd_batch", "allow_cuda_graphs": False},
+            {"search_type": "malsd_batch", "allow_cuda_graphs": True},
+            {"search_type": "maes_batch", "allow_cuda_graphs": False},
+        ],
+    )
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("batch_size", [4, 16])
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_rnnt_beam_decoding_return_best_hypothesis(
+        self, test_audio_filenames, rnnt_model, get_rnnt_encoder_output, beam_config, device, batch_size, beam_size
+    ):
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = rnnt_model.to(device)
+        encoder_output, encoded_lengths = get_rnnt_encoder_output
+        encoder_output, encoded_lengths = encoder_output[:num_samples].to(device), encoded_lengths[:num_samples].to(
+            device
+        )
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedRNNTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=vocab_size,
+            beam_size=beam_size,
+            score_norm=True,
+            return_best_hypothesis=True,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy=beam_config['search_type'],
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            hyps = decoding(encoder_output=encoder_output, encoded_lengths=encoded_lengths)[0]
+
+            check_res_best_hyps(num_samples, hyps)
+            hyps = decode_text_from_hypotheses(hyps, model)
+            print_res_best_hyps(hyps)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "malsd_batch", "allow_cuda_graphs": False},
+            {"search_type": "malsd_batch", "allow_cuda_graphs": True},
+            {"search_type": "maes_batch", "allow_cuda_graphs": False},
+        ],
+    )
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("batch_size", [4])
+    def test_rnnt_beam_decoding_return_nbest(
+        self, test_audio_filenames, rnnt_model, get_rnnt_encoder_output, beam_config, device, beam_size, batch_size
+    ):
+        device = torch.device("cuda")
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = rnnt_model.to(device)
+        encoder_output, encoded_lengths = get_rnnt_encoder_output
+        encoder_output, encoded_lengths = encoder_output[:num_samples].to(device), encoded_lengths[:num_samples].to(
+            device
+        )
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedRNNTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=vocab_size,
+            beam_size=beam_size,
+            score_norm=True,
+            return_best_hypothesis=False,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy=beam_config['search_type'],
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            batch_nbest_hyps = decoding(encoder_output=encoder_output, encoded_lengths=encoded_lengths)[0]
+
+            check_res_nbest_hyps(num_samples, batch_nbest_hyps)
+            batch_nbest_hyps = decode_text_from_nbest_hypotheses(batch_nbest_hyps, model)
+            print_res_nbest_hyps(batch_nbest_hyps)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "malsd_batch", "allow_cuda_graphs": False},
+            {"search_type": "maes_batch", "allow_cuda_graphs": False},
+            {"search_type": "malsd_batch", "allow_cuda_graphs": True},
+        ],
+    )
+    @pytest.mark.parametrize("batch_size", [4])
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("pruning_mode", ["late", "early"])
+    @pytest.mark.parametrize("blank_lm_score_mode", ["no_score", "lm_weighted_full"])
+    def test_rnnt_beam_decoding_kenlm(
+        self,
+        kenlm_model_path,
+        test_audio_filenames,
+        rnnt_model,
+        get_rnnt_encoder_output,
+        beam_config,
+        device,
+        batch_size,
+        beam_size,
+        pruning_mode,
+        blank_lm_score_mode,
+    ):
+        device = torch.device("cuda")
+
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = rnnt_model.to(device)
+        encoder_output, encoded_lengths = get_rnnt_encoder_output
+        encoder_output, encoded_lengths = encoder_output[:num_samples].to(device), encoded_lengths[:num_samples].to(
+            device
+        )
+
+        vocab_size = model.tokenizer.vocab_size
+
+        fusion_models = [NGramGPULanguageModel.from_file(lm_path=kenlm_model_path, vocab_size=vocab_size)]
+        fusion_models_alpha = [0.3]
+
+        decoding = BeamBatchedRNNTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=vocab_size,
+            beam_size=beam_size,
+            score_norm=True,
+            return_best_hypothesis=True,
+            pruning_mode=pruning_mode,
+            blank_lm_score_mode=blank_lm_score_mode,
+            fusion_models=fusion_models,
+            fusion_models_alpha=fusion_models_alpha,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy=beam_config['search_type'],
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            hyps = decoding(encoder_output=encoder_output, encoded_lengths=encoded_lengths)[0]
+
+            check_res_best_hyps(num_samples, hyps)
+            hyps = decode_text_from_hypotheses(hyps, model)
+            print_res_best_hyps(hyps)
+
+
+class TestTDTDecoding:
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "malsd_batch", "allow_cuda_graphs": False},
+            {"search_type": "malsd_batch", "allow_cuda_graphs": True},
+        ],
+    )
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("batch_size", [4, 16])
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_tdt_beam_decoding_return_best_hypothesis(
+        self, test_audio_filenames, tdt_model, get_tdt_encoder_output, beam_config, device, batch_size, beam_size
+    ):
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = tdt_model.to(device)
+        encoder_output, encoded_lengths = get_tdt_encoder_output
+        encoder_output, encoded_lengths = encoder_output[:num_samples].to(device), encoded_lengths[:num_samples].to(
+            device
+        )
+
+        model_config = model.to_config_dict()
+        durations = list(model_config["model_defaults"]["tdt_durations"])
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedTDTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=vocab_size,
+            durations=durations,
+            beam_size=beam_size,
+            score_norm=True,
+            return_best_hypothesis=True,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy=beam_config['search_type'],
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            hyps = decoding(encoder_output=encoder_output, encoded_lengths=encoded_lengths)[0]
+
+            check_res_best_hyps(num_samples, hyps)
+            hyps = decode_text_from_hypotheses(hyps, model)
+            print_res_best_hyps(hyps)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "malsd_batch", "allow_cuda_graphs": False},
+            {"search_type": "malsd_batch", "allow_cuda_graphs": True},
+        ],
+    )
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("batch_size", [4])
+    def test_tdt_beam_decoding_return_nbest(
+        self, test_audio_filenames, tdt_model, get_tdt_encoder_output, beam_config, device, beam_size, batch_size
+    ):
+        device = torch.device("cuda")
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = tdt_model.to(device)
+        encoder_output, encoded_lengths = get_tdt_encoder_output
+        encoder_output, encoded_lengths = encoder_output[:num_samples].to(device), encoded_lengths[:num_samples].to(
+            device
+        )
+
+        model_config = model.to_config_dict()
+        durations = list(model_config["model_defaults"]["tdt_durations"])
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedTDTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=vocab_size,
+            durations=durations,
+            beam_size=beam_size,
+            score_norm=True,
+            return_best_hypothesis=False,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy=beam_config['search_type'],
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            batch_nbest_hyps = decoding(encoder_output=encoder_output, encoded_lengths=encoded_lengths)[0]
+
+            check_res_nbest_hyps(num_samples, batch_nbest_hyps)
+            batch_nbest_hyps = decode_text_from_nbest_hypotheses(batch_nbest_hyps, model)
+            print_res_nbest_hyps(batch_nbest_hyps)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "malsd_batch", "allow_cuda_graphs": False},
+            {"search_type": "malsd_batch", "allow_cuda_graphs": True},
+        ],
+    )
+    @pytest.mark.parametrize("batch_size", [4])
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("pruning_mode", ["late", "early"])
+    @pytest.mark.parametrize("blank_lm_score_mode", ["lm_weighted_full", "no_score"])
+    def test_tdt_beam_decoding_kenlm(
+        self,
+        kenlm_model_path,
+        test_audio_filenames,
+        tdt_model,
+        get_tdt_encoder_output,
+        beam_config,
+        device,
+        batch_size,
+        beam_size,
+        pruning_mode,
+        blank_lm_score_mode,
+    ):
+        device = torch.device("cuda")
+
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = tdt_model.to(device)
+        encoder_output, encoded_lengths = get_tdt_encoder_output
+        encoder_output, encoded_lengths = encoder_output[:num_samples].to(device), encoded_lengths[:num_samples].to(
+            device
+        )
+
+        model_config = model.to_config_dict()
+        durations = list(model_config["model_defaults"]["tdt_durations"])
+
+        vocab_size = model.tokenizer.vocab_size
+
+        fusion_models = [NGramGPULanguageModel.from_file(lm_path=kenlm_model_path, vocab_size=vocab_size)]
+        fusion_models_alpha = [0.3]
+
+        decoding = BeamBatchedTDTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=vocab_size,
+            durations=durations,
+            beam_size=beam_size,
+            score_norm=True,
+            return_best_hypothesis=True,
+            pruning_mode=pruning_mode,
+            blank_lm_score_mode=blank_lm_score_mode,
+            fusion_models=fusion_models,
+            fusion_models_alpha=fusion_models_alpha,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy=beam_config['search_type'],
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            hyps = decoding(encoder_output=encoder_output, encoded_lengths=encoded_lengths)[0]
+
+            check_res_best_hyps(num_samples, hyps)
+            hyps = decode_text_from_hypotheses(hyps, model)
+            print_res_best_hyps(hyps)
+
+
+class TestTransducerCudaGraphBeamDecoding:
+    """
+    Tests CudaGraphs implementations from Transducer models (RNN-T and TDT)
+    """
+
+    @pytest.mark.with_downloads
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA decoder can run only on CUDA")
+    @pytest.mark.parametrize("force_mode", ["no_graphs", "no_while_loops", "full_graph"])
+    @pytest.mark.parametrize("model_type", ["rnnt", "tdt"])
+    def test_stated_stateless(self, test_audio_filenames, rnnt_model, tdt_model, model_type, force_mode: str):
+        """
+        Compares pure Pytorch and with three modes of statefull implementations for double floating point precision.
+            1. Pure pytorch, but statefull implementation: no_graphs
+            2. With CudaGrpahs: no_while_loops and full_graph.
+        """
+        if force_mode == "full_graph":
+            skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+
+        batch_size = 16
+        device = torch.device("cuda")
+        model = rnnt_model.to(device) if model_type == "rnnt" else tdt_model.to(device)
+        decoding_config = copy.deepcopy(model.cfg.decoding)
+
+        with open_dict(decoding_config):
+            decoding_config["strategy"] = "malsd_batch"
+            decoding_config["beam"]["beam_size"] = 4
+            decoding_config["beam"]["return_best_hypothesis"] = False
+            decoding_config["beam"]["allow_cuda_graphs"] = False
+
+        model.change_decoding_strategy(decoding_config)
+
+        actual_hypotheses = model.transcribe(test_audio_filenames, batch_size=batch_size, num_workers=None)
+        actual_transcripts = [[hyp.text for hyp in actual_beam] for actual_beam in actual_hypotheses]
+        actual_scores = [[hyp.score for hyp in actual_beam] for actual_beam in actual_hypotheses]
+        actual_timestamps = [[hyp.timestamp for hyp in actual_beam] for actual_beam in actual_hypotheses]
+
+        # transcribe with use implementation with cuda graphs
+        decoding_config["beam"]["allow_cuda_graphs"] = True
+        model.change_decoding_strategy(decoding_config)
+        model.decoding.decoding.decoding_computer.force_cuda_graphs_mode(mode=force_mode)
+
+        cudagraph_hypotheses = model.transcribe(test_audio_filenames, batch_size=batch_size, num_workers=None)
+        cudagraph_transcripts = [[hyp.text for hyp in cudagraphs_beam] for cudagraphs_beam in cudagraph_hypotheses]
+        cudagraph_scores = [[hyp.score for hyp in cudagraph_beam] for cudagraph_beam in cudagraph_hypotheses]
+        cudagraph_timestamps = [[hyp.timestamp for hyp in cudagraph_beam] for cudagraph_beam in cudagraph_hypotheses]
+
+        for batch_idx in range(min(batch_size, len(test_audio_filenames))):
+            assert len(actual_transcripts[batch_idx]) == len(cudagraph_transcripts[batch_idx])
+            assert cudagraph_scores[batch_idx] == pytest.approx(
+                actual_scores[batch_idx], abs=1e-2
+            ), f"Scores mismatch for batch_idx {batch_idx}"
+            assert (
+                cudagraph_timestamps[batch_idx] == actual_timestamps[batch_idx]
+            ), f"Timestamps mismatch for batch_idx {batch_idx}"
+
+            for actual, fast in zip(actual_transcripts[batch_idx], cudagraph_transcripts[batch_idx]):
+                ref_words = actual.split()
+                hyp_words = fast.split()
+                wer = edit_distance(ref_words, hyp_words)['total'] / max(len(ref_words), 1)
+
+                assert wer <= 1e-3, "Cuda graph beam decoder should match original decoder implementation."
+
+                if actual != fast:
+                    print("Erroneous samples in batch:", batch_idx)
+                    print("Original transcript:", actual)
+                    print("New transcript:", fast)
+
+    @pytest.mark.with_downloads
+    @pytest.mark.skipif(
+        not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()), reason="CUDA decoder can run only on CUDA"
+    )
+    @pytest.mark.parametrize("model_type", ["rnnt", "tdt"])
+    def test_stated_stateless_bf16(self, test_audio_filenames, rnnt_model, tdt_model, model_type):
+        """
+        Checks that we are able to run without errors all decodings in bfloat16.
+        Computational errors accumulate, so just checking if algorithms run without errors
+        """
+        batch_size = 16
+        device = torch.device("cuda")
+        model = rnnt_model.to(device) if model_type == "rnnt" else tdt_model.to(device)
+        decoding_config = copy.deepcopy(model.cfg.decoding)
+
+        # checking pytorch implementation
+        with open_dict(decoding_config):
+            decoding_config["strategy"] = "malsd_batch"
+            decoding_config["beam"]["beam_size"] = 4
+            decoding_config["beam"]["return_best_hypothesis"] = False
+            decoding_config["beam"]["allow_cuda_graphs"] = False
+
+        model.change_decoding_strategy(decoding_config)
+
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
+            model.transcribe(test_audio_filenames, batch_size=batch_size, num_workers=None)
+
+        modes = ["no_graphs", "no_while_loops", "full_graph"]
+        for force_mode in modes:
+            if force_mode == "full_graph":
+                skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+
+            # transcribe with use implementation with cuda graphs
+            decoding_config["beam"]["allow_cuda_graphs"] = True
+            model.change_decoding_strategy(decoding_config)
+            model.decoding.decoding.decoding_computer.force_cuda_graphs_mode(mode=force_mode)
+
+            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
+                model.transcribe(test_audio_filenames, batch_size=batch_size, num_workers=None)
+
+
+class TestCTCDecoding:
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"allow_cuda_graphs": False},
+            {"allow_cuda_graphs": True},
+        ],
+    )
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("batch_size", [4, 16])
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_ctc_beam_decoding_return_best_hypothesis(
+        self, test_audio_filenames, ctc_model, get_ctc_output, beam_config, device, batch_size, beam_size
+    ):
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = ctc_model.to(device)
+        log_probs, encoded_lengths = get_ctc_output
+        log_probs, encoded_lengths = log_probs[:num_samples].to(device), encoded_lengths[:num_samples].to(device)
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedCTCInfer(
+            blank_index=vocab_size,
+            beam_size=beam_size,
+            return_best_hypothesis=True,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy="beam_batch",
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            hyps = decoding(decoder_output=log_probs, decoder_lengths=encoded_lengths)[0]
+
+            check_res_best_hyps(num_samples, hyps)
+            hyps = decode_text_from_hypotheses(hyps, model)
+            print_res_best_hyps(hyps)
+
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"allow_cuda_graphs": False},
+            {"allow_cuda_graphs": True},
+        ],
+    )
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("batch_size", [4])
+    def test_ctc_beam_decoding_return_nbest(
+        self, test_audio_filenames, ctc_model, get_ctc_output, beam_config, device, beam_size, batch_size
+    ):
+        device = torch.device("cuda")
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = ctc_model.to(device)
+        log_probs, encoded_lengths = get_ctc_output
+        log_probs, encoded_lengths = log_probs[:num_samples].to(device), encoded_lengths[:num_samples].to(device)
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedCTCInfer(
+            blank_index=vocab_size,
+            beam_size=beam_size,
+            return_best_hypothesis=False,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy="beam_batch",
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            batch_nbest_hyps = decoding(decoder_output=log_probs, decoder_lengths=encoded_lengths)[0]
+
+            check_res_nbest_hyps(num_samples, batch_nbest_hyps)
+            batch_nbest_hyps = decode_text_from_nbest_hypotheses(batch_nbest_hyps, model)
+            print_res_nbest_hyps(batch_nbest_hyps)
+
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"allow_cuda_graphs": False, "ngram_lm_alpha": 0.3, "beam_beta": 1.0},
+            {"allow_cuda_graphs": False, "ngram_lm_alpha": 0.3, "beam_beta": 1.0},
+        ],
+    )
+    @pytest.mark.parametrize("batch_size", [4])
+    @pytest.mark.parametrize("beam_size", [4])
+    def test_ctc_beam_decoding_kenlm(
+        self,
+        kenlm_model_path,
+        test_audio_filenames,
+        ctc_model,
+        get_ctc_output,
+        beam_config,
+        device,
+        batch_size,
+        beam_size,
+    ):
+        device = torch.device("cuda")
+        beam_config["ngram_lm_model"] = kenlm_model_path
+
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = ctc_model.to(device)
+        decoder_output, decoder_lengths = get_ctc_output
+        decoder_output, decoder_lengths = decoder_output[:num_samples].to(device), decoder_lengths[:num_samples].to(
+            device
+        )
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedCTCInfer(
+            blank_index=vocab_size,
+            beam_size=beam_size,
+            return_best_hypothesis=True,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy="beam_batch",
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            hyps = decoding(decoder_output=decoder_output, decoder_lengths=decoder_lengths)[0]
+
+            check_res_best_hyps(num_samples, hyps)
+            hyps = decode_text_from_hypotheses(hyps, model)
+            print_res_best_hyps(hyps)

tests/collections/asr/decoding/test_batched_hyps.py

@@ -0,0 +1,797 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+import pytest
+import torch
+
+from nemo.collections.asr.parts.submodules.transducer_decoding.batched_hyps import BatchedHyps
+from nemo.collections.asr.parts.utils.rnnt_utils import batched_hyps_to_hypotheses
+from tests.collections.asr.decoding.utils import avoid_sync_operations
+
+DEVICES: List[torch.device] = [torch.device("cpu")]
+
+if torch.cuda.is_available():
+    DEVICES.append(torch.device("cuda"))
+
+if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    DEVICES.append(torch.device("mps"))
+
+# blank id that does not collide with any non-blank label used in the "no blank steps" tests
+NON_COLLIDING_BLANK_ID = 1024
+
+
+class TestBatchedHyps:
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_instantiate(self, device: torch.device):
+        hyps = BatchedHyps(batch_size=2, init_length=3, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        assert torch.is_tensor(hyps.timestamps)
+        # device: for mps device we need to use `type`, not directly compare
+        assert hyps.timestamps.device.type == device.type
+        assert hyps.timestamps.shape == (2, 3)
+        assert hyps.transcript.shape == (2, 3)
+        assert hyps.scores.shape == (2,)
+        assert hyps.current_lengths.tolist() == [0, 0]
+        # optional storage is disabled by default
+        assert hyps.token_durations is None
+        assert hyps.step_confidence is None
+        assert hyps.logits is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("batch_size", [-1, 0])
+    def test_instantiate_incorrect_batch_size(self, batch_size):
+        with pytest.raises(ValueError):
+            _ = BatchedHyps(batch_size=batch_size, init_length=3, blank_id=0)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("init_length", [-1, 0])
+    def test_instantiate_incorrect_init_length(self, init_length):
+        with pytest.raises(ValueError):
+            _ = BatchedHyps(batch_size=1, init_length=init_length, blank_id=0)
+
+    @pytest.mark.unit
+    def test_instantiate_with_logits_requires_logits_dim(self):
+        # `with_logits=True` without `logits_dim` is invalid
+        with pytest.raises(ValueError):
+            _ = BatchedHyps(batch_size=1, init_length=3, blank_id=0, with_logits=True, logits_dim=None)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_instantiate_optional_storage(self, device: torch.device):
+        logits_dim = 7
+        hyps = BatchedHyps(
+            batch_size=2,
+            init_length=3,
+            blank_id=0,
+            logits_dim=logits_dim,
+            device=device,
+            with_durations=True,
+            with_step_confidence=True,
+            with_duration_confidence=True,
+            with_logits=True,
+        )
+        assert hyps.token_durations.shape == (2, 3)
+        # duration confidence makes the confidence tensor store a pair (step + duration) per token
+        assert hyps.step_confidence.shape == (2, 3, 2)
+        assert hyps.logits.shape == (2, 3, logits_dim)
+
+        # without duration confidence the confidence tensor is 2d
+        hyps_no_dur_conf = BatchedHyps(
+            batch_size=2,
+            init_length=3,
+            blank_id=0,
+            device=device,
+            with_step_confidence=True,
+            with_duration_confidence=False,
+        )
+        assert hyps_no_dur_conf.step_confidence.shape == (2, 3)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_results_masked(self, device: torch.device):
+        # batch of size 2, add label for first utterance only (second is inactive)
+        hyps = BatchedHyps(batch_size=2, init_length=1, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, False], device=device),
+            labels=torch.tensor([5, 1], device=device),
+            time_indices=torch.tensor([1, 0], device=device),
+            scores=torch.tensor([0.5, 10.0], device=device),
+        )
+        assert hyps.current_lengths.tolist() == [1, 0]
+        assert hyps.transcript.tolist()[0][:1] == [5]
+        assert hyps.timestamps.tolist()[0][:1] == [1]
+        assert hyps.scores.tolist() == pytest.approx([0.5, 0.0])  # inactive score should be ignored!
+        assert hyps.last_nb_timestamp.tolist() == [1, -1]
+        assert hyps.last_nb_timestamp_lasts.tolist() == [1, 0]
+        assert hyps.last_nb_labels.tolist() == [5, -1]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_multiple_results_masked(self, device: torch.device):
+        # batch of size 2, add label for first utterance, then add labels for both utterances
+        hyps = BatchedHyps(batch_size=2, init_length=1, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, False], device=device),
+            labels=torch.tensor([5, 2], device=device),
+            time_indices=torch.tensor([1, 0], device=device),
+            scores=torch.tensor([0.5, 10.0], device=device),
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([2, 4], device=device),
+            time_indices=torch.tensor([1, 2], device=device),
+            scores=torch.tensor([1.0, 1.0], device=device),
+        )
+        assert hyps.current_lengths.tolist() == [2, 1]
+        assert hyps.transcript.tolist()[0][:2] == [5, 2]
+        assert hyps.transcript.tolist()[1][:1] == [4]
+        assert hyps.timestamps.tolist()[0][:2] == [1, 1]
+        assert hyps.timestamps.tolist()[1][:1] == [2]
+        assert hyps.scores.tolist() == pytest.approx([1.5, 1.0])
+        assert hyps.last_nb_timestamp.tolist() == [1, 2]
+        assert hyps.last_nb_timestamp_lasts.tolist() == [2, 1]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_results_masked_no_checks(self, device: torch.device):
+        # `check_lengths=False` must contain no host<->device synchronization (blocking) operations
+        hyps = BatchedHyps(batch_size=2, init_length=4, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        active_mask = torch.tensor([True, False], device=device)
+        time_indices = torch.tensor([1, 0], device=device)
+        scores = torch.tensor([0.5, 10.0], device=device)
+        labels = torch.tensor([5, 1], device=device)
+        # check there are no blocking operations
+        with avoid_sync_operations(device=device):
+            hyps.add_results_masked_(
+                active_mask=active_mask,
+                labels=labels,
+                time_indices=time_indices,
+                scores=scores,
+                check_lengths=False,
+            )
+        assert hyps.current_lengths.tolist() == [1, 0]
+        assert hyps.transcript.tolist()[0][:1] == [5]
+        assert hyps.timestamps.tolist()[0][:1] == [1]
+        assert hyps.scores.tolist() == pytest.approx([0.5, 0.0])  # inactive score should be ignored!
+        assert hyps.last_nb_timestamp.tolist() == [1, -1]
+        assert hyps.last_nb_timestamp_lasts.tolist() == [1, 0]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_results_masked_reallocates(self, device: torch.device):
+        # init_length is intentionally small; storage must grow transparently when check_lengths=True
+        hyps = BatchedHyps(batch_size=2, init_length=1, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        for step in range(5):
+            hyps.add_results_masked_(
+                active_mask=torch.tensor([True, True], device=device),
+                labels=torch.tensor([step, step + 10], device=device),
+                time_indices=torch.tensor([step, step], device=device),
+                scores=torch.tensor([1.0, 1.0], device=device),
+                check_lengths=True,
+            )
+        assert hyps._max_length >= 5
+        assert hyps.current_lengths.tolist() == [5, 5]
+        assert hyps.transcript.tolist()[0][:5] == [0, 1, 2, 3, 4]
+        assert hyps.transcript.tolist()[1][:5] == [10, 11, 12, 13, 14]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_results_masked_with_blank_steps(self, device: torch.device):
+        # with_blank_steps=True: blank labels are stored in the transcript, but they do NOT advance
+        # the score / last-non-blank tracking. Single utterance for clarity.
+        blank_id = 0
+        hyps = BatchedHyps(batch_size=1, init_length=2, blank_id=blank_id, device=device, with_blank_steps=True)
+        # (label, time, score): two non-blank tokens at t=0, then blank, then non-blank at t=1, then blank
+        steps = [
+            (3, 0, 1.0),  # non-blank
+            (5, 0, 1.5),  # non-blank, same timestamp
+            (blank_id, 0, 0.1),  # blank
+            (7, 1, 2.0),  # non-blank
+            (blank_id, 1, 0.2),  # blank
+        ]
+        for label, time, score in steps:
+            hyps.add_results_masked_(
+                active_mask=torch.tensor([True], device=device),
+                labels=torch.tensor([label], device=device),
+                time_indices=torch.tensor([time], device=device),
+                scores=torch.tensor([score], device=device),
+            )
+        # all steps (including blanks) are stored
+        assert hyps.current_lengths.tolist() == [5]
+        assert hyps.transcript.tolist()[0][:5] == [3, 5, blank_id, 7, blank_id]
+        assert hyps.timestamps.tolist()[0][:5] == [0, 0, 0, 1, 1]
+        # only non-blank scores accumulate: 1.0 + 1.5 + 2.0
+        assert hyps.scores.tolist() == pytest.approx([4.5])
+        # last-non-blank tracking ignores blanks
+        assert hyps.last_nb_timestamp.tolist() == [1]
+        assert hyps.last_nb_timestamp_lasts.tolist() == [1]
+        assert hyps.last_nb_labels.tolist() == [7]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_get_data_without_blank_no_blank_steps(self, device: torch.device):
+        # with_blank_steps=False: data is returned as-is (it never contained blanks)
+        hyps = BatchedHyps(batch_size=2, init_length=2, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([5, 4], device=device),
+            time_indices=torch.tensor([0, 1], device=device),
+            scores=torch.tensor([1.0, 1.0], device=device),
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, False], device=device),
+            labels=torch.tensor([2, 0], device=device),
+            time_indices=torch.tensor([1, 0], device=device),
+            scores=torch.tensor([1.0, 0.0], device=device),
+        )
+        lengths, transcript, timestamps, durations, confidence = hyps.get_data_without_blank()
+        # returned objects are the underlying (unmodified) tensors
+        assert lengths is hyps.current_lengths
+        assert transcript is hyps.transcript
+        assert timestamps is hyps.timestamps
+        assert durations is None
+        assert confidence is None
+        assert lengths.tolist() == [2, 1]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_get_data_without_blank_with_blank_steps(self, device: torch.device):
+        # with_blank_steps=True: blanks are stripped and non-blank order is preserved
+        blank_id = 0
+        hyps = BatchedHyps(batch_size=2, init_length=2, blank_id=blank_id, device=device, with_blank_steps=True)
+        # seq 0: [5, blank, 2, blank]  -> [5, 2]
+        # seq 1: [blank, 4]            -> [4]
+        steps = [
+            # (labels, times, active_mask)
+            ([5, blank_id], [0, 0], [True, True]),
+            ([blank_id, 4], [0, 1], [True, True]),
+            ([2, blank_id], [1, 1], [True, False]),
+            ([blank_id, 0], [1, 0], [True, False]),
+        ]
+        for labels, times, active in steps:
+            hyps.add_results_masked_(
+                active_mask=torch.tensor(active, device=device),
+                labels=torch.tensor(labels, device=device),
+                time_indices=torch.tensor(times, device=device),
+                scores=torch.tensor([1.0, 1.0], device=device),
+            )
+        lengths, transcript, timestamps, durations, confidence = hyps.get_data_without_blank()
+        assert lengths.tolist() == [2, 1]
+        assert transcript[0, :2].tolist() == [5, 2]
+        assert transcript[1, :1].tolist() == [4]
+        assert timestamps[0, :2].tolist() == [0, 1]
+        assert timestamps[1, :1].tolist() == [1]
+        assert durations is None
+        assert confidence is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_get_last_labels(self, device: torch.device):
+        hyps = BatchedHyps(batch_size=2, init_length=2, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        # no labels yet -> pad_id everywhere
+        assert hyps.get_last_labels(pad_id=-1).tolist() == [-1, -1]
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, False], device=device),
+            labels=torch.tensor([5, 1], device=device),
+            time_indices=torch.tensor([0, 0], device=device),
+            scores=torch.tensor([1.0, 0.0], device=device),
+        )
+        assert hyps.get_last_labels(pad_id=-1).tolist() == [5, -1]
+        assert hyps.get_last_labels(pad_id=100).tolist() == [5, 100]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_clear(self, device: torch.device):
+        hyps = BatchedHyps(batch_size=2, init_length=2, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([5, 4], device=device),
+            time_indices=torch.tensor([0, 0], device=device),
+            scores=torch.tensor([1.0, 1.0], device=device),
+        )
+        hyps.clear_()
+        assert hyps.current_lengths.tolist() == [0, 0]
+        assert hyps.scores.tolist() == pytest.approx([0.0, 0.0])
+        assert hyps.transcript.tolist() == [[0, 0], [0, 0]]
+        assert hyps.last_nb_timestamp.tolist() == [-1, -1]
+        assert hyps.last_nb_timestamp_lasts.tolist() == [0, 0]
+        assert hyps.last_nb_labels.tolist() == [-1, -1]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_clone(self, device: torch.device):
+        logits_dim = 7
+        hyps = BatchedHyps(
+            batch_size=2,
+            init_length=2,
+            blank_id=0,
+            logits_dim=logits_dim,
+            device=device,
+            with_durations=True,
+            with_step_confidence=True,
+            with_logits=True,
+            with_blank_steps=True,
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([5, 4], device=device),
+            time_indices=torch.tensor([0, 0], device=device),
+            scores=torch.tensor([1.0, 1.0], device=device),
+            token_durations=torch.tensor([1, 2], device=device),
+            confidence=torch.tensor([0.9, 0.8], device=device),
+            logits=torch.rand((2, logits_dim), device=device),
+        )
+        clone = hyps.clone()
+        # flags carried over
+        assert clone.with_durations and clone.with_step_confidence and clone.with_logits and clone.with_blank_steps
+        assert clone.blank_id == hyps.blank_id
+        # values copied
+        assert clone.current_lengths.tolist() == hyps.current_lengths.tolist()
+        assert torch.equal(clone.transcript, hyps.transcript)
+        assert torch.allclose(clone.logits, hyps.logits)
+        assert torch.allclose(clone.step_confidence, hyps.step_confidence)
+        assert torch.equal(clone.token_durations, hyps.token_durations)
+        # clone is independent of the original
+        hyps.transcript.fill_(0)
+        hyps.scores.fill_(0.0)
+        assert clone.transcript.tolist()[0][:1] == [5]
+        assert clone.scores.tolist() == pytest.approx([1.0, 1.0])
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_merge(self, device: torch.device):
+        # merge two batched hypotheses (basic case: no blank steps, no optional storage)
+        def build(labels_per_step, times_per_step, masks_per_step, scores_per_step):
+            hyps = BatchedHyps(batch_size=2, init_length=4, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+            for labels, times, mask, scores in zip(labels_per_step, times_per_step, masks_per_step, scores_per_step):
+                hyps.add_results_masked_(
+                    active_mask=torch.tensor(mask, device=device),
+                    labels=torch.tensor(labels, device=device),
+                    time_indices=torch.tensor(times, device=device),
+                    scores=torch.tensor(scores, device=device),
+                )
+            return hyps
+
+        # A: seq0=[5, 2], seq1=[4]
+        hyps_a = build(
+            labels_per_step=[[5, 4], [2, 0]],
+            times_per_step=[[0, 0], [1, 0]],
+            masks_per_step=[[True, True], [True, False]],
+            scores_per_step=[[0.5, 0.7], [0.5, 0.0]],
+        )
+        # B: seq0=[7], seq1=[8, 9]
+        hyps_b = build(
+            labels_per_step=[[7, 8], [0, 9]],
+            times_per_step=[[2, 2], [0, 3]],
+            masks_per_step=[[True, True], [False, True]],
+            scores_per_step=[[0.3, 0.3], [0.0, 0.4]],
+        )
+
+        hyps_a.merge_(hyps_b)
+        assert hyps_a.current_lengths.tolist() == [3, 3]
+        assert hyps_a.transcript[0, :3].tolist() == [5, 2, 7]
+        assert hyps_a.transcript[1, :3].tolist() == [4, 8, 9]
+        assert hyps_a.scores.tolist() == pytest.approx([1.3, 1.4])
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_merge_with_logits(self, device: torch.device):
+        # Regression test: merge_ must use dim=1 (sequence axis) for logits scatter/cat,
+        # not dim=-1 (logits-dim axis), which would silently corrupt the data.
+        logits_dim = 5
+        blank_id = NON_COLLIDING_BLANK_ID
+
+        def build_with_logits(labels_per_step, times_per_step, masks_per_step, scores_per_step, logits_per_step):
+            hyps = BatchedHyps(
+                batch_size=2,
+                init_length=4,
+                blank_id=blank_id,
+                logits_dim=logits_dim,
+                device=device,
+                float_dtype=torch.float32,
+                with_logits=True,
+            )
+            for labels, times, mask, scores, logits in zip(
+                labels_per_step, times_per_step, masks_per_step, scores_per_step, logits_per_step
+            ):
+                hyps.add_results_masked_(
+                    active_mask=torch.tensor(mask, device=device),
+                    labels=torch.tensor(labels, device=device),
+                    time_indices=torch.tensor(times, device=device),
+                    scores=torch.tensor(scores, device=device),
+                    logits=logits,
+                )
+            return hyps
+
+        # Fixed logits so we can assert exact values after merge
+        logits_a0 = torch.full((2, logits_dim), 1.0, device=device)  # step for seq0=[5], seq1=[4]
+        logits_a1 = torch.full((2, logits_dim), 2.0, device=device)  # step for seq0=[2], seq1 inactive
+
+        logits_b0 = torch.full((2, logits_dim), 3.0, device=device)  # step for seq0=[7], seq1=[8]
+        logits_b1 = torch.full((2, logits_dim), 4.0, device=device)  # step for seq0 inactive, seq1=[9]
+
+        # A: seq0=[5, 2], seq1=[4]
+        hyps_a = build_with_logits(
+            labels_per_step=[[5, 4], [2, 0]],
+            times_per_step=[[0, 0], [1, 0]],
+            masks_per_step=[[True, True], [True, False]],
+            scores_per_step=[[0.5, 0.7], [0.5, 0.0]],
+            logits_per_step=[logits_a0, logits_a1],
+        )
+        # B: seq0=[7], seq1=[8, 9]
+        hyps_b = build_with_logits(
+            labels_per_step=[[7, 8], [0, 9]],
+            times_per_step=[[2, 2], [0, 3]],
+            masks_per_step=[[True, True], [False, True]],
+            scores_per_step=[[0.3, 0.3], [0.0, 0.4]],
+            logits_per_step=[logits_b0, logits_b1],
+        )
+
+        hyps_a.merge_(hyps_b)
+
+        # Sequence lengths: seq0=2+1=3, seq1=1+2=3
+        assert hyps_a.current_lengths.tolist() == [3, 3]
+
+        # seq0 logits: positions [0,1] from A (value=1 then 2), position [2] from B (value=3)
+        assert torch.allclose(hyps_a.logits[0, 0], torch.full((logits_dim,), 1.0, device=device))
+        assert torch.allclose(hyps_a.logits[0, 1], torch.full((logits_dim,), 2.0, device=device))
+        assert torch.allclose(hyps_a.logits[0, 2], torch.full((logits_dim,), 3.0, device=device))
+
+        # seq1 logits: position [0] from A (value=1), positions [1,2] from B (value=3 then 4)
+        assert torch.allclose(hyps_a.logits[1, 0], torch.full((logits_dim,), 1.0, device=device))
+        assert torch.allclose(hyps_a.logits[1, 1], torch.full((logits_dim,), 3.0, device=device))
+        assert torch.allclose(hyps_a.logits[1, 2], torch.full((logits_dim,), 4.0, device=device))
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_merge_with_logits_triggers_reallocation(self, device: torch.device):
+        # When the combined length exceeds _max_length, merge_ must reallocate logits along dim=1
+        # (sequence axis). With the dim=-1 bug, the reallocation would expand the logits-dim
+        # axis instead, causing a shape mismatch on the subsequent scatter_.
+        logits_dim = 5
+        blank_id = NON_COLLIDING_BLANK_ID
+
+        # Use init_length=1 to force reallocation during merge
+        hyps_a = BatchedHyps(
+            batch_size=1,
+            init_length=1,
+            blank_id=blank_id,
+            logits_dim=logits_dim,
+            device=device,
+            float_dtype=torch.float32,
+            with_logits=True,
+        )
+        hyps_b = BatchedHyps(
+            batch_size=1,
+            init_length=1,
+            blank_id=blank_id,
+            logits_dim=logits_dim,
+            device=device,
+            float_dtype=torch.float32,
+            with_logits=True,
+        )
+
+        logits_a = torch.full((1, logits_dim), 1.0, device=device)
+        logits_b = torch.full((1, logits_dim), 2.0, device=device)
+
+        hyps_a.add_results_masked_(
+            active_mask=torch.tensor([True], device=device),
+            labels=torch.tensor([5], device=device),
+            time_indices=torch.tensor([0], device=device),
+            scores=torch.tensor([1.0], device=device),
+            logits=logits_a,
+        )
+        hyps_b.add_results_masked_(
+            active_mask=torch.tensor([True], device=device),
+            labels=torch.tensor([7], device=device),
+            time_indices=torch.tensor([1], device=device),
+            scores=torch.tensor([1.0], device=device),
+            logits=logits_b,
+        )
+
+        # cur_len=1, other_len=1 -> combined=2 >= init_length=1, so reallocation is triggered
+        hyps_a.merge_(hyps_b)
+
+        assert hyps_a.current_lengths.tolist() == [2]
+        assert hyps_a.logits.shape == (1, hyps_a._max_length, logits_dim)
+        assert torch.allclose(hyps_a.logits[0, 0], torch.full((logits_dim,), 1.0, device=device))
+        assert torch.allclose(hyps_a.logits[0, 1], torch.full((logits_dim,), 2.0, device=device))
+
+
+class TestConvertToHypotheses:
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_no_blank_steps(self, device: torch.device):
+        # with_blank_steps=False: transcript already contains only non-blank labels
+        hyps = BatchedHyps(batch_size=2, init_length=1, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, False], device=device),
+            labels=torch.tensor([5, 0], device=device),
+            time_indices=torch.tensor([1, 0], device=device),
+            scores=torch.tensor([0.5, 0.0], device=device),
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([2, 4], device=device),
+            time_indices=torch.tensor([1, 2], device=device),
+            scores=torch.tensor([1.0, 1.0], device=device),
+        )
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        assert (hypotheses[0].y_sequence == torch.tensor([5, 2], device="cpu")).all()
+        assert (hypotheses[1].y_sequence == torch.tensor([4], device="cpu")).all()
+        assert hypotheses[0].score == pytest.approx(1.5)
+        assert hypotheses[1].score == pytest.approx(1.0)
+        assert (hypotheses[0].timestamp == torch.tensor([1, 1], device="cpu")).all()
+        assert (hypotheses[1].timestamp == torch.tensor([2], device="cpu")).all()
+        # no blank steps -> no alignments / frame confidence
+        assert hypotheses[0].alignments is None
+        assert hypotheses[1].alignments is None
+        assert hypotheses[0].frame_confidence is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_batch_size_arg(self, device: torch.device):
+        # batch_size arg returns only the first `batch_size` hypotheses (CUDA-graph constant batch)
+        hyps = BatchedHyps(batch_size=4, init_length=1, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True, True, True], device=device),
+            labels=torch.tensor([5, 4, 3, 2], device=device),
+            time_indices=torch.tensor([0, 0, 0, 0], device=device),
+            scores=torch.tensor([1.0, 1.0, 1.0, 1.0], device=device),
+        )
+        hypotheses = batched_hyps_to_hypotheses(hyps, batch_size=2)
+        assert len(hypotheses) == 2
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_with_blank_steps_strips_blanks(self, device: torch.device):
+        # with_blank_steps=True but no logits/confidence: blanks must be stripped from y_sequence/timestamps,
+        # while alignments are NOT produced (no logits recorded)
+        blank_id = 0
+        hyps = BatchedHyps(batch_size=2, init_length=2, blank_id=blank_id, device=device, with_blank_steps=True)
+        # seq 0: [5, blank, 2, blank] -> [5, 2]
+        # seq 1: [blank, 4]           -> [4]
+        steps = [
+            ([5, blank_id], [0, 0], [True, True], [0.5, 0.1]),
+            ([blank_id, 4], [0, 1], [True, True], [0.1, 1.0]),
+            ([2, blank_id], [1, 1], [True, False], [1.0, 0.0]),
+            ([blank_id, 0], [1, 0], [True, False], [0.1, 0.0]),
+        ]
+        for labels, times, active, scores in steps:
+            hyps.add_results_masked_(
+                active_mask=torch.tensor(active, device=device),
+                labels=torch.tensor(labels, device=device),
+                time_indices=torch.tensor(times, device=device),
+                scores=torch.tensor(scores, device=device),
+            )
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        assert (hypotheses[0].y_sequence == torch.tensor([5, 2], device="cpu")).all()
+        assert (hypotheses[1].y_sequence == torch.tensor([4], device="cpu")).all()
+        assert (hypotheses[0].timestamp == torch.tensor([0, 1], device="cpu")).all()
+        assert (hypotheses[1].timestamp == torch.tensor([1], device="cpu")).all()
+        # only non-blank scores accumulated
+        assert hypotheses[0].score == pytest.approx(1.5)
+        assert hypotheses[1].score == pytest.approx(1.0)
+        # no logits recorded -> alignments stay None even though blank steps were stored
+        assert hypotheses[0].alignments is None
+        assert hypotheses[1].alignments is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_logits_no_alignments_without_blank_steps(self, device: torch.device):
+        # logits are recorded but with_blank_steps=False -> alignments must NOT be produced
+        logits_dim = 7
+        hyps = BatchedHyps(
+            batch_size=2,
+            init_length=2,
+            blank_id=6,
+            logits_dim=logits_dim,
+            device=device,
+            with_logits=True,
+            with_blank_steps=False,
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([5, 4], device=device),
+            time_indices=torch.tensor([0, 0], device=device),
+            scores=torch.tensor([1.0, 1.0], device=device),
+            logits=torch.rand((2, logits_dim), device=device),
+        )
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        assert hypotheses[0].alignments is None
+        assert hypotheses[1].alignments is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_with_blank_steps_and_logits_alignments(self, device: torch.device):
+        # Reproduces alignment recovery: with_blank_steps=True + with_logits=True
+        batch_size = 2
+        logits_dim = 7
+        blank_index = 6
+        hyps = BatchedHyps(
+            batch_size=batch_size,
+            init_length=1,
+            blank_id=blank_index,
+            logits_dim=logits_dim,
+            device=device,
+            with_logits=True,
+            with_blank_steps=True,
+        )
+        # sequence 0: [[5, blank], [2, blank]] -> [5, 2]
+        # sequence 1: [[blank   ], [4, blank]] -> [4]
+        # one logits row per (batch, add-call); rows belonging to inactive entries are ignored
+        L = [torch.rand((batch_size, logits_dim), device=device) for _ in range(4)]
+        # call0: seq0=5@t0, seq1=blank@t0
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([5, blank_index], device=device),
+            time_indices=torch.tensor([0, 0], device=device),
+            scores=torch.tensor([0.5, 0.1], device=device),
+            logits=L[0],
+        )
+        # call1: seq0=blank@t0, seq1=4@t1
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([blank_index, 4], device=device),
+            time_indices=torch.tensor([0, 1], device=device),
+            scores=torch.tensor([0.1, 1.0], device=device),
+            logits=L[1],
+        )
+        # call2: seq0=2@t1, seq1=blank@t1
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([2, blank_index], device=device),
+            time_indices=torch.tensor([1, 1], device=device),
+            scores=torch.tensor([1.0, 0.1], device=device),
+            logits=L[2],
+        )
+        # call3: seq0=blank@t1, seq1 inactive
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, False], device=device),
+            labels=torch.tensor([blank_index, 0], device=device),
+            time_indices=torch.tensor([1, 0], device=device),
+            scores=torch.tensor([0.1, 0.0], device=device),
+            logits=L[3],
+        )
+
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        assert (hypotheses[0].y_sequence == torch.tensor([5, 2], device="cpu")).all()
+        assert (hypotheses[1].y_sequence == torch.tensor([4], device="cpu")).all()
+        assert hypotheses[0].score == pytest.approx(1.5)
+        assert hypotheses[1].score == pytest.approx(1.0)
+        assert (hypotheses[0].timestamp == torch.tensor([0, 1], device="cpu")).all()
+        assert (hypotheses[1].timestamp == torch.tensor([1], device="cpu")).all()
+
+        # alignments are grouped by timestamp; each entry is a (logits, label) tuple
+        etalon = [
+            [
+                [(L[0][0].cpu(), 5), (L[1][0].cpu(), blank_index)],
+                [(L[2][0].cpu(), 2), (L[3][0].cpu(), blank_index)],
+            ],
+            [
+                [(L[0][1].cpu(), blank_index)],
+                [(L[1][1].cpu(), 4), (L[2][1].cpu(), blank_index)],
+            ],
+        ]
+        for batch_i in range(batch_size):
+            assert len(hypotheses[batch_i].alignments) == len(etalon[batch_i])
+            for t, group_for_timestamp in enumerate(etalon[batch_i]):
+                assert len(hypotheses[batch_i].alignments[t]) == len(group_for_timestamp)
+                for step, (current_logits, label) in enumerate(group_for_timestamp):
+                    assert torch.allclose(hypotheses[batch_i].alignments[t][step][0], current_logits)
+                    assert hypotheses[batch_i].alignments[t][step][1] == label
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_with_durations(self, device: torch.device):
+        # TDT-style: token durations are stored and (with blank steps) stripped together with blanks
+        blank_id = 0
+        hyps = BatchedHyps(
+            batch_size=1,
+            init_length=2,
+            blank_id=blank_id,
+            device=device,
+            with_durations=True,
+            with_blank_steps=True,
+        )
+        # transcript [3, blank, 7, blank] with durations [2, 1, 4, 1] -> [3, 7] with durations [2, 4]
+        steps = [
+            (3, 0, 2, 1.0),
+            (blank_id, 2, 1, 0.1),
+            (7, 3, 4, 2.0),
+            (blank_id, 7, 1, 0.2),
+        ]
+        for label, time, duration, score in steps:
+            hyps.add_results_masked_(
+                active_mask=torch.tensor([True], device=device),
+                labels=torch.tensor([label], device=device),
+                time_indices=torch.tensor([time], device=device),
+                scores=torch.tensor([score], device=device),
+                token_durations=torch.tensor([duration], device=device),
+            )
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        assert (hypotheses[0].y_sequence == torch.tensor([3, 7], device="cpu")).all()
+        assert (hypotheses[0].timestamp == torch.tensor([0, 3], device="cpu")).all()
+        assert (hypotheses[0].token_duration == torch.tensor([2, 4], device="cpu")).all()
+        assert hypotheses[0].score == pytest.approx(3.0)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_with_step_confidence_no_blank_steps(self, device: torch.device):
+        # with_blank_steps=False: per-token confidence is precomputed, no frame_confidence (no blank steps)
+        hyps = BatchedHyps(
+            batch_size=1,
+            init_length=2,
+            blank_id=NON_COLLIDING_BLANK_ID,
+            device=device,
+            with_step_confidence=True,
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True], device=device),
+            labels=torch.tensor([5], device=device),
+            time_indices=torch.tensor([0], device=device),
+            scores=torch.tensor([1.0], device=device),
+            confidence=torch.tensor([0.9], device=device),
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True], device=device),
+            labels=torch.tensor([2], device=device),
+            time_indices=torch.tensor([1], device=device),
+            scores=torch.tensor([1.0], device=device),
+            confidence=torch.tensor([0.8], device=device),
+        )
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        assert hypotheses[0].non_blank_step_confidence_precomputed == pytest.approx([0.9, 0.8])
+        assert hypotheses[0].frame_confidence is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_with_step_confidence_and_blank_steps(self, device: torch.device):
+        # with_blank_steps=True: frame_confidence is grouped per timestamp (incl. blanks),
+        # while non_blank_step_confidence_precomputed holds only non-blank tokens
+        blank_id = 0
+        hyps = BatchedHyps(
+            batch_size=1,
+            init_length=2,
+            blank_id=blank_id,
+            device=device,
+            with_step_confidence=True,
+            with_blank_steps=True,
+        )
+        # transcript [3, blank, 7], confidence [0.9, 0.5, 0.8], timestamps [0, 0, 1]
+        steps = [
+            (3, 0, 0.9, 1.0),
+            (blank_id, 0, 0.5, 0.1),
+            (7, 1, 0.8, 2.0),
+        ]
+        for label, time, confidence, score in steps:
+            hyps.add_results_masked_(
+                active_mask=torch.tensor([True], device=device),
+                labels=torch.tensor([label], device=device),
+                time_indices=torch.tensor([time], device=device),
+                scores=torch.tensor([score], device=device),
+                confidence=torch.tensor([confidence], device=device),
+            )
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        # non-blank tokens only
+        assert hypotheses[0].non_blank_step_confidence_precomputed == pytest.approx([0.9, 0.8])
+        # grouped by timestamp: t=0 has 2 steps (token + blank), t=1 has 1 step
+        frame_confidence = hypotheses[0].frame_confidence
+        assert len(frame_confidence) == 2
+        assert len(frame_confidence[0]) == 2
+        assert len(frame_confidence[1]) == 1
+        flat = [float(c) for group in frame_confidence for c in group]
+        assert flat == pytest.approx([0.9, 0.5, 0.8])

tests/collections/asr/decoding/test_biasing_multi_model.py

@@ -0,0 +1,279 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.collections.asr.parts.context_biasing.biasing_multi_model import (
+    GPUBiasingMultiModel,
+    GPUBiasingMultiModelReference,
+)
+from nemo.collections.asr.parts.context_biasing.boosting_graph_batched import (
+    BoostingTreeModelConfig,
+    GPUBoostingTreeModel,
+)
+from nemo.core.utils.optional_libs import TRITON_AVAILABLE
+
+DEVICES = [torch.device("cpu")]
+if torch.cuda.is_available():
+    DEVICES.append(torch.device("cuda"))
+if hasattr(torch, "mps") and torch.mps.is_available():
+    DEVICES.append(torch.device("mps"))
+
+# Triton only works on CUDA, so only test use_triton=True if Triton is available
+USE_TRITON_OPTIONS = [False, True] if TRITON_AVAILABLE else [False]
+
+
+def create_boosting_model(phrases: list[str], tokenizer, device: torch.device) -> GPUBoostingTreeModel:
+    """Helper to create boosting model from phrases"""
+    cfg = BoostingTreeModelConfig(key_phrases_list=phrases, context_score=1.0)
+    model = GPUBoostingTreeModel.from_config(cfg, tokenizer=tokenizer)
+    return model.to(device)
+
+
+class TestGPUBiasingMultiModel:
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_models_incremental(self, stt_en_conformer_transducer_small, device: torch.device):
+        """Test adding 2 boosting models one-by-one, verifying arcs and states are correctly merged."""
+        tokenizer = stt_en_conformer_transducer_small.tokenizer
+        vocab_size = tokenizer.vocab_size
+
+        # Create empty multi-model
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size).to(device)
+
+        # Initially empty
+        assert multi_model.num_models == 0
+        assert multi_model.has_models() is False
+        assert multi_model.num_states_total == 0
+        assert multi_model.num_arcs_extended_total == 0
+
+        # Create and add first model
+        model1 = create_boosting_model(["hello", "world"], tokenizer, device)
+        model_id1 = multi_model.add_model(model1, alpha=1.0)
+
+        # Verify after first model
+        assert model_id1 == 0
+        assert multi_model.num_models == 1
+        assert multi_model.has_models() is True
+        assert multi_model.model2active[model_id1].item() is True
+        assert multi_model.num_states_total == model1.num_states
+        assert multi_model.num_arcs_extended_total == model1.num_arcs_extended
+        assert multi_model.model2num_states[model_id1].item() == model1.num_states
+        assert multi_model.model2num_arcs_extended[model_id1].item() == model1.num_arcs_extended
+
+        # Create and add second model
+        model2 = create_boosting_model(["test", "one", "two"], tokenizer, device)
+        model_id2 = multi_model.add_model(model2, alpha=1.5)
+
+        # Verify after second model
+        assert model_id2 == 1
+        assert multi_model.num_models == 2
+        assert multi_model.has_models() is True
+        assert multi_model.model2active[model_id1].item() is True
+        assert multi_model.model2active[model_id2].item() is True
+        assert multi_model.num_states_total == model1.num_states + model2.num_states
+        assert multi_model.num_arcs_extended_total == model1.num_arcs_extended + model2.num_arcs_extended
+
+        # Verify offsets
+        assert multi_model.model2states_offset[model_id1].item() == 0
+        assert multi_model.model2states_offset[model_id2].item() == model1.num_states
+        assert multi_model.model2arcs_offset[model_id1].item() == 0
+        assert multi_model.model2arcs_offset[model_id2].item() == model1.num_arcs_extended
+
+        # Verify init states work
+        init_states = multi_model.get_init_states(batch_size=4, bos=True)
+        assert init_states.shape == (4,)
+        assert init_states.device.type == device.type
+
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_then_remove_model(self, stt_en_conformer_transducer_small, device: torch.device):
+        """Test adding 2 models then removing the first one."""
+        tokenizer = stt_en_conformer_transducer_small.tokenizer
+        vocab_size = tokenizer.vocab_size
+
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size).to(device)
+
+        # Add two models
+        model1 = create_boosting_model(["alpha", "beta"], tokenizer, device)
+        model2 = create_boosting_model(["gamma", "delta"], tokenizer, device)
+
+        model_id1 = multi_model.add_model(model1, alpha=1.0)
+        model_id2 = multi_model.add_model(model2, alpha=2.0)
+
+        # Store counts before removal
+        model1_num_states = model1.num_states
+        model1_num_arcs = model1.num_arcs_extended
+        total_states_before = multi_model.num_states_total
+        total_arcs_before = multi_model.num_arcs_extended_total
+
+        assert multi_model.model2active[model_id1].item() is True
+        assert multi_model.model2active[model_id2].item() is True
+
+        # Remove first model
+        multi_model.remove_model(model_id1)
+
+        # Verify removal
+        assert model_id1 in multi_model.free_ids
+        assert multi_model.model2active[model_id1].item() is False
+        assert multi_model.model2active[model_id2].item() is True
+        assert multi_model.model2alpha[model_id1].item() == 0.0
+        assert multi_model.model2alpha[model_id2].item() == 2.0
+
+        # Verify state/arc counts decreased
+        assert multi_model.num_states_total == total_states_before - model1_num_states
+        assert multi_model.num_arcs_extended_total == total_arcs_before - model1_num_arcs
+
+        # Verify model2 offset updated (shifted left)
+        assert multi_model.model2states_offset[model_id2].item() == 0
+        assert multi_model.model2arcs_offset[model_id2].item() == 0
+
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_model_id_reuse(self, stt_en_conformer_transducer_small, device):
+        """Test that removed model IDs are reused."""
+        tokenizer = stt_en_conformer_transducer_small.tokenizer
+        vocab_size = tokenizer.vocab_size
+
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size).to(device)
+
+        # Add model1 -> id=0
+        model1 = create_boosting_model(["first"], tokenizer, device)
+        model_id1 = multi_model.add_model(model1)
+        assert model_id1 == 0
+
+        # Add model2 -> id=1
+        model2 = create_boosting_model(["second"], tokenizer, device)
+        model_id2 = multi_model.add_model(model2)
+        assert model_id2 == 1
+
+        # Remove model1
+        multi_model.remove_model(model_id1)
+        assert model_id1 in multi_model.free_ids
+
+        # Add model3 -> should reuse id=0
+        model3 = create_boosting_model(["third"], tokenizer, device)
+        model_id3 = multi_model.add_model(model3)
+        assert model_id3 == model_id1  # Reused ID
+        assert model_id1 not in multi_model.free_ids  # No longer free
+
+        # Verify model3 is active
+        assert multi_model.model2active[model_id3].item() is True
+
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("batch_size", [1, 4])
+    @pytest.mark.parametrize("use_triton", USE_TRITON_OPTIONS)
+    @pytest.mark.parametrize("bos", [True, False])
+    def test_advance_matches_reference(
+        self, stt_en_conformer_transducer_small, device: torch.device, batch_size: int, use_triton, bos: bool
+    ):
+        """Verify GPUBiasingMultiModel produces same scores/states as reference implementation."""
+        tokenizer = stt_en_conformer_transducer_small.tokenizer
+        vocab_size = tokenizer.vocab_size
+
+        # Create both implementations
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size, use_triton=use_triton).to(device)
+        reference = GPUBiasingMultiModelReference(vocab_size=vocab_size).to(device)
+
+        # Create boosting models with same phrases
+        phrases1 = ["hello world", "test"]
+        phrases2 = ["neural", "network"]
+
+        model1_mm = create_boosting_model(phrases1, tokenizer, device)
+        model1_ref = create_boosting_model(phrases1, tokenizer, device)
+        model2_mm = create_boosting_model(phrases2, tokenizer, device)
+        model2_ref = create_boosting_model(phrases2, tokenizer, device)
+
+        # Add models to both with same alpha values
+        alpha1, alpha2 = 1.0, 1.5
+        model_id1_mm = multi_model.add_model(model1_mm, alpha=alpha1)
+        model_id1_ref = reference.add_model(model1_ref, alpha=alpha1)
+        model_id2_mm = multi_model.add_model(model2_mm, alpha=alpha2)
+        model_id2_ref = reference.add_model(model2_ref, alpha=alpha2)
+
+        assert model_id1_mm == model_id1_ref
+        assert model_id2_mm == model_id2_ref
+
+        # Get initial states
+        states_mm = multi_model.get_init_states(batch_size=batch_size, bos=bos)
+        states_ref = reference.get_init_states(batch_size=batch_size, bos=bos)
+
+        # Create model_ids tensor with alternating models
+        model_ids = torch.tensor(
+            [model_id1_mm if i % 2 == 0 else model_id2_mm for i in range(batch_size)],
+            dtype=torch.long,
+            device=device,
+        )
+
+        # Call advance on both
+        scores_mm, next_states_mm = multi_model.advance(states_mm, model_ids)
+        scores_ref, next_states_ref = reference.advance(states_ref, model_ids)
+
+        # Verify shapes
+        assert scores_mm.shape == (batch_size, vocab_size)
+        assert next_states_mm.shape == (batch_size, vocab_size)
+        assert scores_ref.shape == (batch_size, vocab_size)
+        assert next_states_ref.shape == (batch_size, vocab_size)
+
+        # Verify scores and states match
+        assert torch.allclose(
+            scores_mm, scores_ref, atol=1e-5
+        ), f"Scores mismatch: max diff = {(scores_mm - scores_ref).abs().max()}"
+        assert torch.equal(next_states_mm, next_states_ref), "Next states mismatch"
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_empty_multi_model(self, device: torch.device):
+        """Test behavior of empty multi-model."""
+        vocab_size = 100
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size, use_triton=False).to(device)
+
+        # Verify empty state
+        assert multi_model.has_models() is False
+        assert multi_model.num_models == 0
+        assert multi_model.num_states_total == 0
+        assert multi_model.num_arcs_extended_total == 0
+
+        # get_init_states should work and return START_STATE
+        init_states = multi_model.get_init_states(batch_size=4, bos=True)
+        assert init_states.shape == (4,)
+        assert (init_states == GPUBiasingMultiModel.START_STATE).all()
+
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_get_alphas(self, stt_en_conformer_transducer_small, device: torch.device):
+        """Per-stream alpha lookup returns model weight or 0 for invalid ids."""
+        tokenizer = stt_en_conformer_transducer_small.tokenizer
+        vocab_size = tokenizer.vocab_size
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size).to(device)
+        model1 = create_boosting_model(["hello"], tokenizer, device)
+        model2 = create_boosting_model(["world"], tokenizer, device)
+        model_id1 = multi_model.add_model(model1, alpha=1.0)
+        model_id2 = multi_model.add_model(model2, alpha=2.5)
+
+        model_ids = torch.tensor([model_id1, model_id2, -1, model_id1], device=device, dtype=torch.long)
+        alphas = multi_model.get_alphas(model_ids)
+
+        assert alphas.shape == (4,)
+        assert alphas[0].item() == pytest.approx(1.0)
+        assert alphas[1].item() == pytest.approx(2.5)
+        assert alphas[2].item() == pytest.approx(0.0)
+        assert alphas[3].item() == pytest.approx(1.0)

tests/collections/asr/decoding/test_ctc_decoding.py

@@ -0,0 +1,487 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import os
+from functools import cached_property, lru_cache
+from pathlib import Path
+
+import pytest
+import torch
+from kaldialign import edit_distance
+from omegaconf import DictConfig, open_dict
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.mixins import mixins
+from nemo.collections.asr.parts.submodules.ctc_decoding import (
+    CTCBPEDecoding,
+    CTCBPEDecodingConfig,
+    CTCDecoding,
+    CTCDecodingConfig,
+)
+from nemo.collections.asr.parts.submodules.ngram_lm.ngram_lm_batched import NGramGPULanguageModel
+from nemo.collections.asr.parts.utils.asr_confidence_utils import ConfidenceConfig
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.core.utils.cuda_python_utils import skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported
+from tests.collections.asr.decoding.test_timestamps import BaseTimestampsTest
+
+
+@pytest.fixture(scope="module")
+def audio_file(test_data_dir):
+    return os.path.join(test_data_dir, "asr/test/an4/wav/cen3-mjwl-b.wav")
+
+
+CTC_MODEL = "nvidia/stt_en_conformer_ctc_small"
+
+
+@pytest.fixture(scope="module")
+def kenlm_model_path(tmp_path_factory, test_data_dir):
+    lm_path = Path(test_data_dir) / "asr/kenlm_ngram_lm/parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+    assert os.path.exists(lm_path), f"LM file not found: {lm_path}"
+    lm_nemo_path = tmp_path_factory.mktemp("lm") / f"{lm_path.name}.nemo"
+    NGramGPULanguageModel.from_file(lm_path, vocab_size=1024).save_to(f"{lm_nemo_path}")
+    return f"{lm_nemo_path}"
+
+
+@pytest.fixture(scope="module")
+def ctc_model():
+    model = ASRModel.from_pretrained(model_name=CTC_MODEL, map_location="cpu")
+    model.eval()
+    return model
+
+
+def char_vocabulary():
+    return [' ', 'a', 'b', 'c', 'd', 'e', 'f', '.']
+
+
+@pytest.fixture()
+@lru_cache(maxsize=8)
+def tmp_tokenizer(test_data_dir):
+    cfg = DictConfig({'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'})
+
+    class _TmpASRBPE(mixins.ASRBPEMixin):
+        def register_artifact(self, _, vocab_path):
+            return vocab_path
+
+    asrbpe = _TmpASRBPE()
+    asrbpe._setup_tokenizer(cfg)
+    return asrbpe.tokenizer
+
+
+class TestCTCDecoding:
+    @pytest.mark.unit
+    def test_constructor(self):
+        cfg = CTCDecodingConfig()
+        vocab = char_vocabulary()
+        decoding = CTCDecoding(decoding_cfg=cfg, vocabulary=vocab)
+        assert decoding is not None
+
+    @pytest.mark.unit
+    def test_constructor_subword(self, tmp_tokenizer):
+        cfg = CTCBPEDecodingConfig()
+        decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+        assert decoding is not None
+
+    @pytest.mark.unit
+    def test_char_decoding_greedy_forward(
+        self,
+    ):
+        cfg = CTCDecodingConfig(strategy='greedy')
+        vocab = char_vocabulary()
+        decoding = CTCDecoding(decoding_cfg=cfg, vocabulary=vocab)
+
+        B, T = 4, 20
+        V = len(char_vocabulary()) + 1
+        input_signal = torch.randn(size=(B, T, V))
+        length = torch.randint(low=1, high=T, size=[B])
+
+        with torch.no_grad():
+            hypotheses = decoding.ctc_decoder_predictions_tensor(
+                input_signal, length, fold_consecutive=True, return_hypotheses=False
+            )
+            texts = [hyp.text for hyp in hypotheses]
+
+            for text in texts:
+                assert isinstance(text, str)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('alignments', [False, True])
+    @pytest.mark.parametrize('timestamps', [False, True])
+    def test_char_decoding_greedy_forward_hypotheses(self, alignments, timestamps):
+        cfg = CTCDecodingConfig(strategy='greedy', preserve_alignments=alignments, compute_timestamps=timestamps)
+        vocab = char_vocabulary()
+        decoding = CTCDecoding(decoding_cfg=cfg, vocabulary=vocab)
+
+        B, T = 4, 20
+        V = len(char_vocabulary()) + 1
+        input_signal = torch.randn(size=(B, T, V))
+        length = torch.randint(low=1, high=T, size=[B])
+
+        with torch.no_grad():
+            hyps = decoding.ctc_decoder_predictions_tensor(
+                input_signal, length, fold_consecutive=True, return_hypotheses=True
+            )
+
+            for idx, hyp in enumerate(hyps):
+                assert isinstance(hyp, Hypothesis)
+                assert torch.is_tensor(hyp.y_sequence)
+                assert isinstance(hyp.text, str)
+
+                # alignments check
+                if alignments:
+                    assert hyp.alignments is not None
+                    assert isinstance(hyp.alignments, tuple)
+                    assert len(hyp.alignments[0]) == length[idx]
+                    assert len(hyp.alignments[1]) == length[idx]
+
+                # timestamps check
+                if timestamps:
+                    BaseTimestampsTest.check_char_timestamps(hyp, decoding)
+
+    @pytest.mark.unit
+    def test_subword_decoding_greedy_forward(self, tmp_tokenizer):
+        cfg = CTCBPEDecodingConfig(strategy='greedy')
+        decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+
+        B, T = 4, 20
+        V = decoding.tokenizer.tokenizer.vocab_size + 1
+        input_signal = torch.randn(size=(B, T, V))
+        length = torch.randint(low=1, high=T, size=[B])
+
+        with torch.no_grad():
+            hypotheses = decoding.ctc_decoder_predictions_tensor(
+                input_signal, length, fold_consecutive=True, return_hypotheses=False
+            )
+            texts = [hyp.text for hyp in hypotheses]
+
+            for text in texts:
+                assert isinstance(text, str)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('alignments', [False, True])
+    @pytest.mark.parametrize('timestamps', [False, True])
+    @pytest.mark.pleasefixme
+    def test_subword_decoding_greedy_forward_hypotheses(self, tmp_tokenizer, alignments, timestamps):
+        cfg = CTCBPEDecodingConfig(strategy='greedy', preserve_alignments=alignments, compute_timestamps=timestamps)
+        decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+
+        B, T = 4, 20
+        V = decoding.tokenizer.tokenizer.vocab_size + 1
+        input_signal = torch.randn(size=(B, T, V))
+        length = torch.randint(low=1, high=T, size=[B])
+
+        with torch.no_grad():
+            hyps = decoding.ctc_decoder_predictions_tensor(
+                input_signal, length, fold_consecutive=True, return_hypotheses=True
+            )
+
+            for idx, hyp in enumerate(hyps):
+                assert isinstance(hyp, Hypothesis)
+                assert torch.is_tensor(hyp.y_sequence)
+                assert isinstance(hyp.text, str)
+
+                # alignments check
+                if alignments:
+                    assert hyp.alignments is not None
+                    assert isinstance(hyp.alignments, tuple)
+                    assert len(hyp.alignments[0]) == length[idx]
+                    assert len(hyp.alignments[1]) == length[idx]
+
+                # timestamps check
+                if timestamps:
+                    BaseTimestampsTest.check_subword_timestamps(hyp, decoding)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('alignments', [False, True])
+    @pytest.mark.parametrize('timestamps', [False, True])
+    @pytest.mark.parametrize('preserve_frame_confidence', [False, True])
+    @pytest.mark.parametrize('length_is_none', [False, True])
+    @pytest.mark.parametrize(
+        "logprobs_device",
+        [
+            torch.device("cpu"),
+            pytest.param(
+                torch.device("cuda"),
+                marks=pytest.mark.skipif(
+                    not torch.cuda.is_available(),
+                    reason='CUDA required for test.',
+                ),
+            ),
+        ],
+    )
+    @pytest.mark.parametrize(
+        "length_device",
+        [
+            torch.device("cpu"),
+            pytest.param(
+                torch.device("cuda"),
+                marks=pytest.mark.skipif(
+                    not torch.cuda.is_available(),
+                    reason='CUDA required for test.',
+                ),
+            ),
+        ],
+    )
+    def test_batched_decoding_logprobs(
+        self,
+        tmp_tokenizer,
+        alignments,
+        timestamps,
+        preserve_frame_confidence,
+        length_is_none,
+        logprobs_device,
+        length_device,
+    ):
+        cfg = CTCBPEDecodingConfig(
+            strategy='greedy',
+            preserve_alignments=alignments,
+            compute_timestamps=timestamps,
+            confidence_cfg=ConfidenceConfig(preserve_frame_confidence=preserve_frame_confidence),
+        )
+        unbatched_decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+
+        cfg.strategy = 'greedy_batch'
+        batched_decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+
+        torch.manual_seed(1)
+        B, T = 4, 20
+        V = unbatched_decoding.tokenizer.tokenizer.vocab_size + 1
+        input_signal = torch.randn(size=(B, T, V), device=logprobs_device)
+        # Set the blank index to a very high probability to make sure
+        # that we always handle at least a few blanks.
+        input_signal[:, 0, unbatched_decoding.tokenizer.tokenizer.vocab_size] = 1000
+        input_signal[:, 1, unbatched_decoding.tokenizer.tokenizer.vocab_size] = 1000
+        if length_is_none:
+            length = None
+        else:
+            length = torch.randint(low=1, high=T, size=[B], device=length_device)
+
+        with torch.inference_mode():
+            hyps = unbatched_decoding.ctc_decoder_predictions_tensor(
+                input_signal, length, fold_consecutive=True, return_hypotheses=True
+            )
+
+            batched_hyps = batched_decoding.ctc_decoder_predictions_tensor(
+                input_signal, length, fold_consecutive=True, return_hypotheses=True
+            )
+
+            assert len(hyps) == len(batched_hyps) == B
+            for hyp, batched_hyp in zip(hyps, batched_hyps):
+                assert torch.abs(hyp.score - batched_hyp.score) <= 1e-5
+                assert torch.all(hyp.y_sequence == batched_hyp.y_sequence)
+                if timestamps:
+                    assert hyp.timestamp == batched_hyp.timestamp
+                if alignments:
+                    assert torch.all(hyp.alignments[0] == batched_hyp.alignments[0])
+                    assert torch.all(hyp.alignments[1] == batched_hyp.alignments[1])
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('timestamps', [False, True])
+    @pytest.mark.parametrize('length_is_none', [False, True])
+    @pytest.mark.parametrize(
+        "labels_device",
+        [
+            torch.device("cpu"),
+            pytest.param(
+                torch.device("cuda"),
+                marks=pytest.mark.skipif(
+                    not torch.cuda.is_available(),
+                    reason='CUDA required for test.',
+                ),
+            ),
+        ],
+    )
+    @pytest.mark.parametrize(
+        "length_device",
+        [
+            torch.device("cpu"),
+            pytest.param(
+                torch.device("cuda"),
+                marks=pytest.mark.skipif(
+                    not torch.cuda.is_available(),
+                    reason='CUDA required for test.',
+                ),
+            ),
+        ],
+    )
+    def test_batched_decoding_labels(self, tmp_tokenizer, timestamps, length_is_none, labels_device, length_device):
+        cfg = CTCBPEDecodingConfig(strategy='greedy', compute_timestamps=timestamps)
+        unbatched_decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+        cfg.strategy = 'greedy_batch'
+        batched_decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+
+        torch.manual_seed(1)
+        B, T = 4, 20
+        V = unbatched_decoding.tokenizer.tokenizer.vocab_size + 1
+        input_labels = torch.randint(V, size=(B, T), device=labels_device)
+        # Set some indices to blank to make sure that we always handle
+        # at least a few blanks.
+        input_labels[:, 0] = unbatched_decoding.tokenizer.tokenizer.vocab_size
+        input_labels[:, 1] = unbatched_decoding.tokenizer.tokenizer.vocab_size
+        if length_is_none:
+            length = None
+        else:
+            length = torch.randint(low=1, high=T, size=[B], device=length_device)
+
+        with torch.inference_mode():
+            hyps = unbatched_decoding.ctc_decoder_predictions_tensor(
+                input_labels, length, fold_consecutive=True, return_hypotheses=True
+            )
+
+            batched_hyps = batched_decoding.ctc_decoder_predictions_tensor(
+                input_labels, length, fold_consecutive=True, return_hypotheses=True
+            )
+
+            assert len(hyps) == len(batched_hyps) == B
+            for hyp, batched_hyp in zip(hyps, batched_hyps):
+                assert abs(hyp.score - batched_hyp.score) <= 1e-5
+                assert torch.all(hyp.y_sequence == batched_hyp.y_sequence)
+                if timestamps:
+                    assert hyp.timestamp == batched_hyp.timestamp
+
+
+class TestCTCTimestamps(BaseTimestampsTest):
+    """CTC-specific timestamp tests that inherit from BaseTimestampsTest"""
+
+    @cached_property
+    def decoding_char(self):
+        cfg = CTCDecodingConfig()
+        vocab = char_vocabulary()
+        decoding = CTCDecoding(decoding_cfg=cfg, vocabulary=vocab)
+        return decoding
+
+    @cached_property
+    def decoding_subword_wpe(self):
+        cfg = CTCBPEDecodingConfig(compute_timestamps=True)
+        decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=self.tmp_tokenizer)
+        return decoding
+
+    @cached_property
+    def decoding_subword_bpe(self):
+        cfg = CTCBPEDecodingConfig(compute_timestamps=True)
+        decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=self.bpe_tokenizer)
+        return decoding
+
+    @pytest.mark.unit
+    def test_word_offsets_subword_wpe(self, tmp_tokenizer):
+        self.tmp_tokenizer = tmp_tokenizer
+        super().test_word_offsets_subword_wpe()
+
+    @pytest.mark.unit
+    def test_word_offsets_subword_wpe_other_delimiter(self, tmp_tokenizer):
+        self.tmp_tokenizer = tmp_tokenizer
+        super().test_word_offsets_subword_wpe_other_delimiter()
+
+
+class TestCTCGreedyDecodingWithNGPU_LM:
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    def test_ctc_decoding_gpulm(
+        self,
+        audio_file,
+        kenlm_model_path,
+        ctc_model,
+    ):
+        device = torch.device("cuda")
+        model = ctc_model.to(device)
+
+        gt_hyp = model.transcribe([audio_file], num_workers=None)
+
+        decoding_config = copy.deepcopy(model.cfg.decoding)
+        with open_dict(model.decoding.cfg) as cfg:
+            cfg.greedy["ngram_lm_model"] = kenlm_model_path
+            cfg.greedy["ngram_lm_alpha"] = 0.0
+            model.change_decoding_strategy(cfg)
+        lm_hyp = model.transcribe([audio_file], num_workers=None)
+
+        assert gt_hyp[0].text == lm_hyp[0].text
+        assert abs(gt_hyp[0].score - lm_hyp[0].score) <= 1e-3
+
+        with open_dict(model.decoding.cfg) as cfg:
+            cfg.greedy["ngram_lm_model"] = kenlm_model_path
+            cfg.greedy["ngram_lm_alpha"] = 10.0
+        model.change_decoding_strategy(cfg)
+        lm_hyp = model.transcribe([audio_file], num_workers=None)
+        assert gt_hyp[0].text != lm_hyp[0].text
+        assert abs(gt_hyp[0].score - lm_hyp[0].score) > 1e-3
+
+        model.change_decoding_strategy(decoding_config)
+
+
+class TestCTCGreedyDecodingCudaGrpahs:
+    """
+    Tests CudaGraphs implementations from CTC models greedy decoding
+    """
+
+    @pytest.mark.with_downloads
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA decoder can run only on CUDA")
+    @pytest.mark.parametrize("force_mode", ["no_graphs", "no_while_loops", "full_graph"])
+    def test_stated_stateless(self, audio_file, kenlm_model_path, ctc_model, force_mode: str):
+        """
+        Compares pure Pytorch and with three modes of statefull implementations for double floating point precision.
+            1. Pure pytorch, but statefull implementation: no_graphs
+            2. With CudaGrpahs: no_while_loops and full_graph.
+        """
+        if force_mode == "full_graph":
+            skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+
+        device = torch.device("cuda")
+        model = ctc_model.to(device)
+        decoding_config = copy.deepcopy(model.cfg.decoding)
+
+        with open_dict(model.decoding.cfg) as cfg:
+            cfg.greedy["ngram_lm_model"] = kenlm_model_path
+            cfg.greedy["ngram_lm_alpha"] = 0.2
+            cfg.greedy["allow_cuda_graphs"] = False
+
+            model.change_decoding_strategy(cfg)
+
+        actual_hypotheses = model.transcribe([audio_file], num_workers=None)
+        actual_transcripts = [hyp.text for hyp in actual_hypotheses]
+        actual_scores = [hyp.score for hyp in actual_hypotheses]
+        actual_timestamps = [hyp.timestamp for hyp in actual_hypotheses]
+
+        # transcribe with use implementation with cuda graphs
+        model.decoding.cfg["greedy"]["allow_cuda_graphs"] = True
+        model.change_decoding_strategy(model.decoding.cfg)
+        model.decoding.decoding.force_cuda_graphs_mode(mode=force_mode)
+
+        cudagraph_hypotheses = model.transcribe([audio_file], num_workers=None)
+        cudagraph_transcripts = [hyp.text for hyp in cudagraph_hypotheses]
+        cudagraph_scores = [hyp.score for hyp in cudagraph_hypotheses]
+        cudagraph_timestamps = [hyp.timestamp for hyp in cudagraph_hypotheses]
+
+        for batch_idx in range(len(actual_transcripts)):
+            assert len(actual_transcripts[batch_idx]) == len(cudagraph_transcripts[batch_idx])
+            assert cudagraph_scores[batch_idx] == pytest.approx(
+                actual_scores[batch_idx], abs=1e-2
+            ), f"Scores mismatch for batch_idx {batch_idx}"
+            assert (
+                cudagraph_timestamps[batch_idx] == actual_timestamps[batch_idx]
+            ), f"Timestamps mismatch for batch_idx {batch_idx}"
+
+            ref_words = actual_transcripts[batch_idx].split()
+            hyp_words = cudagraph_transcripts[batch_idx].split()
+            wer = edit_distance(ref_words, hyp_words)['total'] / max(len(ref_words), 1)
+
+            assert wer <= 1e-3, "Cuda graph greedy decoder should match original decoder implementation."
+
+            for actual, fast in zip(actual_transcripts[batch_idx], cudagraph_transcripts[batch_idx]):
+                if actual != fast:
+                    print("Erroneous samples in batch:", batch_idx)
+                    print("Original transcript:", actual)
+                    print("New transcript:", fast)
+
+        model.change_decoding_strategy(decoding_config)

tests/collections/asr/decoding/test_cuda_graph_rnnt_greedy_decoding.py

@@ -0,0 +1,368 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import copy
+import glob
+import types
+
+import lightning.pytorch as ptl
+import pytest
+import torch
+from kaldialign import edit_distance
+from omegaconf import DictConfig, open_dict
+
+from nemo.core.config.pytorch_lightning import TrainerConfig
+from nemo.core.utils.cuda_python_utils import skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported
+
+
+# These tests move the model to CUDA before calling transcribe(), so avoid forking DataLoader workers afterwards.
+CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS = 0
+
+
+def test_forced_full_graph_compile_does_not_fallback():
+    from nemo.collections.asr.parts.submodules.transducer_decoding.rnnt_label_looping import (
+        GreedyBatchedRNNTLabelLoopingComputer,
+    )
+
+    accelerator_error = getattr(torch, "AcceleratorError", RuntimeError)
+    computer = GreedyBatchedRNNTLabelLoopingComputer.__new__(GreedyBatchedRNNTLabelLoopingComputer)
+    computer.cuda_graphs_allow_fallback = False
+
+    with pytest.raises(RuntimeError, match="Full CUDA graph decoding failed"):
+        computer._raise_or_warn_no_while_loop_cuda_graphs(accelerator_error("CUDA error: invalid argument"))
+
+
+def test_conditional_node_restores_previous_stream_on_body_error(monkeypatch):
+    from nemo.core.utils import cuda_python_utils
+
+    if not cuda_python_utils.CUDA_PYTHON_AVAILABLE:
+        pytest.skip("cuda-python is required to test with_conditional_node")
+
+    class FakeStream:
+        def __init__(self, name):
+            self.name = name
+            self.cuda_stream = name
+
+    class FakeTorchCuda:
+        def __init__(self):
+            self.parent_stream = FakeStream("parent")
+            self.body_stream = FakeStream("body")
+            self.current = self.parent_stream
+            self.set_calls = []
+
+        def current_stream(self, device=None):
+            return self.current
+
+        def Stream(self, device=None):
+            return self.body_stream
+
+        def set_stream(self, stream):
+            self.current = stream
+            self.set_calls.append(stream)
+
+    class FakeCudart:
+        cudaStreamCaptureStatus = types.SimpleNamespace(cudaStreamCaptureStatusActive="active")
+        cudaStreamUpdateCaptureDependenciesFlags = types.SimpleNamespace(cudaStreamSetCaptureDependencies="set")
+        cudaStreamCaptureMode = types.SimpleNamespace(cudaStreamCaptureModeThreadLocal="thread_local")
+
+        def __init__(self):
+            self.ended_streams = []
+
+        def cudaStreamGetCaptureInfo(self, stream):
+            return ("active", None, "graph", ["dependency"])
+
+        def cudaStreamUpdateCaptureDependencies(self, *args):
+            return ()
+
+        def cudaStreamBeginCaptureToGraph(self, *args):
+            return ()
+
+        def cudaStreamEndCapture(self, stream):
+            self.ended_streams.append(stream)
+            return ()
+
+    class FakeCuda:
+        CUgraphNodeType = types.SimpleNamespace(CU_GRAPH_NODE_TYPE_CONDITIONAL="conditional")
+        CUgraphConditionalNodeType = types.SimpleNamespace(CU_GRAPH_COND_TYPE_WHILE="while")
+
+        class CUgraphNodeParams:
+            def __init__(self):
+                self.conditional = types.SimpleNamespace(phGraph_out=["body_graph"])
+
+        def cuGraphAddNode(self, *args):
+            return ("node",)
+
+        def cuCtxGetCurrent(self):
+            return ("ctx",)
+
+        def cuLaunchKernel(self, *args):
+            return ()
+
+    fake_torch_cuda = FakeTorchCuda()
+    fake_cudart = FakeCudart()
+    fake_args = types.SimpleNamespace(ctypes=types.SimpleNamespace(data=1234))
+    fake_handle = types.SimpleNamespace(getPtr=lambda: 5678)
+
+    monkeypatch.setattr(cuda_python_utils, "cu_call", lambda result: result)
+    monkeypatch.setattr(cuda_python_utils, "cuda", FakeCuda())
+    monkeypatch.setattr(cuda_python_utils, "cudart", fake_cudart)
+    monkeypatch.setattr(cuda_python_utils, "cuda_python_version", "13.0.0")
+    monkeypatch.setattr(cuda_python_utils.torch, "cuda", fake_torch_cuda)
+
+    with pytest.raises(RuntimeError, match="body failed"):
+        with cuda_python_utils.with_conditional_node("kernel", fake_args, fake_handle, device="cuda"):
+            assert fake_torch_cuda.current_stream(device="cuda") is fake_torch_cuda.body_stream
+            raise RuntimeError("body failed")
+
+    assert fake_torch_cuda.current_stream(device="cuda") is fake_torch_cuda.parent_stream
+    assert fake_torch_cuda.set_calls == [fake_torch_cuda.body_stream, fake_torch_cuda.parent_stream]
+    assert fake_cudart.ended_streams == ["body"]
+
+
+@pytest.mark.with_downloads
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA decoder can run only on CUDA")
+@pytest.mark.parametrize(
+    ("model_name", "batch_size", "enable_bfloat16"),
+    [
+        ("stt_en_fastconformer_transducer_large", 8, False),
+        ("stt_en_fastconformer_transducer_large", 8, True),
+    ],
+)
+@pytest.mark.parametrize("loop_labels", [False, True])
+def test_cuda_graph_rnnt_greedy_decoder(model_name, batch_size, enable_bfloat16, loop_labels: bool, request):
+    if not loop_labels:
+        skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+    if enable_bfloat16 and not torch.cuda.is_bf16_supported():
+        pytest.skip("bfloat16 is not supported")
+
+    device = torch.device("cuda")
+    nemo_model = request.getfixturevalue(model_name).to(device)
+    decoding_config = copy.deepcopy(nemo_model.cfg.decoding)
+
+    with open_dict(decoding_config):
+        decoding_config["greedy"]["max_symbols"] = 5
+        decoding_config["greedy"]["loop_labels"] = loop_labels
+        decoding_config["greedy"]["use_cuda_graph_decoder"] = False
+
+    nemo_model.change_decoding_strategy(decoding_config)
+    audio_filepaths = glob.glob("tests/.data/asr/test/an4/wav/*.wav")
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=enable_bfloat16):
+        actual_hypotheses = nemo_model.transcribe(
+            audio_filepaths, batch_size=batch_size, num_workers=CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS
+        )
+
+    actual_transcripts = [hyp.text for hyp in actual_hypotheses]
+    actual_y_sequences = [hyp.y_sequence for hyp in actual_hypotheses]
+
+    decoding_config["greedy"]["use_cuda_graph_decoder"] = True
+
+    nemo_model.change_decoding_strategy(decoding_config)
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=enable_bfloat16):
+        fast_hypotheses = nemo_model.transcribe(
+            audio_filepaths, batch_size=batch_size, num_workers=CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS
+        )
+
+    fast_transcripts = [hyp.text for hyp in fast_hypotheses]
+    fast_y_sequences = [hyp.y_sequence for hyp in fast_hypotheses]
+
+    total_dist = sum(
+        edit_distance(r.split(), h.split())['total'] for r, h in zip(actual_transcripts, fast_transcripts)
+    )
+    total_words = sum(len(r.split()) for r in actual_transcripts)
+    wer = total_dist / total_words if total_words > 0 else 0.0
+    y_sequence_eq = [torch.equal(act_y, fast_y) for (act_y, fast_y) in zip(actual_y_sequences, fast_y_sequences)]
+
+    assert wer <= 1e-3, "Cuda graph greedy decoder should match original decoder implementation."
+    assert all(y_sequence_eq), "Cuda graph greedy decoder should match original decoder implementation."
+
+    for actual, fast in zip(actual_transcripts, fast_transcripts):
+        if actual != fast:
+            print("erroneous samples:")
+            print("Original transcript:", actual)
+            print("New transcript:", fast)
+
+
+@pytest.mark.with_downloads
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA decoder can run only on CUDA")
+@pytest.mark.parametrize("force_mode", ["no_graphs", "no_while_loops", "full_graph"])
+@pytest.mark.parametrize("enable_bfloat16", [False, True])
+def test_loop_labels_cuda_graph_rnnt_greedy_decoder_forced_mode(
+    stt_en_fastconformer_transducer_large, force_mode: str, enable_bfloat16: bool
+):
+    """
+    Testing Label-Looping algorithm with CUDA graphs in forced mode.
+    This test guarantees that we check that the fallback behavior is working.
+    NB: Since it is impossible to directly debug CUDA graphs, when making changes,
+    start testing and debugging the code with forced "no_graphs" mode.
+    """
+    if enable_bfloat16 and not torch.cuda.is_bf16_supported():
+        pytest.skip("bfloat16 is not supported")
+
+    if force_mode == "full_graph":
+        skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+
+    batch_size = 16
+    device = torch.device("cuda")
+    nemo_model = stt_en_fastconformer_transducer_large.to(device)
+    decoding_config = copy.deepcopy(nemo_model.cfg.decoding)
+
+    with open_dict(decoding_config):
+        decoding_config["greedy"]["max_symbols"] = 5
+        decoding_config["greedy"]["loop_labels"] = True
+        decoding_config["greedy"]["use_cuda_graph_decoder"] = False
+        # test that alignments and confidence do not introduce failures
+        decoding_config["greedy"]["preserve_alignments"] = True
+        decoding_config["greedy"]["preserve_frame_confidence"] = True
+
+    nemo_model.change_decoding_strategy(decoding_config)
+    audio_filepaths = glob.glob("tests/.data/asr/test/an4/wav/*.wav")
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=enable_bfloat16):
+        actual_hypotheses = nemo_model.transcribe(
+            audio_filepaths, batch_size=batch_size, num_workers=CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS
+        )
+    actual_transcripts = [hyp.text for hyp in actual_hypotheses]
+
+    # transcribe with use implementation with cuda graphs
+    decoding_config["greedy"]["use_cuda_graph_decoder"] = True
+    nemo_model.change_decoding_strategy(decoding_config)
+    backup_cuda_graph_mode = nemo_model.decoding.decoding.decoding_computer.cuda_graphs_mode
+    try:
+        nemo_model.decoding.decoding.decoding_computer.force_cuda_graphs_mode(mode=force_mode)
+
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=enable_bfloat16):
+            fast_hypotheses = nemo_model.transcribe(
+                audio_filepaths, batch_size=batch_size, num_workers=CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS
+            )
+        fast_transcripts = [hyp.text for hyp in fast_hypotheses]
+
+        total_dist = sum(
+            edit_distance(r.split(), h.split())['total'] for r, h in zip(actual_transcripts, fast_transcripts)
+        )
+        total_words = sum(len(r.split()) for r in actual_transcripts)
+        wer = total_dist / total_words if total_words > 0 else 0.0
+
+        assert wer <= 1e-3, "Cuda graph greedy decoder should match original decoder implementation."
+
+        for actual, fast in zip(actual_transcripts, fast_transcripts):
+            if actual != fast:
+                print("erroneous samples:")
+                print("Original transcript:", actual)
+                print("New transcript:", fast)
+    finally:
+        nemo_model.decoding.decoding.decoding_computer.force_cuda_graphs_mode(mode=backup_cuda_graph_mode)
+
+
+@pytest.mark.with_downloads
+@pytest.mark.skipif(
+    not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
+    reason="Test requires CUDA device with bf16 support",
+)
+@pytest.mark.parametrize("is_tdt", [False, True])
+def test_loop_labels_cuda_graph_ddp_mixed_precision(
+    tmp_path_factory,
+    an4_train_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    is_tdt: bool,
+):
+    """CUDA graphs with DDP and mixed precision have bugs. We need to test that validation works with these settings."""
+    batch_size = 16
+
+    # instantiate trainer with bf16 mixed precision
+    trainer_cfg = TrainerConfig(devices=[0], accelerator="cuda", strategy="ddp", max_epochs=1, precision="bf16-mixed")
+    trainer = ptl.Trainer(**DictConfig(trainer_cfg))
+
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+
+    # setup validation data
+    val_ds_cfg = model.cfg.validation_ds
+    with open_dict(val_ds_cfg):
+        val_ds_cfg.manifest_filepath = [str(an4_train_manifest_corrected)]
+        val_ds_cfg.batch_size = batch_size
+        val_ds_cfg.is_tarred = False
+        val_ds_cfg.use_lhotse = False
+        if is_tdt:
+            # TDT model has config with missing mandatory values, this results in errors when setting up validation data
+            # we set all the mandatory values to dummy values
+            model.cfg.train_ds.tarred_audio_filepaths = None
+            model.cfg.train_ds.manifest_filepath = None
+            model.cfg.test_ds.manifest_filepath = None
+            model.cfg.tokenizer.dir = None
+
+    model.setup_multiple_validation_data(val_ds_cfg)
+
+    # validate using trainer
+    val_results = trainer.validate(model)
+    wer = val_results[0]["val_wer"]
+
+    # explicitly free resources, then test conditions
+    trainer._teardown()
+    # teardown from the trainer is not enough, and problem with CPU will still exist, related issue:
+    # https://github.com/Lightning-AI/pytorch-lightning/issues/18803)
+    # solution is to destroy torch process group explicitly
+    torch.distributed.destroy_process_group()
+
+    assert wer <= 0.1, f"WER is too high: {wer}"
+
+
+@pytest.mark.with_downloads
+@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() < 2, reason="Test requires 2 GPUs")
+@pytest.mark.parametrize("loop_labels", [False, True])
+def test_change_devices(loop_labels: bool, stt_en_fastconformer_transducer_large):
+    if not loop_labels:
+        skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+
+    first_device = torch.device("cuda:0")
+    second_device = torch.device("cuda:1")
+
+    batch_size = 8
+
+    nemo_model = stt_en_fastconformer_transducer_large.to(second_device)
+    decoding_config = copy.deepcopy(nemo_model.cfg.decoding)
+
+    with open_dict(decoding_config):
+        decoding_config["greedy"]["max_symbols"] = 5
+        decoding_config["greedy"]["loop_labels"] = loop_labels
+        decoding_config["greedy"]["use_cuda_graph_decoder"] = True
+
+    nemo_model.change_decoding_strategy(decoding_config)
+
+    # Test that the model can run successfully when it is first
+    # initialized on second_device and then transferred to
+    # true_device
+    nemo_model.to(first_device)
+    audio_filepaths = glob.glob("tests/.data/asr/test/an4/wav/*.wav")
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
+        second_device_hypotheses = nemo_model.transcribe(
+            audio_filepaths, batch_size=batch_size, num_workers=CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS
+        )
+    second_device_transcripts = [hyp.text for hyp in second_device_hypotheses]
+
+    # Test that the model can run successfully back on second_device
+    # after having been first run on first_device. Because the
+    # decoder's data structures are lazily initialized, this activates
+    # slightly different code than the first case (where the decoder
+    # has not run at all), so we want to exercise both cases.
+    nemo_model.to(second_device)
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
+        first_device_hypotheses = nemo_model.transcribe(
+            audio_filepaths, batch_size=batch_size, num_workers=CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS
+        )
+    first_device_transcripts = [hyp.text for hyp in first_device_hypotheses]
+    # Sanity check: The device we run on should not change execution
+    # output.
+    assert first_device_transcripts == second_device_transcripts

tests/collections/asr/decoding/test_multi_task_decoding.py

@@ -0,0 +1,342 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from nemo.collections.asr.modules.transformer.transformer import TransformerDecoderNM
+from nemo.collections.asr.modules.transformer.transformer_generators import (
+    BeamSearchSequenceGenerator,
+    BeamSearchSequenceGeneratorWithFusionModels,
+    GreedySequenceGenerator,
+)
+from nemo.collections.asr.parts.context_biasing import GPUBoostingTreeModel
+from nemo.collections.asr.parts.submodules.multitask_beam_decoding import TransformerAEDBeamInfer
+from nemo.collections.asr.parts.submodules.multitask_greedy_decoding import TransformerAEDGreedyInfer
+from nemo.collections.asr.parts.submodules.ngram_lm import NGramGPULanguageModel
+from nemo.collections.asr.parts.submodules.token_classifier import TokenClassifier
+
+
+@pytest.fixture()
+def deterministic_rng():
+    state = torch.get_rng_state()
+    torch.manual_seed(0)
+    yield
+    torch.set_rng_state(state)
+
+
+@pytest.fixture()
+def decoder_nm(deterministic_rng):
+    return TransformerDecoderNM(
+        vocab_size=8,
+        hidden_size=2,
+        num_layers=1,
+        inner_size=4,
+        num_attention_heads=1,
+        max_sequence_length=32,
+    ).eval()
+
+
+@pytest.fixture()
+def nnet(decoder_nm):
+    ans = (
+        decoder_nm.embedding,
+        decoder_nm.decoder,
+        TokenClassifier(hidden_size=2, num_classes=8),
+    )
+    ans = tuple(m.eval() for m in ans)
+    return ans
+
+
+@pytest.fixture()
+def inputs():
+    B, T, C = 1, 5, 2
+    return (
+        torch.tensor([[1]], dtype=torch.long),  # decoder_input_ids
+        torch.ones(B, T, C, dtype=torch.float),  # encoder_hidden_states
+        torch.ones(B, T, dtype=torch.float),  # encoder_input_mask
+    )
+
+
+@pytest.fixture()
+def tokenizer():
+    tok = Mock()
+    tok.pad = 0
+    tok.bos = 1
+    tok.eos = 2
+    return tok
+
+
+@pytest.mark.parametrize('with_confidence', [False, True])
+@pytest.mark.parametrize('return_xattn_scores', [False, True])
+def test_greedy_decoding(inputs, nnet, deterministic_rng, with_confidence, return_xattn_scores):
+    gen = GreedySequenceGenerator(
+        *nnet, return_xattn_scores=return_xattn_scores, preserve_step_confidence=with_confidence
+    )
+    output = gen(*inputs)
+
+    assert len(output) == 4
+    best_path, hypotheses, confidence, xattn_list = output
+
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+    assert best_path.shape == (1, 25)
+    if return_xattn_scores:
+        assert len(xattn_list) == len(nnet[1].layers)
+        assert xattn_list[0].shape == (1, 1, 24, 5)
+    else:
+        assert xattn_list is None
+
+    assert hypotheses is None
+
+    if with_confidence:
+        assert confidence is not None
+        assert torch.is_tensor(confidence)
+        assert confidence.shape == best_path.shape
+    else:
+        assert confidence is None
+
+
+@pytest.mark.parametrize('return_xattn_scores', [False, True])
+def test_temperature_sampling_decoding(inputs, nnet, return_xattn_scores):
+    gen = GreedySequenceGenerator(*nnet, return_xattn_scores=return_xattn_scores, temperature=10.0, n_samples=2)
+    output = gen(*inputs)
+
+    assert len(output) == 4
+    best_path, hypotheses, _, xatt_list = output
+
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+    assert best_path.shape[0] == 1
+
+    assert isinstance(hypotheses, list)
+    assert len(hypotheses) == 1
+    (seq0,) = hypotheses
+    assert seq0.shape[0] == 2
+    assert (seq0[0] != seq0[1]).any()
+
+    if return_xattn_scores:
+        assert len(xatt_list) == len(nnet[1].layers)
+        assert xatt_list[0].shape == (2, 1, 24, 5)
+    else:
+        assert xatt_list is None
+
+
+def test_beam_decoding_beam_scores_false(inputs, nnet):
+    gen = BeamSearchSequenceGenerator(*nnet, beam_size=2)
+    output = gen(*inputs, return_beam_scores=False)
+
+    assert len(output) == 1
+    (best_path,) = output
+
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+    assert best_path.shape == (26,)
+
+
+@pytest.mark.parametrize('return_xattn_scores', [False, True])
+def test_beam_decoding_beam_scores_true(inputs, nnet, return_xattn_scores):
+    gen = BeamSearchSequenceGenerator(*nnet, return_xattn_scores=return_xattn_scores, beam_size=2)
+    output = gen(*inputs, return_beam_scores=True)
+
+    assert len(output) == 4
+    beam_paths, scores, best_path, xatt_scores_list = output
+
+    assert beam_paths is not None
+    assert isinstance(beam_paths, list)
+    assert len(beam_paths) == 1
+    (beam_paths_seq0,) = beam_paths
+    assert torch.is_tensor(beam_paths_seq0)
+    assert beam_paths_seq0.shape == (2, 26)
+
+    assert scores is not None
+    assert isinstance(scores, list)
+    assert len(scores) == 1
+    (scores_seq0,) = scores
+    assert torch.is_tensor(scores_seq0)
+    assert scores_seq0.shape == (2,)
+
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+    assert best_path.shape == (1, 26)
+
+    if return_xattn_scores:
+        assert xatt_scores_list is not None
+        assert isinstance(xatt_scores_list, list)
+        assert torch.is_tensor(xatt_scores_list[0])
+        assert xatt_scores_list[0].shape == (1, 1, 25, 5)
+    else:
+        assert xatt_scores_list is None
+
+
+def test_beam_decoding_beam_scores_true_with_fusion_models(inputs, nnet):
+    """Test decoding with dummy unigram LM and boosting tree"""
+    # load dummy ngpu-lm
+    lm = NGramGPULanguageModel.dummy_unigram_lm(vocab_size=8)
+
+    # load dummy boosting tree
+    boosting_tree = GPUBoostingTreeModel.dummy_boosting_tree(vocab_size=8)
+
+    fusion_models = [lm, boosting_tree]
+    fusion_models_alpha = [0.2, 0.2]
+
+    gen = BeamSearchSequenceGeneratorWithFusionModels(
+        *nnet,
+        return_xattn_scores=True,
+        fusion_models=fusion_models,
+        fusion_models_alpha=fusion_models_alpha,
+        beam_size=2,
+    )
+    output = gen(*inputs, return_beam_scores=True)
+
+    assert len(output) == 4
+    beam_paths, scores, best_path, xatt_scores_list = output
+
+    assert beam_paths is not None
+    assert isinstance(beam_paths, list)
+    assert len(beam_paths) == 1
+    (beam_paths_seq0,) = beam_paths
+    assert torch.is_tensor(beam_paths_seq0)
+    assert beam_paths_seq0.shape == (2, 26)
+
+    assert scores is not None
+    assert isinstance(scores, list)
+    assert len(scores) == 1
+    (scores_seq0,) = scores
+    assert torch.is_tensor(scores_seq0)
+    assert scores_seq0.shape == (2,)
+
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+    assert best_path.shape == (1, 26)
+
+    assert xatt_scores_list is not None
+    assert isinstance(xatt_scores_list, list)
+    assert torch.is_tensor(xatt_scores_list[0])
+    assert xatt_scores_list[0].shape == (1, 1, 25, 5)
+
+
+@pytest.fixture()
+def prompted_inputs():
+    B, T, C = 1, 5, 2
+    return (
+        torch.tensor([[1, 0, 2, 3, 4]], dtype=torch.long),  # prompt
+        torch.ones(B, T, C, dtype=torch.float),  # encoder_hidden_states
+        torch.ones(B, T, dtype=torch.float),  # encoder_input_mask
+    )
+
+
+def test_transformer_aed_beam_infer_strips_prompt(prompted_inputs, decoder_nm, nnet, tokenizer):
+    decoder_input_ids, encoder_hidden_states, encoder_input_mask = prompted_inputs
+    *_, classifier = nnet
+
+    # Run the actual top-level module used by MultiTask AED model for decoding.
+    # This module is expected to trim the prompt from the beginning, and eos and pad from the end.
+    gen = TransformerAEDBeamInfer(decoder_nm, classifier, tokenizer)
+    ans = gen(
+        encoder_hidden_states=encoder_hidden_states,
+        encoder_input_mask=encoder_input_mask,
+        decoder_input_ids=decoder_input_ids,
+    )
+    best_path = ans[0][0].y_sequence
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+
+    # Now run the underlying beam search generator that doesn't trim anything.
+    *_, (untrimmed,), _ = gen.beam_search(*prompted_inputs, return_beam_scores=True)
+    assert untrimmed is not None
+    assert torch.is_tensor(untrimmed)
+
+    # Check that the expected trimming has indeed been done.
+    torch.testing.assert_close(
+        untrimmed[decoder_input_ids.shape[1] :], best_path
+    )  # stripped the prompt from the beggining
+
+
+def test_transformer_aed_greedy_infer_strips_prompt(prompted_inputs, decoder_nm, nnet, tokenizer):
+    decoder_input_ids, encoder_hidden_states, encoder_input_mask = prompted_inputs
+    decoder_input_ids = torch.tensor([[1, 0, 2, 3, 4]], dtype=torch.long)  # prompt
+    *_, classifier = nnet
+
+    # Run the actual top-level module used by MultiTask AED model for decoding.
+    # This module is expected to trim the prompt from the beginning, and eos and pad from the end.
+    gen = TransformerAEDGreedyInfer(decoder_nm, classifier, tokenizer)
+    ans = gen(
+        encoder_hidden_states=encoder_hidden_states,
+        encoder_input_mask=encoder_input_mask,
+        decoder_input_ids=decoder_input_ids,
+    )
+    best_path = ans[0][0].y_sequence
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+
+    # Now run the underlying beam search generator that doesn't trim anything.
+    (untrimmed,), _, _, _ = gen.greedy_search(*prompted_inputs)
+    assert untrimmed is not None
+    assert torch.is_tensor(untrimmed)
+
+    # Check that the expected trimming has indeed been done.
+    torch.testing.assert_close(
+        untrimmed[decoder_input_ids.shape[1] :], best_path
+    )  # stripped the prompt from the beggining
+
+
+def test_transformer_aed_beam_infer_trims_xatt_scores(prompted_inputs, decoder_nm, nnet, tokenizer):
+    decoder_input_ids, encoder_hidden_states, encoder_input_mask = prompted_inputs
+    *_, classifier = nnet
+
+    # Run the actual top-level module used by MultiTask AED model for decoding.
+    # This module is expected to trim eos and pads in xatt from the end.
+    gen = TransformerAEDBeamInfer(decoder_nm, classifier, tokenizer, return_xattn_scores=True)
+    ans = gen(
+        encoder_hidden_states=encoder_hidden_states,
+        encoder_input_mask=encoder_input_mask,
+        decoder_input_ids=decoder_input_ids,
+    )
+    hyp = ans[0][0]
+
+    assert hyp.xatt_scores is not None
+    seq_len = hyp.y_sequence.shape[0]
+    decoder_input_ids_len = decoder_input_ids.shape[1]
+    total_expected_len = seq_len + decoder_input_ids_len - 1
+
+    # Check that the expected trimming has indeed been done.
+    for layer_scores in hyp.xatt_scores:
+        assert layer_scores.shape[1] == total_expected_len
+
+
+def test_transformer_aed_greedy_infer_trims_xatt_scores(prompted_inputs, decoder_nm, nnet, tokenizer):
+    decoder_input_ids, encoder_hidden_states, encoder_input_mask = prompted_inputs
+    *_, classifier = nnet
+
+    # Run the actual top-level module used by MultiTask AED model for decoding.
+    # This module is expected to trim eos and pads in xatt from the end.
+    gen = TransformerAEDGreedyInfer(decoder_nm, classifier, tokenizer, return_xattn_scores=True)
+    ans = gen(
+        encoder_hidden_states=encoder_hidden_states,
+        encoder_input_mask=encoder_input_mask,
+        decoder_input_ids=decoder_input_ids,
+    )
+    hyp = ans[0][0]
+
+    assert hyp.xatt_scores is not None
+    seq_len = hyp.y_sequence.shape[0]
+    decoder_input_ids_len = decoder_input_ids.shape[1]
+    total_expected_len = seq_len + decoder_input_ids_len - 1
+
+    # Check that the expected trimming has indeed been done.
+    for layer_scores in hyp.xatt_scores:
+        assert layer_scores.shape[1] == total_expected_len

tests/collections/asr/decoding/test_multi_task_streaming_decoding.py

@@ -0,0 +1,192 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+from kaldialign import edit_distance
+from omegaconf import OmegaConf
+from tqdm.auto import tqdm
+
+from nemo.collections.asr.models.aed_multitask_models import lens_to_mask
+from nemo.collections.asr.parts.submodules.aed_decoding import (
+    GreedyBatchedStreamingAEDComputer,
+    return_decoder_input_ids,
+)
+from nemo.collections.asr.parts.submodules.multitask_decoding import (
+    AEDStreamingDecodingConfig,
+    MultiTaskDecodingConfig,
+)
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+from nemo.collections.asr.parts.utils.streaming_utils import ContextSize
+from tests.collections.asr.decoding.utils import load_audio, make_preprocessor_deterministic
+
+DEVICES = [torch.device("cpu")]
+if torch.cuda.is_available():
+    DEVICES.append(torch.device("cuda:0"))
+
+if torch.mps.is_available():
+    DEVICES.append(torch.device("mps"))
+
+
+def get_batch_encoder_outputs_from_records(records, model, device):
+    """Helper function to get encoder outputs for a batch of manifest records"""
+    local_batch_size = len(records)
+    filenames = [record["audio_filepath"] for record in records]
+    audio_filepaths = filenames[:local_batch_size]
+
+    with torch.no_grad():
+        make_preprocessor_deterministic(model)
+        model.eval()
+
+        all_inputs, all_lengths = [], []
+        for audio_file in tqdm(audio_filepaths, desc="Loading audio files"):
+            audio_tensor, _ = load_audio(audio_file)
+            all_inputs.append(audio_tensor)
+            all_lengths.append(torch.tensor(audio_tensor.shape[0], dtype=torch.int64))
+
+        input_batch = torch.nn.utils.rnn.pad_sequence(all_inputs, batch_first=True).to(
+            device=device, dtype=torch.float32
+        )
+        length_batch = torch.tensor(all_lengths, dtype=torch.int64).to(device)
+        # get encoder output using full audio signal
+        _, encoded_length, encoded_output, _ = model(input_signal=input_batch, input_signal_length=length_batch)
+
+    return encoded_output, encoded_length
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize(
+    "device,use_cuda_graph_decoder",
+    [(device, False) for device in DEVICES] + [(device, True) for device in DEVICES if device.type == "cuda"],
+)
+@pytest.mark.parametrize("decoding_policy", ["waitk", "alignatt"])
+@pytest.mark.parametrize("chunk_size", [3, 4])
+@pytest.mark.parametrize("batch_size", [4])
+def test_multi_task_streaming_decoding(
+    tmp_path_factory,
+    an4_val_manifest_corrected,
+    canary_180m_flash,
+    device: torch.device,
+    use_cuda_graph_decoder: bool,
+    decoding_policy: str,
+    chunk_size: int,
+    batch_size: int,
+):
+    """Test streaming decoding with multi-task model for different decoding policies"""
+    model = canary_180m_flash
+    model.eval()
+    model.to(device=device)
+
+    # setup streaming decoding config
+    streaming_decoding_cfg = AEDStreamingDecodingConfig()
+    streaming_decoding_cfg.streaming_policy = decoding_policy
+    streaming_decoding_cfg.chunk_secs = 1
+    streaming_decoding_cfg.right_context_secs = 0.0
+    streaming_decoding_cfg.batch_size = batch_size
+    streaming_decoding_cfg.prompt = OmegaConf.create({'pnc': 'no', 'task': 'asr'})
+
+    context_encoder_frames = ContextSize(
+        left=100,
+        chunk=chunk_size,
+        right=0.0,
+    )
+
+    # setup decoding strategy
+    if hasattr(model, 'change_decoding_strategy'):
+        multitask_decoding = MultiTaskDecodingConfig()
+        multitask_decoding.strategy = "greedy"
+        model.change_decoding_strategy(multitask_decoding)
+
+    manifest = read_manifest(an4_val_manifest_corrected)
+
+    all_hyps = []
+    tokens_frame_alignment = []
+    predicted_token_ids = []
+    decoding_computer = GreedyBatchedStreamingAEDComputer(
+        model,
+        frame_chunk_size=chunk_size,
+        decoding_cfg=streaming_decoding_cfg,
+    )
+
+    with torch.no_grad(), torch.inference_mode():
+        for i in range(0, len(manifest), batch_size):
+            encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+                manifest[i : i + batch_size], model=model, device=device
+            )
+            local_batch_size = encoder_output_len.shape[0]
+            decoder_input_ids = return_decoder_input_ids(streaming_decoding_cfg, model)
+            model_state = GreedyBatchedStreamingAEDComputer.initialize_aed_model_state(
+                asr_model=model,
+                decoder_input_ids=decoder_input_ids,
+                batch_size=local_batch_size,
+                context_encoder_frames=context_encoder_frames,
+                chunk_secs=streaming_decoding_cfg.chunk_secs,
+                right_context_secs=streaming_decoding_cfg.right_context_secs,
+            )
+
+            # decode encoder output by chunks, passing state between decoder invocations
+            for t in range(0, encoder_output.shape[1], chunk_size):
+                current_len = torch.full_like(encoder_output_len, fill_value=t + chunk_size)
+                current_len = torch.minimum(current_len, encoder_output_len)
+                model_state.is_last_chunk_batch = current_len >= encoder_output_len
+                encoder_input_mask = lens_to_mask(current_len, encoder_output[:, : t + chunk_size].shape[1]).to(
+                    encoder_output.dtype
+                )
+
+                model_state = decoding_computer(
+                    encoder_output=encoder_output[:, : t + chunk_size],
+                    encoder_output_len=current_len,
+                    encoder_input_mask=encoder_input_mask,
+                    prev_batched_state=model_state,
+                )
+            # get final results for each sample in the batch
+            for j in range(local_batch_size):
+                transcription_idx = model_state.pred_tokens_ids[
+                    j, model_state.decoder_input_ids.size(-1) : model_state.current_context_lengths[j]
+                ]
+                transcription = model.tokenizer.ids_to_text(transcription_idx.tolist()).strip()
+                all_hyps.append(transcription)
+                tokens_frame_alignment.append(model_state.tokens_frame_alignment[j])
+                predicted_token_ids.append(
+                    model_state.pred_tokens_ids[
+                        j, model_state.decoder_input_ids.size(-1) : model_state.current_context_lengths[j]
+                    ]
+                )
+
+    # compare decoding results with reference transcripts
+    ref_transcripts = [item['text'] for item in manifest]
+    assert (
+        edit_distance(ref_transcripts, all_hyps)['total'] <= len(ref_transcripts) * 0.1
+    )  # Expected WER is less than 10%
+
+    # compute latency
+    audio_encoder_fs = 80  # in ms
+    laal_list = None
+    if decoding_policy == "waitk":
+        laal_list = decoding_computer.compute_waitk_lagging(
+            manifest, predicted_token_ids, context_encoder_frames, audio_encoder_fs, BOW_PREFIX="\u2581"
+        )
+    elif decoding_policy == "alignatt":
+        laal_list = decoding_computer.compute_alignatt_lagging(
+            manifest,
+            predicted_token_ids,
+            tokens_frame_alignment,
+            context_encoder_frames,
+            audio_encoder_fs,
+            BOW_PREFIX="\u2581",
+        )
+    else:
+        raise ValueError(f"Decoding policy {decoding_policy} is not supported")
+    laal = sum(laal_list) / len(laal_list)
+    assert 300 <= laal <= 900  # Expected LAAL is between 300ms and 900ms depending on the decoding policy

tests/collections/asr/decoding/test_rnnt_alignments.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from pathlib import Path
+from typing import Union
+
+import pytest
+import torch.cuda
+from examples.asr.transcribe_speech import TranscriptionConfig
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.models import EncDecRNNTBPEModel
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.collections.asr.parts.utils.transcribe_utils import prepare_audio_data
+
+DEVICES = []
+
+if torch.cuda.is_available():
+    DEVICES.append('cuda')
+
+if torch.mps.is_available():
+    DEVICES.append('mps')
+
+
+@pytest.fixture(scope="module")
+def stt_en_conformer_transducer_small_model():
+    model = EncDecRNNTBPEModel.from_pretrained(model_name="stt_en_conformer_transducer_small", map_location="cpu")
+    return model
+
+
+def get_rnnt_alignments(
+    strategy: str,
+    manifest_path: Union[Path, str],
+    model: EncDecRNNTBPEModel,
+    loop_labels: bool = True,
+    use_cuda_graph_decoder=False,
+    device="cuda",
+) -> list[Hypothesis]:
+    cfg = OmegaConf.structured(TranscriptionConfig())
+    cfg.rnnt_decoding.confidence_cfg.preserve_frame_confidence = True
+    cfg.rnnt_decoding.confidence_cfg.exclude_blank = False
+    cfg.rnnt_decoding.preserve_alignments = True
+    cfg.rnnt_decoding.strategy = strategy
+    if cfg.rnnt_decoding.strategy == "greedy_batch":
+        cfg.rnnt_decoding.greedy.loop_labels = loop_labels
+        cfg.rnnt_decoding.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+    cfg.dataset_manifest = str(manifest_path)
+    filepaths = prepare_audio_data(cfg)[0][:8]  # selecting 8 files only
+    # NB: 9th file has the same transcription but a bit different alignment for batched/non-batched decoding
+
+    model = model.to(device)
+    model.change_decoding_strategy(cfg.rnnt_decoding)
+
+    transcriptions: list[Hypothesis] = model.transcribe(
+        audio=filepaths,
+        batch_size=cfg.batch_size,
+        num_workers=cfg.num_workers,
+        return_hypotheses=True,
+        channel_selector=cfg.channel_selector,
+    )
+
+    for transcription in transcriptions:
+        for align_elem, frame_confidence in zip(transcription.alignments, transcription.frame_confidence):
+            assert len(align_elem) == len(frame_confidence)  # frame confidences have to match alignments
+            assert len(align_elem) > 0  # no empty alignments
+            for idx, pred in enumerate(align_elem):
+                if idx < len(align_elem) - 1:
+                    assert pred[1].item() != model.decoder.blank_idx  # all except last have to be non-blank
+                else:
+                    assert pred[1].item() == model.decoder.blank_idx  # last one has to be blank
+    return transcriptions
+
+
+@pytest.fixture(autouse=True)
+def cleanup_local_folder():
+    """Overriding global fixture to make sure it's not applied for this test.
+
+    Otherwise, there will be errors in the CI in github.
+    """
+    return
+
+
+# TODO: add the same tests for multi-blank RNNT decoding
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("loop_labels", [True, False])
+@pytest.mark.parametrize("use_cuda_graph_decoder", [True, False])
+@pytest.mark.with_downloads
+def test_rnnt_alignments(
+    loop_labels: bool,
+    use_cuda_graph_decoder: bool,
+    device: str,
+    an4_val_manifest_corrected,
+    stt_en_conformer_transducer_small_model,
+):
+    if use_cuda_graph_decoder and device != "cuda":
+        pytest.skip("CUDA decoder works only with CUDA")
+    if not loop_labels and use_cuda_graph_decoder:
+        pytest.skip("Frame-Looping algorithm with CUDA graphs does not yet support alignments")
+    # using greedy as baseline and comparing all other configurations to it
+    ref_transcriptions = get_rnnt_alignments(
+        "greedy",
+        manifest_path=an4_val_manifest_corrected,
+        model=stt_en_conformer_transducer_small_model,
+        device=device,
+    )
+    transcriptions = get_rnnt_alignments(
+        "greedy_batch",
+        loop_labels=loop_labels,
+        use_cuda_graph_decoder=use_cuda_graph_decoder,
+        manifest_path=an4_val_manifest_corrected,
+        model=stt_en_conformer_transducer_small_model,
+        device=device,
+    )
+    # comparing that label sequence in alignments is exactly the same
+    # we can't compare logits as well, because they are expected to be
+    # slightly different in batched and single-sample mode
+    assert len(ref_transcriptions) == len(transcriptions)
+    for ref_transcription, transcription in zip(ref_transcriptions, transcriptions):
+        for ref_align_elem, align_elem in zip(ref_transcription.alignments, transcription.alignments):
+            assert len(ref_align_elem) == len(align_elem)
+            for ref_pred, pred in zip(ref_align_elem, align_elem):
+                assert ref_pred[1].item() == pred[1].item()

tests/collections/asr/decoding/test_rnnt_decoding.py

@@ -0,0 +1,661 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import os
+from functools import cached_property, lru_cache
+from pathlib import Path
+from typing import Optional
+
+import pytest
+import torch
+from omegaconf import DictConfig
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.modules import RNNTDecoder, RNNTJoint
+from nemo.collections.asr.parts.context_biasing import BoostingTreeModelConfig, GPUBoostingTreeModel
+from nemo.collections.asr.parts.mixins import mixins
+from nemo.collections.asr.parts.submodules import rnnt_beam_decoding
+from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode
+from nemo.collections.asr.parts.submodules import tdt_beam_decoding
+from nemo.collections.asr.parts.submodules.ngram_lm import NGramGPULanguageModel
+from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTBPEDecoding, RNNTDecoding, RNNTDecodingConfig
+from nemo.collections.asr.parts.utils import rnnt_utils
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+from tests.collections.asr.decoding.test_timestamps import BaseTimestampsTest
+
+NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cpu_is_supported(
+    __NUMBA_MINIMUM_VERSION__
+) or numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+
+
+def char_vocabulary():
+    return [' ', 'a', 'b', 'c', 'd', 'e', 'f', '.']
+
+
+@pytest.fixture()
+@lru_cache(maxsize=8)
+def tmp_tokenizer(test_data_dir):
+    cfg = DictConfig({'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'})
+
+    class _TmpASRBPE(mixins.ASRBPEMixin):
+        def register_artifact(self, _, vocab_path):
+            return vocab_path
+
+    asrbpe = _TmpASRBPE()
+    asrbpe._setup_tokenizer(cfg)
+    return asrbpe.tokenizer
+
+
+@lru_cache(maxsize=2)
+def get_rnnt_decoder(vocab_size, decoder_output_size=4):
+    prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1}
+    torch.manual_seed(0)
+    decoder = RNNTDecoder(prednet=prednet_cfg, vocab_size=vocab_size)
+    decoder.freeze()
+    return decoder
+
+
+@lru_cache(maxsize=2)
+def get_rnnt_joint(vocab_size, vocabulary=None, encoder_output_size=4, decoder_output_size=4, joint_output_shape=4):
+    jointnet_cfg = {
+        'encoder_hidden': encoder_output_size,
+        'pred_hidden': decoder_output_size,
+        'joint_hidden': joint_output_shape,
+        'activation': 'relu',
+    }
+    torch.manual_seed(0)
+    joint = RNNTJoint(jointnet_cfg, vocab_size, vocabulary=vocabulary)
+    joint.freeze()
+    return joint
+
+
+@lru_cache(maxsize=1)
+def get_model_encoder_output(data_dir, model_name):
+    # Import inside function to avoid issues with dependencies
+    import librosa
+
+    audio_filepath = os.path.join(data_dir, 'asr', 'test', 'an4', 'wav', 'cen3-fjlp-b.wav')
+
+    with torch.no_grad():
+        model = ASRModel.from_pretrained(model_name, map_location='cpu')  # type: ASRModel
+        model.preprocessor.featurizer.dither = 0.0
+        model.preprocessor.featurizer.pad_to = 0
+        model.eval()
+
+        audio, sr = librosa.load(path=audio_filepath, sr=16000, mono=True)
+
+        input_signal = torch.tensor(audio, dtype=torch.float32).unsqueeze(0)
+        input_signal_length = torch.tensor([len(audio)], dtype=torch.int32)
+
+        encoded, encoded_len = model(input_signal=input_signal, input_signal_length=input_signal_length)
+
+    return model, encoded, encoded_len
+
+
+def decode_text_from_greedy_hypotheses(hyps, decoding):
+    decoded_hyps = decoding.decode_hypothesis(hyps)  # type: List[str]
+
+    return decoded_hyps
+
+
+def decode_text_from_nbest_hypotheses(hyps, decoding):
+    hypotheses = []
+    all_hypotheses = []
+
+    for nbest_hyp in hyps:  # type: rnnt_utils.NBestHypotheses
+        n_hyps = nbest_hyp.n_best_hypotheses  # Extract all hypotheses for this sample
+        decoded_hyps = decoding.decode_hypothesis(n_hyps)  # type: List[str]
+
+        hypotheses.append(decoded_hyps[0])  # best hypothesis
+        all_hypotheses.append(decoded_hyps)
+
+    return hypotheses, all_hypotheses
+
+
+def check_beam_decoding(test_data_dir, beam_config):
+    beam_size = beam_config.pop("beam_size", 1)
+    model, encoded, encoded_len = get_model_encoder_output(test_data_dir, 'nvidia/parakeet-tdt_ctc-110m')
+
+    model_config = model.to_config_dict()
+    durations = list(model_config["model_defaults"]["tdt_durations"])
+
+    beam = tdt_beam_decoding.BeamTDTInfer(
+        model.decoder,
+        model.joint,
+        beam_size=beam_size,
+        return_best_hypothesis=False,
+        durations=durations,
+        **beam_config,
+    )
+
+    enc_out = encoded
+    enc_len = encoded_len
+
+    with torch.no_grad():
+        hyps: rnnt_utils.Hypothesis = beam(encoder_output=enc_out, encoded_lengths=enc_len)[0]
+        _, all_hyps = decode_text_from_nbest_hypotheses(hyps, model.decoding)
+        all_hyps = all_hyps[0]
+
+        print("Beam search algorithm :", beam_config['search_type'])
+        for idx, hyp_ in enumerate(all_hyps):
+            print("Hyp index", idx + 1, "text :", hyp_.text)
+
+            assert len(hyp_.timestamp) > 0
+            print("Timesteps", hyp_.timestamp)
+            print()
+
+
+def check_tdt_greedy_decoding(
+    test_data_dir,
+    use_cuda_graph_decoder: bool,
+    lm_path: Optional[str | Path] = None,
+    boosting_tree: Optional[BoostingTreeModelConfig] = None,
+    enable_per_stream_biasing: bool = False,
+):
+    model, encoded, encoded_len = get_model_encoder_output(test_data_dir, 'nvidia/parakeet-tdt_ctc-110m')
+
+    model_config = model.to_config_dict()
+
+    fusion_models, fusion_models_alpha = None, None
+    if lm_path or boosting_tree:
+        fusion_models = []
+        fusion_models_alpha = []
+
+    if lm_path:
+        fusion_models.append(NGramGPULanguageModel.from_file(lm_path=lm_path, vocab_size=model.decoder.blank_idx))
+        fusion_models_alpha.append(0.5)
+    if boosting_tree:
+        fusion_models.append(GPUBoostingTreeModel.from_config(boosting_tree, tokenizer=model.tokenizer))
+        fusion_models_alpha.append(0.5)
+
+    decoding_algo = greedy_decode.GreedyBatchedTDTInfer(
+        model.decoder,
+        model.joint,
+        blank_index=model.decoder.blank_idx,
+        durations=list(model_config["model_defaults"]["tdt_durations"]),
+        max_symbols_per_step=10,
+        preserve_alignments=False,
+        preserve_frame_confidence=False,
+        use_cuda_graph_decoder=use_cuda_graph_decoder,
+        fusion_models=fusion_models,
+        fusion_models_alpha=fusion_models_alpha,
+        enable_per_stream_biasing=enable_per_stream_biasing,
+    )
+
+    enc_out = encoded
+    enc_len = encoded_len
+
+    with torch.no_grad():
+        hyps: rnnt_utils.Hypothesis = decoding_algo(encoder_output=enc_out, encoded_lengths=enc_len)[0]
+        all_hyps = decode_text_from_greedy_hypotheses(hyps, model.decoding)
+
+        print("Decoding result")
+        for idx, hyp_ in enumerate(all_hyps):
+            print(f"Hyp index {idx + 1} | text : {hyp_.text}")
+            assert len(hyp_.timestamp) > 0
+            print("Timesteps", hyp_.timestamp)
+            print()
+
+
+class TestRNNTDecoding:
+    @pytest.mark.unit
+    def test_constructor(self):
+        cfg = RNNTDecodingConfig()
+        vocab = char_vocabulary()
+        decoder = get_rnnt_decoder(vocab_size=len(vocab))
+        joint = get_rnnt_joint(vocab_size=len(vocab))
+        decoding = RNNTDecoding(decoding_cfg=cfg, decoder=decoder, joint=joint, vocabulary=vocab)
+        assert decoding is not None
+
+    @pytest.mark.unit
+    def test_constructor_subword(self, tmp_tokenizer):
+        cfg = RNNTDecodingConfig()
+        vocab = tmp_tokenizer.vocab
+        decoder = get_rnnt_decoder(vocab_size=len(vocab))
+        joint = get_rnnt_joint(vocab_size=len(vocab))
+        decoding = RNNTBPEDecoding(decoding_cfg=cfg, decoder=decoder, joint=joint, tokenizer=tmp_tokenizer)
+        assert decoding is not None
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    def test_greedy_decoding_preserve_alignments(self, test_data_dir):
+        model, encoded, encoded_len = get_model_encoder_output(test_data_dir, 'stt_en_conformer_transducer_small')
+
+        beam = greedy_decode.GreedyRNNTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=model.joint.num_classes_with_blank - 1,
+            max_symbols_per_step=5,
+            preserve_alignments=True,
+        )
+
+        enc_out = encoded
+        enc_len = encoded_len
+
+        with torch.no_grad():
+            hyps = beam(encoder_output=enc_out, encoded_lengths=enc_len)[0]  # type: rnnt_utils.Hypothesis
+            hyp = decode_text_from_greedy_hypotheses(hyps, model.decoding)
+            hyp = hyp[0]
+
+            assert hyp.alignments is not None
+
+            # Use the following commented print statements to check
+            # the alignment of other algorithms compared to the default
+            print("Text", hyp.text)
+            for t in range(len(hyp.alignments)):
+                t_u = []
+                for u in range(len(hyp.alignments[t])):
+                    logp, label = hyp.alignments[t][u]
+                    assert torch.is_tensor(logp)
+                    assert torch.is_tensor(label)
+
+                    t_u.append(int(label))
+
+                print(f"Tokens at timestamp {t} = {t_u}")
+            print()
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize("loop_labels", [True, False])
+    def test_batched_greedy_decoding_preserve_alignments(self, test_data_dir, loop_labels: bool):
+        """Test batched greedy decoding using non-batched decoding as a reference"""
+        model, encoded, encoded_len = get_model_encoder_output(test_data_dir, 'stt_en_conformer_transducer_small')
+
+        search_algo = greedy_decode.GreedyBatchedRNNTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=model.joint.num_classes_with_blank - 1,
+            max_symbols_per_step=5,
+            preserve_alignments=True,
+            loop_labels=loop_labels,
+        )
+
+        etalon_search_algo = greedy_decode.GreedyRNNTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=model.joint.num_classes_with_blank - 1,
+            max_symbols_per_step=5,
+            preserve_alignments=True,
+        )
+
+        enc_out = encoded
+        enc_len = encoded_len
+
+        with torch.no_grad():
+            hyps: list[rnnt_utils.Hypothesis] = search_algo(encoder_output=enc_out, encoded_lengths=enc_len)[0]
+            hyp = decode_text_from_greedy_hypotheses(hyps, model.decoding)[0]
+            etalon_hyps: list[rnnt_utils.Hypothesis] = etalon_search_algo(
+                encoder_output=enc_out, encoded_lengths=enc_len
+            )[0]
+            etalon_hyp = decode_text_from_greedy_hypotheses(etalon_hyps, model.decoding)[0]
+
+            assert hyp.alignments is not None
+            assert etalon_hyp.alignments is not None
+
+            assert hyp.text == etalon_hyp.text
+            assert len(hyp.alignments) == len(etalon_hyp.alignments)
+
+            for t in range(len(hyp.alignments)):
+                t_u = []
+                for u in range(len(hyp.alignments[t])):
+                    logp, label = hyp.alignments[t][u]
+                    assert torch.is_tensor(logp)
+                    assert torch.is_tensor(label)
+                    etalon_logp, etalon_label = etalon_hyp.alignments[t][u]
+                    assert label == etalon_label
+                    assert torch.allclose(logp, etalon_logp, atol=1e-4, rtol=1e-4)
+
+                    t_u.append(int(label))
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "greedy"},
+            {
+                "search_type": "default",
+                "beam_size": 2,
+            },
+            {
+                "search_type": "alsd",
+                "alsd_max_target_len": 0.5,
+                "beam_size": 2,
+            },
+            {
+                "search_type": "tsd",
+                "tsd_max_sym_exp_per_step": 3,
+                "beam_size": 2,
+            },
+            {"search_type": "maes", "maes_num_steps": 2, "maes_expansion_beta": 2, "beam_size": 2},
+            {"search_type": "maes", "maes_num_steps": 3, "maes_expansion_beta": 1, "beam_size": 2},
+        ],
+    )
+    def test_rnnt_beam_decoding_preserve_alignments(self, test_data_dir, beam_config):
+        beam_size = beam_config.pop("beam_size", 1)
+        model, encoded, encoded_len = get_model_encoder_output(test_data_dir, 'stt_en_conformer_transducer_small')
+        beam = rnnt_beam_decoding.BeamRNNTInfer(
+            model.decoder,
+            model.joint,
+            beam_size=beam_size,
+            return_best_hypothesis=False,
+            preserve_alignments=True,
+            **beam_config,
+        )
+
+        enc_out = encoded
+        enc_len = encoded_len
+        blank_id = torch.tensor(model.joint.num_classes_with_blank - 1, dtype=torch.int32)
+
+        with torch.no_grad():
+            hyps = beam(encoder_output=enc_out, encoded_lengths=enc_len)[0]  # type: rnnt_utils.Hypothesis
+            hyp, all_hyps = decode_text_from_nbest_hypotheses(hyps, model.decoding)
+            hyp = hyp[0]  # best hypothesis
+            all_hyps = all_hyps[0]
+
+            assert hyp.alignments is not None
+
+            if beam_config['search_type'] == 'alsd':
+                assert len(all_hyps) <= int(beam_config['alsd_max_target_len'] * float(enc_len[0]))
+
+            print("Beam search algorithm :", beam_config['search_type'])
+            # Use the following commented print statements to check
+            # the alignment of other algorithms compared to the default
+            for idx, hyp_ in enumerate(all_hyps):  # type: (int, rnnt_utils.Hypothesis)
+                print("Hyp index", idx + 1, "text :", hyp_.text)
+
+                # Alignment length (T) must match audio length (T)
+                # NOTE: increase length threshold to two to prevent intermittent failures when a word is split into subwords
+                assert abs(len(hyp_.alignments) - enc_len[0]) <= 2  # 1
+
+                for t in range(len(hyp_.alignments)):
+                    t_u = []
+                    for u in range(len(hyp_.alignments[t])):
+                        logp, label = hyp_.alignments[t][u]
+                        assert torch.is_tensor(logp)
+                        assert torch.is_tensor(label)
+
+                        t_u.append(int(label))
+
+                    # Blank token must be the last token in the current
+                    if len(t_u) > 1:
+                        assert t_u[-1] == blank_id
+
+                        # No blank token should be present in the current timestamp other than at the end
+                        for token in t_u[:-1]:
+                            assert token != blank_id
+
+                    print(f"Tokens at timestamp {t} = {t_u}")
+                print()
+
+                assert len(hyp_.timestamp) > 0
+                print("Timesteps", hyp_.timestamp)
+                print()
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "model_name, decoding_strategy",
+        [
+            ("stt_en_conformer_transducer_small", "greedy"),
+            ("stt_en_conformer_transducer_small", "greedy_batch"),
+            ("stt_en_conformer_transducer_small", "beam"),
+            # ("stt_en_conformer_transducer_small", "tsd"),
+            ("stt_en_conformer_transducer_small", "alsd"),
+            ("nvidia/parakeet-tdt_ctc-110m", "greedy"),
+            ("nvidia/parakeet-tdt_ctc-110m", "greedy_batch"),
+        ],
+    )
+    def test_subword_decoding_compute_timestamps(self, test_data_dir, decoding_strategy, model_name):
+
+        model, encoded, encoded_len = get_model_encoder_output(test_data_dir, model_name)
+
+        cfg = DictConfig(model.cfg.decoding)
+        cfg['strategy'] = decoding_strategy
+        cfg['preserve_alignments'] = True
+        cfg['compute_timestamps'] = True
+
+        decoding = RNNTBPEDecoding(
+            decoding_cfg=cfg, decoder=model.decoder, joint=model.joint, tokenizer=model.tokenizer
+        )
+
+        hyps = decoding.rnnt_decoder_predictions_tensor(encoded, encoded_len, return_hypotheses=True)
+        if isinstance(hyps[0], list):
+            BaseTimestampsTest.check_subword_timestamps(hyps[0][0], decoding)
+        else:
+            BaseTimestampsTest.check_subword_timestamps(hyps[0], decoding)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "model_name, decoding_strategy",
+        [
+            ("stt_en_conformer_transducer_small", "greedy"),
+            ("stt_en_conformer_transducer_small", "greedy_batch"),
+            ("stt_en_conformer_transducer_small", "beam"),
+            # ("stt_en_conformer_transducer_small", "tsd"),
+            ("stt_en_conformer_transducer_small", "alsd"),
+            ("nvidia/parakeet-tdt_ctc-110m", "greedy"),
+            ("nvidia/parakeet-tdt_ctc-110m", "greedy_batch"),
+        ],
+    )
+    def test_char_decoding_compute_timestamps(self, test_data_dir, decoding_strategy, model_name):
+
+        model, encoded, encoded_len = get_model_encoder_output(test_data_dir, model_name)
+
+        cfg = DictConfig(model.cfg.decoding)
+        cfg['strategy'] = decoding_strategy
+        cfg['preserve_alignments'] = True
+        cfg['compute_timestamps'] = True
+
+        vocab = [t[0] for t in model.tokenizer.vocab]
+
+        decoding = RNNTDecoding(decoding_cfg=cfg, decoder=model.decoder, joint=model.joint, vocabulary=vocab)
+
+        hyps = decoding.rnnt_decoder_predictions_tensor(encoded, encoded_len, return_hypotheses=True)
+
+        if isinstance(hyps[0], list):
+            BaseTimestampsTest.check_char_timestamps(hyps[0][0], decoding)
+        else:
+            BaseTimestampsTest.check_char_timestamps(hyps[0], decoding)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize("use_cuda_graph_decoder", [True, False])
+    @pytest.mark.parametrize("use_lm", [True, False])
+    @pytest.mark.parametrize("use_boosting_tree", [True, False])
+    @pytest.mark.parametrize("enable_per_stream_biasing", [True, False])
+    def test_tdt_greedy_decoding(
+        self,
+        test_data_dir,
+        use_cuda_graph_decoder: bool,
+        use_lm: bool,
+        use_boosting_tree: bool,
+        enable_per_stream_biasing: bool,
+    ):
+        kenlm_model_path = Path(test_data_dir) / "asr/kenlm_ngram_lm/parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+        boosting_tree = BoostingTreeModelConfig(key_phrases_list=["hello", "nvidia"]) if use_boosting_tree else None
+        check_tdt_greedy_decoding(
+            test_data_dir,
+            use_cuda_graph_decoder=use_cuda_graph_decoder,
+            lm_path=kenlm_model_path if use_lm else None,
+            boosting_tree=boosting_tree,
+            enable_per_stream_biasing=enable_per_stream_biasing,
+        )
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {
+                "search_type": "default",
+                "beam_size": 2,
+            },
+            {"search_type": "maes", "maes_num_steps": 2, "maes_expansion_beta": 2, "beam_size": 2},
+            {"search_type": "maes", "maes_num_steps": 2, "maes_expansion_beta": 1, "beam_size": 4},
+        ],
+    )
+    def test_tdt_beam_decoding(self, test_data_dir, beam_config):
+        check_beam_decoding(test_data_dir, beam_config)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {
+                "search_type": "maes",
+                "maes_num_steps": 2,
+                "maes_expansion_beta": 1,
+                "beam_size": 4,
+                "ngram_lm_alpha": 0.3,
+            },
+        ],
+    )
+    def test_tdt_beam_decoding_with_kenlm(self, test_data_dir, beam_config):
+        # skipping if kenlm is not installed
+        pytest.importorskip("kenlm", reason="Skipping test because 'kenlm' is not installed.")
+
+        kenlm_model_path = os.path.join(
+            test_data_dir, "asr", "kenlm_ngram_lm", "parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+        )
+        beam_config["ngram_lm_model"] = kenlm_model_path
+        check_beam_decoding(test_data_dir, beam_config)
+
+
+class TestRNNTTimestamps(BaseTimestampsTest):
+    """RNNT-specific timestamp tests that inherit from BaseTimestampsTest"""
+
+    def _convert_offsets(self, offsets):
+        result = copy.deepcopy(offsets)
+        for offset in result:
+            offset['char'] = [offset['char']]
+        return result
+
+    @property
+    def char_offsets_chars(self):
+        return self._convert_offsets(super().char_offsets_chars)
+
+    @property
+    def char_offsets_wpe(self):
+        return self._convert_offsets(super().char_offsets_wpe)
+
+    @property
+    def char_offsets_bpe(self):
+        return self._convert_offsets(super().char_offsets_bpe)
+
+    @property
+    def encoded_char_offsets_bpe(self):
+        return self._convert_offsets(super().encoded_char_offsets_bpe)
+
+    @cached_property
+    def decoding_char(self):
+        cfg = RNNTDecodingConfig()
+        vocab = char_vocabulary()
+        decoder = get_rnnt_decoder(vocab_size=len(vocab))
+        joint = get_rnnt_joint(vocab_size=len(vocab))
+        decoding = RNNTDecoding(decoding_cfg=cfg, decoder=decoder, joint=joint, vocabulary=vocab)
+        return decoding
+
+    @cached_property
+    def decoding_subword_wpe(self):
+        cfg = RNNTDecodingConfig()
+        vocab = self.tmp_tokenizer.vocab
+        decoder = get_rnnt_decoder(vocab_size=len(vocab))
+        joint = get_rnnt_joint(vocab_size=len(vocab))
+        decoding = RNNTBPEDecoding(decoding_cfg=cfg, decoder=decoder, joint=joint, tokenizer=self.tmp_tokenizer)
+        return decoding
+
+    @cached_property
+    def decoding_subword_bpe(self):
+        vocab = self.bpe_tokenizer.vocab
+        cfg = RNNTDecodingConfig()
+        decoder = get_rnnt_decoder(vocab_size=len(vocab))
+        joint = get_rnnt_joint(vocab_size=len(vocab))
+        decoding = RNNTBPEDecoding(decoding_cfg=cfg, decoder=decoder, joint=joint, tokenizer=self.bpe_tokenizer)
+        return decoding
+
+    @pytest.mark.unit
+    def test_word_offsets_subword_wpe(self, tmp_tokenizer):
+        self.tmp_tokenizer = tmp_tokenizer
+        super().test_word_offsets_subword_wpe()
+
+    @pytest.mark.unit
+    def test_word_offsets_subword_wpe_other_delimiter(self, tmp_tokenizer):
+        self.tmp_tokenizer = tmp_tokenizer
+        super().test_word_offsets_subword_wpe_other_delimiter()
+
+
+@pytest.mark.unit
+@pytest.mark.with_downloads
+def test_transcribe_timestamps_no_decoder_reinstantiation(stt_en_fastconformer_transducer_large, test_data_dir):
+    """
+    Test that calling transcribe with timestamps=True multiple times
+    does not reinstantiate the decoder.
+
+    Regression test for the fix that avoids calling change_decoding_strategy()
+    when compute_timestamps is already set to the desired value.
+    """
+    model = stt_en_fastconformer_transducer_large
+    audio_file = os.path.join(test_data_dir, "asr/test/an4/wav/cen3-mjwl-b.wav")
+
+    # First call - may change decoding strategy
+    _ = model.transcribe(audio_file, timestamps=True)
+
+    # Get reference to decoding algorithm after first call
+    decoding_after_first_call = model.decoding.decoding
+
+    # Second call - should NOT reinstantiate decoder
+    _ = model.transcribe(audio_file, timestamps=True)
+
+    # Verify decoder is the same object (not reinstantiated)
+    assert model.decoding.decoding is decoding_after_first_call, (
+        "Decoder was reinstantiated on second transcribe call with timestamps=True. "
+        "This indicates change_decoding_strategy() was called unnecessarily."
+    )

tests/collections/asr/decoding/test_streaming_beam_decoding.py

@@ -0,0 +1,600 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Streaming beam-search (MALSD + MAES) decoding tests.
+
+Beam-search analogue of ``test_streaming_decoding.py``. Exercises the streaming
+path of the batched beam-search computers by manually feeding the encoder
+output to ``model.decoding.decoding.decoding_computer`` in chunks and threading
+``prev_batched_state`` (a ``BatchedBeamState``):
+
+  - :func:`test_malsd_streaming_batched_state` -- covers RNNT and TDT MALSD
+    (:class:`ModifiedALSDBatchedRNNTComputer`, :class:`ModifiedALSDBatchedTDTComputer`)
+    across the eager path and both captured-graph variants (``full_graph`` and
+    ``no_while_loops``).
+  - :func:`test_malsd_streaming_batched_state_with_word_boosting` -- same MALSD matrix
+    but with a ``GPUBoostingTreeModel`` fusion model plugged in (``boosting_tree.
+    key_phrases_list``); exercises cross-chunk restoration of per-beam fusion states
+    in :meth:`_init_decoding_state`.
+  - :func:`test_maes_streaming_batched_state` -- covers RNNT MAES
+    (:class:`ModifiedAESBatchedRNNTComputer`); MAES is RNNT-only and pure-PyTorch,
+    so there is no ``is_tdt`` / ``cuda_graphs_mode`` axis.
+
+Per-chunk results are merged into a single ``BatchedBeamHyps`` via
+``flatten_()`` + ``merge_(..., is_chunk_continuation=True,
+boundary_prev_ptr=...)`` -- the same accumulation pattern used by the
+cache-aware / chunked streaming inference scripts.
+
+Streamed transcripts are asserted to be identical to the non-streaming
+reference produced by ``model.transcribe`` with the same beam settings:
+beam search with ``prev_batched_state`` is chunk-invariant because all
+cross-chunk per-beam state (scores, ``last_label``, decoded lengths,
+decoder + fusion states, ``last_timestamp_lasts``, ...) is preserved across
+boundaries.
+"""
+
+import copy
+from typing import Optional
+
+import pytest
+import torch
+from omegaconf import open_dict
+from tqdm.auto import tqdm
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.context_biasing.biasing_multi_model import BiasingRequestItemConfig
+from nemo.collections.asr.parts.context_biasing.boosting_graph_batched import BoostingTreeModelConfig
+from nemo.collections.asr.parts.submodules.transducer_decoding.label_looping_base import BatchedBeamState
+from nemo.collections.asr.parts.utils.batched_beam_decoding_utils import BatchedBeamHyps
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+from tests.collections.asr.decoding.utils import load_audio, make_preprocessor_deterministic
+
+
+def get_devices_for_testing(use_cpu_always: bool = False) -> list[torch.device]:
+    devices = [torch.device("cpu")] if use_cpu_always else []
+    if torch.cuda.is_available():
+        devices.append(torch.device("cuda:0"))
+
+    if torch.mps.is_available():
+        devices.append(torch.device("mps"))
+
+    if len(devices) == 0:
+        # no fast device for testing, add CPU
+        devices.append(torch.device("cpu"))
+    return devices
+
+
+DEVICES = get_devices_for_testing(use_cpu_always=True)
+
+
+def _make_device_param_matrix() -> list:
+    """
+    Build the ``(device, cuda_graphs_mode)`` parametrize entries with explicit, readable
+    pytest IDs (``cpu-no-graphs``, ``cuda-full-graph``, ``cuda-no-while-loops``, ...) so the
+    test matrix shows which device + graph-mode pair is exercised instead of opaque
+    ``device0`` / ``device1`` ids.
+
+    ``cuda_graphs_mode`` is ``None`` for the eager path or one of the
+    :class:`ModifiedALSDBatched{RNNT,TDT}Computer.CudaGraphsMode` string values
+    (``"full_graph"`` / ``"no_while_loops"``) for the two captured-graph variants. The test
+    uses ``force_cuda_graphs_mode`` to pin the variant explicitly so each captured path is
+    actually exercised (otherwise ``maybe_enable_cuda_graphs`` would always pick
+    ``full_graph`` whenever conditional nodes are supported).
+
+    Coverage:
+      - every device in ``DEVICES`` with ``cuda_graphs_mode=None`` (eager path).
+      - every CUDA device additionally with both ``"full_graph"`` and ``"no_while_loops"``.
+    """
+    entries: list = []
+    for device in DEVICES:
+        entries.append(pytest.param(device, None, id=f"{device.type}-no-graphs"))
+    for device in DEVICES:
+        if device.type == "cuda":
+            entries.append(pytest.param(device, "full_graph", id=f"{device.type}-full-graph"))
+            entries.append(pytest.param(device, "no_while_loops", id=f"{device.type}-no-while-loops"))
+    return entries
+
+
+DEVICE_PARAM_MATRIX = _make_device_param_matrix()
+
+
+def _make_maes_device_param_matrix() -> list:
+    """Build readable ``device`` parametrize entries for MAES tests.
+
+    MAES has no CUDA-graphs path (it's pure-PyTorch), so the matrix is just one entry per
+    available device with explicit IDs (``cpu``, ``cuda``, ...) instead of opaque
+    ``device0`` / ``device1``.
+    """
+    return [pytest.param(device, id=device.type) for device in DEVICES]
+
+
+MAES_DEVICE_PARAM_MATRIX = _make_maes_device_param_matrix()
+
+
+def get_model_encoder_output(
+    test_audio_filenames,
+    num_samples: int,
+    model: ASRModel,
+    device: torch.device = torch.device("cpu"),
+    dtype: torch.dtype = torch.float32,
+):
+    audio_filepaths = test_audio_filenames[:num_samples]
+
+    with torch.no_grad():
+        make_preprocessor_deterministic(model)
+        model.eval()
+
+        all_inputs, all_lengths = [], []
+        for audio_file in tqdm(audio_filepaths, desc="Loading audio files"):
+            audio_tensor, _ = load_audio(audio_file)
+            all_inputs.append(audio_tensor)
+            all_lengths.append(torch.tensor(audio_tensor.shape[0], dtype=torch.int64))
+
+        input_batch = torch.nn.utils.rnn.pad_sequence(all_inputs, batch_first=True).to(device=device, dtype=dtype)
+        length_batch = torch.tensor(all_lengths, dtype=torch.int64).to(device)
+
+        encoded_outputs, encoded_length = model(input_signal=input_batch, input_signal_length=length_batch)
+
+    return encoded_outputs, encoded_length
+
+
+def get_batch_encoder_outputs_from_records(records, model, device):
+    """Helper function to get encoder outputs for a batch of manifest records"""
+    filenames = [record["audio_filepath"] for record in records]
+    local_batch_size = len(filenames)
+    encoder_output, encoder_output_len = get_model_encoder_output(
+        test_audio_filenames=filenames, model=model, num_samples=local_batch_size, device=device
+    )
+    return encoder_output, encoder_output_len
+
+
+def _configure_malsd_decoding(
+    model: ASRModel,
+    cuda_graphs_mode: Optional[str],
+    beam_size: int,
+    max_symbols: int,
+    key_phrases_list: Optional[list[str]] = None,
+    boosting_tree_alpha: float = 1.0,
+    enable_per_stream_biasing: bool = False,
+) -> None:
+    """Switch ``model`` to the ``malsd_batch`` beam-search strategy used by the streaming tests.
+
+    ``cuda_graphs_mode`` is the CudaGraphsMode string (``"full_graph"`` or
+    ``"no_while_loops"``) when CUDA graphs should be used, or ``None`` for the eager path.
+    When non-None we also call ``force_cuda_graphs_mode`` to pin the variant after the
+    decoding strategy is swapped in -- ``maybe_enable_cuda_graphs`` would otherwise auto-pick
+    ``full_graph`` whenever conditional nodes are supported, leaving the ``no_while_loops``
+    branch effectively untested.
+
+    When ``key_phrases_list`` is given, a ``boosting_tree`` fusion model is plugged in
+    (``BoostingTreeModelConfig.key_phrases_list``) so the streaming path exercises
+    cross-chunk fusion-state restoration in :meth:`_init_decoding_state`.
+    """
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "malsd_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.beam.beam_size = beam_size
+        decoding_cfg.beam.max_symbols = max_symbols
+        decoding_cfg.beam.allow_cuda_graphs = cuda_graphs_mode is not None
+        decoding_cfg.beam.return_best_hypothesis = True
+        decoding_cfg.beam.score_norm = True
+        if key_phrases_list is not None:
+            decoding_cfg.beam.boosting_tree = {"key_phrases_list": list(key_phrases_list)}
+            decoding_cfg.beam.boosting_tree_alpha = boosting_tree_alpha
+        if enable_per_stream_biasing:
+            decoding_cfg.beam.enable_per_stream_biasing = True
+    model.change_decoding_strategy(decoding_cfg)
+    if cuda_graphs_mode is not None:
+        model.decoding.decoding.decoding_computer.force_cuda_graphs_mode(cuda_graphs_mode)
+
+
+def _configure_maes_decoding(
+    model: ASRModel,
+    beam_size: int,
+    maes_num_steps: int,
+    maes_expansion_beta: int,
+    maes_expansion_gamma: float,
+) -> None:
+    """Switch ``model`` to the ``maes_batch`` beam-search strategy used by the streaming tests.
+
+    MAES is RNNT-only and currently pure-PyTorch (CUDA graphs are not implemented;
+    ``allow_cuda_graphs`` is accepted only for API parity with MALSD and is ignored by
+    :class:`ModifiedAESBatchedRNNTComputer`).
+    """
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "maes_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.beam.beam_size = beam_size
+        decoding_cfg.beam.maes_num_steps = maes_num_steps
+        decoding_cfg.beam.maes_expansion_beta = maes_expansion_beta
+        decoding_cfg.beam.maes_expansion_gamma = maes_expansion_gamma
+        decoding_cfg.beam.allow_cuda_graphs = False
+        decoding_cfg.beam.return_best_hypothesis = True
+        decoding_cfg.beam.score_norm = True
+    model.change_decoding_strategy(decoding_cfg)
+
+
+def _reset_decoding_computer_state(model: ASRModel) -> None:
+    decoding_computer = model.decoding.decoding.decoding_computer
+    if hasattr(decoding_computer, "reset_cuda_graphs_state"):
+        decoding_computer.reset_cuda_graphs_state()
+
+
+def _decode_malsd_encoder_in_chunks(
+    decoding_computer,
+    encoder_output: torch.Tensor,
+    encoder_output_len: torch.Tensor,
+    chunk_size: int,
+    multi_biasing_ids: Optional[torch.Tensor] = None,
+) -> BatchedBeamHyps:
+    encoder_output = encoder_output.transpose(1, 2)
+    state: Optional[BatchedBeamState] = None
+    current_batched_hyps: BatchedBeamHyps | None = None
+    decode_kwargs = {}
+    if multi_biasing_ids is not None:
+        decode_kwargs["multi_biasing_ids"] = multi_biasing_ids
+
+    for t in range(0, encoder_output.shape[1], chunk_size):
+        rest_len = encoder_output_len - t
+        current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+        current_len = torch.minimum(current_len, rest_len)
+        current_len = torch.maximum(current_len, torch.zeros_like(current_len))
+        chunk_batched_hyps, state = decoding_computer(
+            x=encoder_output[:, t : t + chunk_size],
+            out_len=current_len,
+            prev_batched_state=state,
+            **decode_kwargs,
+        )
+        chunk_root_ptrs = chunk_batched_hyps.flatten_()
+        if current_batched_hyps is None:
+            current_batched_hyps = chunk_batched_hyps
+        else:
+            current_batched_hyps.merge_(
+                chunk_batched_hyps,
+                is_chunk_continuation=True,
+                boundary_prev_ptr=chunk_root_ptrs,
+            )
+
+    assert current_batched_hyps is not None
+    return current_batched_hyps
+
+
+def _register_per_stream_biasing(
+    decoding_computer,
+    tokenizer,
+    boost_texts: list[str],
+    device: torch.device,
+    boosting_model_alpha: float = 10.0,
+) -> tuple[torch.Tensor, list[BiasingRequestItemConfig | None]]:
+    batch_size = len(boost_texts)
+    multi_biasing_ids = torch.full([batch_size], fill_value=-1, dtype=torch.long, device=device)
+    biasing_requests: list[BiasingRequestItemConfig | None] = []
+
+    for batch_idx, boost_text in enumerate(boost_texts):
+        if not boost_text:
+            biasing_requests.append(None)
+            continue
+        request = BiasingRequestItemConfig(
+            boosting_model_cfg=BoostingTreeModelConfig(key_phrases_list=[boost_text], unk_score=-100),
+            boosting_model_alpha=boosting_model_alpha,
+        )
+        request.add_to_multi_model(
+            tokenizer=tokenizer,
+            biasing_multi_model=decoding_computer.biasing_multi_model,
+        )
+        if request.multi_model_id is not None:
+            multi_biasing_ids[batch_idx] = request.multi_model_id
+        biasing_requests.append(request)
+
+    return multi_biasing_ids, biasing_requests
+
+
+def _unregister_per_stream_biasing(decoding_computer, biasing_requests: list[BiasingRequestItemConfig | None]) -> None:
+    for request in biasing_requests:
+        if request is not None and request.multi_model_id is not None:
+            decoding_computer.biasing_multi_model.remove_model(request.multi_model_id)
+            request.multi_model_id = None
+
+
+def _run_malsd_streaming_manifest(
+    model: ASRModel,
+    manifest_path,
+    device: torch.device,
+    chunk_size: int,
+    batch_size: int,
+    boost_texts: Optional[list[str]] = None,
+    boosting_model_alpha: float = 10.0,
+) -> list[str]:
+    manifest = read_manifest(manifest_path)
+    decoding_computer = model.decoding.decoding.decoding_computer
+    all_transcripts: list[str] = []
+
+    with torch.no_grad(), torch.inference_mode():
+        for i in range(0, len(manifest), batch_size):
+            batch_records = manifest[i : i + batch_size]
+            encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+                batch_records, model=model, device=device
+            )
+
+            multi_biasing_ids = None
+            biasing_requests: list[BiasingRequestItemConfig | None] = []
+            if boost_texts is not None:
+                assert decoding_computer.biasing_multi_model is not None
+                batch_boost_texts = boost_texts[i : i + batch_size]
+                multi_biasing_ids, biasing_requests = _register_per_stream_biasing(
+                    decoding_computer,
+                    model.tokenizer,
+                    batch_boost_texts,
+                    device,
+                    boosting_model_alpha=boosting_model_alpha,
+                )
+
+            batched_hyps = _decode_malsd_encoder_in_chunks(
+                decoding_computer,
+                encoder_output,
+                encoder_output_len,
+                chunk_size,
+                multi_biasing_ids=multi_biasing_ids,
+            )
+            if boost_texts is not None:
+                _unregister_per_stream_biasing(decoding_computer, biasing_requests)
+
+            all_transcripts.extend(
+                model.tokenizer.ids_to_text(hyp.y_sequence.tolist())
+                for hyp in batched_hyps.to_hyps_list(score_norm=True)
+            )
+
+    return all_transcripts
+
+
+def _run_streaming_batched_state(
+    model: ASRModel,
+    manifest_path,
+    device: torch.device,
+    chunk_size: int,
+    batch_size: int,
+) -> tuple[list[str], list[str]]:
+    """Drive the model's beam-search ``decoding_computer`` chunk-by-chunk and return
+    ``(ref_transcripts, streaming_transcripts)``.
+
+    Shared between the MALSD and MAES streaming tests: both decoders return a
+    ``(BatchedBeamHyps, None, BatchedBeamState)`` triple and accept
+    ``prev_batched_state`` for cross-chunk state threading. The per-chunk results are
+    flattened and merged into a single accumulator via ``flatten_()`` +
+    ``merge_(..., is_chunk_continuation=True, boundary_prev_ptr=...)`` -- the same
+    accumulation pattern used by the cache-aware / chunked streaming inference scripts.
+    """
+    manifest = read_manifest(manifest_path)
+    transcriptions = model.transcribe(audio=str(manifest_path.absolute()), batch_size=batch_size)
+    ref_transcripts = [hyp.text for hyp in transcriptions]
+
+    all_hyps = []
+    decoding_computer = model.decoding.decoding.decoding_computer
+    with torch.no_grad(), torch.inference_mode():
+        for i in range(0, len(manifest), batch_size):
+            encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+                manifest[i : i + batch_size], model=model, device=device
+            )
+            state: Optional[BatchedBeamState] = None
+            current_batched_hyps: BatchedBeamHyps | None = None
+            encoder_output = encoder_output.transpose(1, 2)  # (B, T, D)
+            for t in range(0, encoder_output.shape[1], chunk_size):
+                rest_len = encoder_output_len - t
+                current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+                current_len = torch.minimum(current_len, rest_len)
+                current_len = torch.maximum(current_len, torch.zeros_like(current_len))
+                chunk_batched_hyps, state = decoding_computer(
+                    x=encoder_output[:, t : t + chunk_size],
+                    out_len=current_len,
+                    prev_batched_state=state,
+                )
+                # Flatten this chunk's prefix tree and thread the cross-chunk beam
+                # permutation (``root_ptrs``) into the accumulator so the final
+                # ``flatten_sort_`` walks back through the right beam history.
+                chunk_root_ptrs = chunk_batched_hyps.flatten_()
+                if current_batched_hyps is None:
+                    current_batched_hyps = chunk_batched_hyps
+                else:
+                    current_batched_hyps.merge_(
+                        chunk_batched_hyps,
+                        is_chunk_continuation=True,
+                        boundary_prev_ptr=chunk_root_ptrs,
+                    )
+            assert current_batched_hyps is not None
+            # ``to_hyps_list`` mutates the prefix tree via ``flatten_sort_``, but we're done
+            # with ``current_batched_hyps`` here so an in-place call is fine.
+            all_hyps.extend(current_batched_hyps.to_hyps_list(score_norm=True))
+
+    streaming_transcripts = [model.tokenizer.ids_to_text(hyp.y_sequence.tolist()) for hyp in all_hyps]
+    return ref_transcripts, streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize("device,cuda_graphs_mode", DEVICE_PARAM_MATRIX)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("beam_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_malsd_streaming_batched_state(
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    cuda_graphs_mode: Optional[str],
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    beam_size: int,
+    max_symbols: int,
+):
+    """Streaming MALSD decoding with batched beam state passed across chunks."""
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    _configure_malsd_decoding(model, cuda_graphs_mode, beam_size=beam_size, max_symbols=max_symbols)
+
+    ref_transcripts, streaming_transcripts = _run_streaming_batched_state(
+        model=model,
+        manifest_path=an4_val_manifest_corrected,
+        device=device,
+        chunk_size=chunk_size,
+        batch_size=batch_size,
+    )
+    assert ref_transcripts == streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize("device", MAES_DEVICE_PARAM_MATRIX)
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("beam_size", [4])
+@pytest.mark.parametrize("maes_num_steps", [2])
+@pytest.mark.parametrize("maes_expansion_beta", [2])
+@pytest.mark.parametrize("maes_expansion_gamma", [2.3])
+def test_maes_streaming_batched_state(
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    device: torch.device,
+    chunk_size: int,
+    batch_size: int,
+    beam_size: int,
+    maes_num_steps: int,
+    maes_expansion_beta: int,
+    maes_expansion_gamma: float,
+):
+    """Streaming MAES decoding with batched beam state passed across chunks.
+
+    MAES is RNNT-only and pure-PyTorch (no CUDA graphs path), so the device matrix is
+    just the set of available devices and there is no ``cuda_graphs_mode`` / ``is_tdt``
+    parameter.
+    """
+    model = stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    _configure_maes_decoding(
+        model,
+        beam_size=beam_size,
+        maes_num_steps=maes_num_steps,
+        maes_expansion_beta=maes_expansion_beta,
+        maes_expansion_gamma=maes_expansion_gamma,
+    )
+
+    ref_transcripts, streaming_transcripts = _run_streaming_batched_state(
+        model=model,
+        manifest_path=an4_val_manifest_corrected,
+        device=device,
+        chunk_size=chunk_size,
+        batch_size=batch_size,
+    )
+    assert ref_transcripts == streaming_transcripts
+
+
+# Phrases chosen from the AN4 val transcripts so the boosting tree is actually exercised
+# (an empty / unseen list collapses to the no-boosting path and would not test fusion-state
+# restoration across chunks).
+_WB_KEY_PHRASES: list[str] = ["nineteen", "forty", "fifty", "repeat", "stop", "yes"]
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize("device,cuda_graphs_mode", DEVICE_PARAM_MATRIX)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("beam_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_malsd_streaming_batched_state_with_word_boosting(
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    cuda_graphs_mode: Optional[str],
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    beam_size: int,
+    max_symbols: int,
+):
+    """Streaming MALSD with word-boosting (``boosting_tree``) is chunk-invariant.
+
+    Adds a ``GPUBoostingTreeModel`` fusion model on top of the standard streaming MALSD
+    test. The reference (``model.transcribe``) and the streaming path are configured
+    identically, so the boosting tree's per-beam fusion states must be restored across
+    chunks via ``_init_decoding_state`` for the two transcripts to match.
+    """
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    _configure_malsd_decoding(
+        model,
+        cuda_graphs_mode,
+        beam_size=beam_size,
+        max_symbols=max_symbols,
+        key_phrases_list=_WB_KEY_PHRASES,
+    )
+
+    ref_transcripts, streaming_transcripts = _run_streaming_batched_state(
+        model=model,
+        manifest_path=an4_val_manifest_corrected,
+        device=device,
+        chunk_size=chunk_size,
+        batch_size=batch_size,
+    )
+    assert ref_transcripts == streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize("device,cuda_graphs_mode", DEVICE_PARAM_MATRIX)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("beam_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_malsd_streaming_boosting_with_ref_transcripts(
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    cuda_graphs_mode: Optional[str],
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    beam_size: int,
+    max_symbols: int,
+):
+    """Metamorphic test analogous to ``test_label_looping_streaming_boosting_with_ref_transcripts``."""
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    _configure_malsd_decoding(model, cuda_graphs_mode, beam_size, max_symbols)
+    ref_transcripts = [
+        hyp.text for hyp in model.transcribe(audio=str(an4_val_manifest_corrected.absolute()), batch_size=batch_size)
+    ]
+
+    _configure_malsd_decoding(model, cuda_graphs_mode, beam_size, max_symbols, enable_per_stream_biasing=True)
+    _reset_decoding_computer_state(model)
+
+    streaming_transcripts = _run_malsd_streaming_manifest(
+        model,
+        an4_val_manifest_corrected,
+        device,
+        chunk_size,
+        batch_size,
+        boost_texts=ref_transcripts,
+    )
+    assert ref_transcripts == streaming_transcripts

tests/collections/asr/decoding/test_streaming_buffer.py

@@ -0,0 +1,469 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for `StreamingBatchedAudioBuffer` and accompanying helper
+classes defined in
+`nemo.collections.asr.parts.utils.streaming_utils`.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+import torch
+
+from nemo.collections.asr.parts.utils.streaming_utils import (
+    ContextSize,
+    ContextSizeBatch,
+    DynamicLengthTensor,
+    StreamingBatchedAudioBuffer,
+)
+
+# -----------------------------------------------------------------------------
+# Helper constants / fixtures
+# -----------------------------------------------------------------------------
+
+DEVICES: list[torch.device] = [torch.device("cpu")]
+if torch.cuda.is_available():
+    DEVICES.append(torch.device("cuda:0"))
+
+
+def _create_audio_batch(batch_size: int, length: int, device: torch.device, dtype: torch.dtype = torch.float32):
+    """Create a dummy audio batch of shape (batch_size, length)."""
+    # Use a simple ramp signal to ease debugging.
+    vals = torch.arange(batch_size * length, device=device, dtype=dtype)
+    return vals.view(batch_size, length)
+
+
+def _make_chunk(batch_size: int, length: int, channels: int, start: float, device: torch.device) -> torch.Tensor:
+    """Create a deterministic chunk of shape (batch_size, length, channels)."""
+    n = batch_size * length * channels
+    return (start + torch.arange(n, device=device, dtype=torch.float32)).view(batch_size, length, channels)
+
+
+def _make_ndim_chunk(
+    batch_size: int, length: int, dim_shape: list[int], start: float, device: torch.device
+) -> torch.Tensor:
+    """Create a deterministic chunk of shape (batch_size, length, *dim_shape)."""
+    n = batch_size * length * math.prod(dim_shape)
+    return (start + torch.arange(n, device=device, dtype=torch.float32)).view(batch_size, length, *dim_shape)
+
+
+# -----------------------------------------------------------------------------
+# Tests for ContextSize and ContextSizeBatch
+# -----------------------------------------------------------------------------
+
+
+class TestContextSize:
+    @pytest.mark.unit
+    def test_context_size_total_and_subsample(self):
+        ctx = ContextSize(left=4, chunk=2, right=1)
+        assert ctx.total() == 7
+
+        half_ctx = ctx.subsample(factor=2)
+        assert isinstance(half_ctx, ContextSize)
+        assert half_ctx.left == 2 and half_ctx.chunk == 1 and half_ctx.right == 0
+        assert half_ctx.total() == math.floor(7 / 2)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_context_size_batch_total_and_subsample(self, device: torch.device):
+        left = torch.tensor([4, 4], dtype=torch.long, device=device)
+        chunk = torch.tensor([2, 2], dtype=torch.long, device=device)
+        right = torch.tensor([2, 2], dtype=torch.long, device=device)
+        batch_ctx = ContextSizeBatch(left=left, chunk=chunk, right=right)
+
+        # total() should equal element-wise sum
+        expected_total = left + chunk + right
+        assert torch.equal(batch_ctx.total(), expected_total)
+
+        # After subsampling by 2 each component should be halved (floor division)
+        half_ctx = batch_ctx.subsample(2)
+        assert torch.equal(half_ctx.left, left // 2)
+        assert torch.equal(half_ctx.chunk, chunk // 2)
+        assert torch.equal(half_ctx.right, right // 2)
+
+
+# -----------------------------------------------------------------------------
+# Tests for StreamingBatchedAudioBuffer
+# -----------------------------------------------------------------------------
+
+
+class TestStreamingBatchedAudioBuffer:
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_streaming_batched_audio_buffer(self, device: torch.device):
+        batch_size = 2
+        expected_ctx = ContextSize(left=4, chunk=2, right=1)  # total = 7
+        buffer = StreamingBatchedAudioBuffer(
+            batch_size=batch_size,
+            context_samples=expected_ctx,
+            dtype=torch.float32,
+            device=device,
+        )
+
+        # ------------------------------------------------------------------
+        # First add : chunk + right (filling initial buffer)
+        # ------------------------------------------------------------------
+        first_len = expected_ctx.chunk + expected_ctx.right  # 3
+        audio_batch = _create_audio_batch(batch_size, first_len, device)
+        audio_lens = torch.full(
+            [
+                batch_size,
+            ],
+            first_len,
+            dtype=torch.long,
+            device=device,
+        )
+        buffer.add_audio_batch_(
+            audio_batch=audio_batch,
+            audio_lengths=audio_lens,
+            is_last_chunk=False,
+            is_last_chunk_batch=torch.zeros(batch_size, dtype=torch.bool, device=device),
+        )
+
+        # Validate context sizes
+        assert buffer.context_size.left == 0
+        assert buffer.context_size.chunk == expected_ctx.chunk
+        assert buffer.context_size.right == expected_ctx.right
+        assert buffer.samples.shape[1] == first_len  # No truncation yet
+
+        # ------------------------------------------------------------------
+        # Second add : only chunk length
+        # ------------------------------------------------------------------
+        chunk_len = expected_ctx.chunk  # 2
+        audio_batch = _create_audio_batch(batch_size, chunk_len, device)
+        audio_lens.fill_(chunk_len)
+        buffer.add_audio_batch_(
+            audio_batch=audio_batch,
+            audio_lengths=audio_lens,
+            is_last_chunk=False,
+            is_last_chunk_batch=torch.zeros(batch_size, dtype=torch.bool, device=device),
+        )
+
+        # After second add, left should have grown by previous chunk (2)
+        assert buffer.context_size.left == 2
+        assert buffer.context_size.chunk == expected_ctx.chunk
+        assert buffer.context_size.right == expected_ctx.right
+        assert buffer.samples.shape[1] == 5  # 2 (left) + 2 (chunk) + 1 (right)
+
+        # ------------------------------------------------------------------
+        # Third add : another chunk, buffer should now reach full capacity (7)
+        # ------------------------------------------------------------------
+        buffer.add_audio_batch_(
+            audio_batch=audio_batch,
+            audio_lengths=audio_lens,
+            is_last_chunk=False,
+            is_last_chunk_batch=torch.zeros(batch_size, dtype=torch.bool, device=device),
+        )
+
+        assert buffer.samples.shape[1] == expected_ctx.total()
+        assert buffer.context_size.total() == expected_ctx.total()
+
+        # ------------------------------------------------------------------
+        # Fourth add : buffer overflows by 2 samples; implementation should
+        # drop the excess from the left context.
+        # ------------------------------------------------------------------
+        buffer.add_audio_batch_(
+            audio_batch=audio_batch,
+            audio_lengths=audio_lens,
+            is_last_chunk=False,
+            is_last_chunk_batch=torch.zeros(batch_size, dtype=torch.bool, device=device),
+        )
+
+        # Buffer length remains constant (total context size)
+        assert buffer.samples.shape[1] == expected_ctx.total()
+        assert buffer.context_size.total() == expected_ctx.total()
+
+        # Left context should have been clipped by 2 samples (from 6 to 4)
+        assert buffer.context_size.left == expected_ctx.left  # 4
+
+        # ------------------------------------------------------------------
+        # Final add : mark last chunk with shorter length; right context
+        # should go to 0 afterwards.
+        # ------------------------------------------------------------------
+        last_len = 1
+        audio_batch = _create_audio_batch(batch_size, last_len, device)
+        audio_lens.fill_(last_len)
+        buffer.add_audio_batch_(
+            audio_batch=audio_batch,
+            audio_lengths=audio_lens,
+            is_last_chunk=True,
+            is_last_chunk_batch=torch.ones(batch_size, dtype=torch.bool, device=device),
+        )
+
+        # After last chunk, right context must be zero and total size preserved
+        assert buffer.context_size.right == 0
+        assert buffer.context_size.total() == expected_ctx.total()
+        assert buffer.samples.shape[1] == expected_ctx.total()
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_streaming_batched_audio_buffer_raises_on_too_long_chunk(self, device: torch.device):
+        """`add_audio_batch_` should raise if provided chunk is larger than chunk + right."""
+
+        expected_ctx = ContextSize(left=0, chunk=2, right=1)
+        buffer = StreamingBatchedAudioBuffer(
+            batch_size=1,
+            context_samples=expected_ctx,
+            dtype=torch.float32,
+            device=device,
+        )
+
+        # Attempt to add a chunk that is too long (4 > 3)
+        too_long_chunk_size = expected_ctx.chunk + expected_ctx.right + 1
+        audio = _create_audio_batch(1, too_long_chunk_size, device)
+        audio_lens = torch.tensor([too_long_chunk_size], dtype=torch.long, device=device)
+
+        with pytest.raises(ValueError):
+            buffer.add_audio_batch_(
+                audio_batch=audio,
+                audio_lengths=audio_lens,
+                is_last_chunk=False,
+                is_last_chunk_batch=torch.tensor([False], dtype=torch.bool, device=device),
+            )
+
+
+# -----------------------------------------------------------------------------
+# Tests for DynamicLengthTensor
+# -----------------------------------------------------------------------------
+
+
+class TestDynamicLengthTensor:
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize(
+        "dim_shape, expected_dim_shape",
+        [
+            (None, []),
+            (3, [3]),
+            ([4, 5], [4, 5]),
+        ],
+    )
+    def test_init(self, device, dim_shape, expected_dim_shape):
+        batch_size, init_length = 2, 5
+        t = DynamicLengthTensor(
+            batch_size=batch_size,
+            init_length=init_length,
+            dim_shape=dim_shape,
+            device=device,
+            dtype=torch.float32,
+        )
+
+        assert t.dim_shape == expected_dim_shape
+        assert list(t.data.shape) == [batch_size, init_length, *expected_dim_shape]
+        assert list(t.lengths.shape) == [batch_size]
+        assert t.lengths.dtype == torch.long
+        assert t.data.dtype == torch.float32
+        # Freshly created storage is zeroed and reports no content.
+        assert torch.count_nonzero(t.lengths) == 0
+        assert torch.count_nonzero(t.data) == 0
+
+    @pytest.mark.unit
+    def test_init_minimum_length(self):
+        """`init_length` is clamped to at least 1 so doubling-based growth works."""
+        t = DynamicLengthTensor(batch_size=2, init_length=0, dim_shape=1)
+        assert t._max_length == 1
+        assert t.data.shape[1] == 1
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_append_with_lengths(self, device):
+        """Per-batch `lengths` control how many frames become valid for each item."""
+        t = DynamicLengthTensor(batch_size=2, init_length=4, dim_shape=1, device=device, dtype=torch.float32)
+
+        # First chunk: batch item 0 keeps 2 frames, item 1 keeps 1 frame.
+        chunk1 = _make_chunk(batch_size=2, length=2, channels=1, start=10.0, device=device)
+        # chunk1 == [[[10], [11]], [[12], [13]]]
+        t.append_(data=chunk1, lengths=torch.tensor([2, 1], device=device))
+        assert t.lengths.tolist() == [2, 1]
+
+        # Second chunk: item 0 keeps 1 frame, item 1 keeps 2 frames. Item 1 should
+        # overwrite the previously written "garbage" frame at position 1.
+        chunk2 = _make_chunk(batch_size=2, length=2, channels=1, start=30.0, device=device)
+        # chunk2 == [[[30], [31]], [[32], [33]]]
+        t.append_(data=chunk2, lengths=torch.tensor([1, 2], device=device))
+        assert t.lengths.tolist() == [3, 3]
+
+        # Valid frames are everything up to the per-item length.
+        item0 = t.data[0, : t.lengths[0], 0].tolist()
+        item1 = t.data[1, : t.lengths[1], 0].tolist()
+        assert item0 == [10.0, 11.0, 30.0]
+        assert item1 == [12.0, 32.0, 33.0]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_append_without_lengths(self, device):
+        """Without `lengths`, every frame in the chunk is appended for all items."""
+        t = DynamicLengthTensor(batch_size=2, init_length=2, dim_shape=1, device=device, dtype=torch.float32)
+
+        chunk = _make_chunk(batch_size=2, length=3, channels=1, start=0.0, device=device)
+        t.append_(data=chunk)
+
+        assert t.lengths.tolist() == [3, 3]
+        assert torch.equal(t.data[:, :3], chunk)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("dim_shape", [[], [3], [2, 3], [2, 3, 4]])
+    def test_append_without_lengths_multidim(self, device, dim_shape):
+        """Append must scatter the whole feature vector for arbitrary trailing `dim_shape`."""
+        batch_size = 2
+        # init_length < appended length so this also exercises the growth path with multi-dim shapes.
+        t = DynamicLengthTensor(
+            batch_size=batch_size, init_length=2, dim_shape=dim_shape, device=device, dtype=torch.float32
+        )
+
+        chunk = _make_ndim_chunk(batch_size, length=3, dim_shape=dim_shape, start=0.0, device=device)
+        t.append_(data=chunk)
+
+        assert t.lengths.tolist() == [3, 3]
+        assert torch.equal(t.data[:, :3], chunk)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_append_with_lengths_multidim(self, device):
+        """Per-batch `lengths` must place the full multi-dim feature vectors at the right offsets."""
+        dim_shape = [2, 3]
+        t = DynamicLengthTensor(batch_size=2, init_length=4, dim_shape=dim_shape, device=device, dtype=torch.float32)
+
+        # First chunk: item 0 keeps 2 frames, item 1 keeps 1 frame.
+        chunk1 = _make_ndim_chunk(2, length=2, dim_shape=dim_shape, start=10.0, device=device)
+        t.append_(data=chunk1, lengths=torch.tensor([2, 1], device=device))
+        assert t.lengths.tolist() == [2, 1]
+
+        # Second chunk: item 0 keeps 1 frame, item 1 keeps 2 frames. Item 1's second frame
+        # must overwrite the previously written "garbage" frame at position 1.
+        chunk2 = _make_ndim_chunk(2, length=2, dim_shape=dim_shape, start=100.0, device=device)
+        t.append_(data=chunk2, lengths=torch.tensor([1, 2], device=device))
+        assert t.lengths.tolist() == [3, 3]
+
+        # Item 0: chunk1[0, 0], chunk1[0, 1], chunk2[0, 0]; each is a full (2, 3) feature vector.
+        assert torch.equal(t.data[0, 0], chunk1[0, 0])
+        assert torch.equal(t.data[0, 1], chunk1[0, 1])
+        assert torch.equal(t.data[0, 2], chunk2[0, 0])
+        # Item 1: chunk1[1, 0], chunk2[1, 0] (overwrites garbage at pos 1), chunk2[1, 1].
+        assert torch.equal(t.data[1, 0], chunk1[1, 0])
+        assert torch.equal(t.data[1, 1], chunk2[1, 0])
+        assert torch.equal(t.data[1, 2], chunk2[1, 1])
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_growth_preserves_data(self, device):
+        """Appending more than the initial capacity reallocates and keeps content."""
+        t = DynamicLengthTensor(batch_size=1, init_length=2, dim_shape=1, device=device, dtype=torch.float32)
+        initial_capacity = t._max_length
+
+        big_len = 10
+        chunk = _make_chunk(batch_size=1, length=big_len, channels=1, start=0.0, device=device)
+        t.append_(data=chunk, lengths=torch.tensor([big_len], device=device))
+
+        assert t._max_length > initial_capacity
+        assert t._max_length >= big_len
+        assert t.lengths.tolist() == [big_len]
+        assert t.data[0, :big_len, 0].tolist() == list(range(big_len))
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_incremental_appends_double_capacity(self, device):
+        """Repeated single-frame appends grow capacity geometrically (amortized O(1))."""
+        n_appends = 9
+        t = DynamicLengthTensor(batch_size=1, init_length=1, dim_shape=1, device=device, dtype=torch.float32)
+
+        capacities = []
+        for i in range(n_appends):
+            frame = torch.full((1, 1, 1), float(i), device=device)
+            t.append_(data=frame)
+            capacities.append(t._max_length)
+
+        # Everything that was appended is retained, in order.
+        assert t.lengths.tolist() == [n_appends]
+        assert t.data[0, :n_appends, 0].tolist() == [float(i) for i in range(n_appends)]
+        # Capacity is always at least what is stored, and grew far less than linearly.
+        assert t._max_length >= n_appends
+        assert len(set(capacities)) < n_appends  # capacity reused across appends, not bumped every time
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_clear(self, device):
+        """`clear_` resets both lengths and storage to zero while keeping capacity."""
+        t = DynamicLengthTensor(batch_size=2, init_length=4, dim_shape=1, device=device, dtype=torch.float32)
+        t.append_(data=_make_chunk(2, 3, 1, start=1.0, device=device), lengths=torch.tensor([3, 3], device=device))
+
+        capacity_before = t._max_length
+        t.clear_()
+
+        assert t.lengths.tolist() == [0, 0]
+        assert torch.count_nonzero(t.data) == 0
+        assert t._max_length == capacity_before
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_merge(self, device):
+        """`merge_` concatenates another tensor's content along the length dim."""
+        a = DynamicLengthTensor(batch_size=2, init_length=2, dim_shape=1, device=device, dtype=torch.float32)
+        a.append_(data=_make_chunk(2, 2, 1, start=1.0, device=device), lengths=torch.tensor([2, 2], device=device))
+
+        b = DynamicLengthTensor(batch_size=2, init_length=2, dim_shape=1, device=device, dtype=torch.float32)
+        b.append_(data=_make_chunk(2, 2, 1, start=100.0, device=device), lengths=torch.tensor([1, 2], device=device))
+
+        a_item0_before = a.data[0, : a.lengths[0], 0].tolist()
+        a_item1_before = a.data[1, : a.lengths[1], 0].tolist()
+        b_item0 = b.data[0, : b.lengths[0], 0].tolist()
+        b_item1 = b.data[1, : b.lengths[1], 0].tolist()
+
+        ret = a.merge_(b)
+        assert ret is a  # in-place, returns self
+        assert a.lengths.tolist() == [3, 4]
+        assert a.data[0, : a.lengths[0], 0].tolist() == a_item0_before + b_item0
+        assert a.data[1, : a.lengths[1], 0].tolist() == a_item1_before + b_item1
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_clone_is_independent(self, device):
+        """`clone` returns a deep copy: same data/lengths, but independent storage."""
+        t = DynamicLengthTensor(batch_size=2, init_length=4, dim_shape=3, device=device, dtype=torch.float32)
+        t.append_(data=_make_chunk(2, 2, 3, start=1.0, device=device), lengths=torch.tensor([2, 1], device=device))
+
+        clone = t.clone()
+        assert clone is not t
+        assert clone.dim_shape == t.dim_shape
+        assert clone.data.shape == t.data.shape
+        assert torch.equal(clone.lengths, t.lengths)
+        assert torch.equal(clone.data, t.data)
+
+        # Mutating the clone must not affect the original.
+        clone.append_(
+            data=_make_chunk(2, 1, 3, start=50.0, device=device), lengths=torch.tensor([1, 1], device=device)
+        )
+        assert clone.lengths.tolist() == [3, 2]
+        assert t.lengths.tolist() == [2, 1]
+
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA to verify cross-device move")
+    def test_to_device(self):
+        """`to_device` moves the underlying storage (not just the bookkeeping attr)."""
+        t = DynamicLengthTensor(
+            batch_size=2, init_length=4, dim_shape=1, device=torch.device("cpu"), dtype=torch.float32
+        )
+        t.append_(data=_make_chunk(2, 2, 1, start=1.0, device=torch.device("cpu")), lengths=torch.tensor([2, 2]))
+
+        ret = t.to_device("cuda:0")
+        assert ret is t
+        assert t.device == "cuda:0"
+        assert t.data.device.type == "cuda"
+        assert t.lengths.device.type == "cuda"
+        # Content survives the move.
+        assert t.data[0, :2, 0].tolist() == [1.0, 2.0]

tests/collections/asr/decoding/test_streaming_decoding.py

@@ -0,0 +1,608 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from typing import Optional
+
+import pytest
+import torch
+import torch.nn.functional as F
+from omegaconf import open_dict
+from tqdm.auto import tqdm
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.context_biasing.biasing_multi_model import BiasingRequestItemConfig
+from nemo.collections.asr.parts.context_biasing.boosting_graph_batched import BoostingTreeModelConfig
+from nemo.collections.asr.parts.submodules.transducer_decoding.label_looping_base import (
+    BatchedLabelLoopingState,
+    GreedyBatchedLabelLoopingComputerBase,
+)
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+from nemo.collections.asr.parts.utils.rnnt_utils import BatchedHyps, Hypothesis, batched_hyps_to_hypotheses
+from tests.collections.asr.decoding.utils import load_audio, make_preprocessor_deterministic
+
+
+def get_devices_for_testing(use_cpu_always: bool = False) -> list[torch.device]:
+    devices = [torch.device("cpu")] if use_cpu_always else []
+    if torch.cuda.is_available():
+        devices.append(torch.device("cuda:0"))
+
+    if torch.mps.is_available():
+        devices.append(torch.device("mps"))
+
+    if len(devices) == 0:
+        # no fast device for testing, add CPU
+        devices.append(torch.device("cpu"))
+    return devices
+
+
+DEVICES = get_devices_for_testing(use_cpu_always=False)
+
+
+def get_model_encoder_output(
+    test_audio_filenames,
+    num_samples: int,
+    model: ASRModel,
+    device: torch.device = torch.device("cpu"),
+    dtype: torch.dtype = torch.float32,
+):
+    audio_filepaths = test_audio_filenames[:num_samples]
+
+    with torch.no_grad():
+        make_preprocessor_deterministic(model)
+        model.eval()
+
+        all_inputs, all_lengths = [], []
+        for audio_file in tqdm(audio_filepaths, desc="Loading audio files"):
+            audio_tensor, _ = load_audio(audio_file)
+            all_inputs.append(audio_tensor)
+            all_lengths.append(torch.tensor(audio_tensor.shape[0], dtype=torch.int64))
+
+        input_batch = torch.nn.utils.rnn.pad_sequence(all_inputs, batch_first=True).to(device=device, dtype=dtype)
+        length_batch = torch.tensor(all_lengths, dtype=torch.int64).to(device)
+
+        encoded_outputs, encoded_length = model(input_signal=input_batch, input_signal_length=length_batch)
+
+    return encoded_outputs, encoded_length
+
+
+def get_batch_encoder_outputs_from_records(records, model, device):
+    """Helper function to get encoder outputs for a batch of manifest records"""
+    filenames = [record["audio_filepath"] for record in records]
+    local_batch_size = len(filenames)
+    encoder_output, encoder_output_len = get_model_encoder_output(
+        test_audio_filenames=filenames, model=model, num_samples=local_batch_size, device=device
+    )
+    return encoder_output, encoder_output_len
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize(
+    "device,use_cuda_graph_decoder",
+    [(device, False) for device in DEVICES] + [(device, True) for device in DEVICES if device.type == "cuda"],
+)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_label_looping_streaming_batched_state(
+    tmp_path_factory,
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    use_cuda_graph_decoder: bool,
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    max_symbols: int,
+):
+    """Test streaming decoding with batched state"""
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "greedy_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+        decoding_cfg.greedy.max_symbols = max_symbols
+    model.change_decoding_strategy(decoding_cfg)
+
+    manifest = read_manifest(an4_val_manifest_corrected)
+    transcriptions = model.transcribe(audio=str(an4_val_manifest_corrected.absolute()), batch_size=batch_size)
+    ref_transcripts = [hyp.text for hyp in transcriptions]
+
+    all_hyps = []
+    decoding_computer: GreedyBatchedLabelLoopingComputerBase = model.decoding.decoding.decoding_computer
+    with torch.no_grad(), torch.inference_mode():
+        for i in range(0, len(manifest), batch_size):
+            encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+                manifest[i : i + batch_size], model=model, device=device
+            )
+            local_batch_size = encoder_output_len.shape[0]
+            # decode encoder output by chunks, passing state between decoder invocations
+            state: Optional[BatchedLabelLoopingState] = None
+            batched_hyps: BatchedHyps | None = None
+            encoder_output = encoder_output.transpose(1, 2)
+            for t in range(0, encoder_output.shape[1], chunk_size):
+                rest_len = encoder_output_len - t
+                current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+                current_len = torch.minimum(current_len, rest_len)
+                current_len = torch.maximum(current_len, torch.zeros_like(current_len))
+                batched_hyps_chunk, state = decoding_computer(
+                    x=encoder_output[:, t : t + chunk_size],
+                    out_len=current_len,
+                    prev_batched_state=state,
+                )
+                if batched_hyps is None:
+                    batched_hyps = batched_hyps_chunk
+                else:
+                    batched_hyps.merge_(batched_hyps_chunk)
+            assert batched_hyps is not None
+            all_hyps.extend(batched_hyps_to_hypotheses(batched_hyps, batch_size=local_batch_size))
+
+    streaming_transcripts = []
+    for hyp in all_hyps:
+        streaming_transcripts.append(model.tokenizer.ids_to_text(hyp.y_sequence.tolist()))
+    assert ref_transcripts == streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize(
+    "device,use_cuda_graph_decoder",
+    [(device, False) for device in DEVICES] + [(device, True) for device in DEVICES if device.type == "cuda"],
+)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_label_looping_streaming_partial_hypotheses(
+    tmp_path_factory,
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    use_cuda_graph_decoder: bool,
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    max_symbols: int,
+):
+    """Test streaming decoding with partial hypotheses"""
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "greedy_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+        decoding_cfg.greedy.max_symbols = max_symbols
+    model.change_decoding_strategy(decoding_cfg)
+
+    manifest = read_manifest(an4_val_manifest_corrected)
+    transcriptions = model.transcribe(audio=str(an4_val_manifest_corrected.absolute()), batch_size=batch_size)
+    ref_transcripts = [hyp.text for hyp in transcriptions]
+
+    all_hyps = []
+    rnnt_infer = model.decoding.decoding
+    with torch.no_grad(), torch.inference_mode():
+        for i in range(0, len(manifest), batch_size):
+            encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+                manifest[i : i + batch_size], model=model, device=device
+            )
+            # decode encoder output by chunks, passing state between decoder invocations
+            hyps: list[Hypothesis] | None = None
+            for t in range(0, encoder_output.shape[2], chunk_size):
+                rest_len = encoder_output_len - t
+                current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+                current_len = torch.minimum(current_len, rest_len)
+                current_len = torch.maximum(current_len, torch.zeros_like(current_len))
+                (hyps,) = rnnt_infer(
+                    encoder_output=encoder_output[:, :, t : t + chunk_size],
+                    encoded_lengths=current_len,
+                    partial_hypotheses=hyps,
+                )
+            # free up memory by resetting decoding state
+            for hyp in hyps:
+                hyp.clean_decoding_state_()
+            all_hyps.extend(hyps)
+    streaming_transcripts = []
+    for hyp in all_hyps:
+        streaming_transcripts.append(model.tokenizer.ids_to_text(hyp.y_sequence.tolist()))
+    assert ref_transcripts == streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize(
+    "device,use_cuda_graph_decoder",
+    [(device, False) for device in DEVICES] + [(device, True) for device in DEVICES if device.type == "cuda"],
+)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_label_looping_continuous_streaming_batched_state(
+    tmp_path_factory,
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    use_cuda_graph_decoder: bool,
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    max_symbols: int,
+):
+    """Test streaming continuos decoding with partial hypotheses"""
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "greedy_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+        decoding_cfg.greedy.max_symbols = max_symbols
+    model.change_decoding_strategy(decoding_cfg)
+
+    manifest = read_manifest(an4_val_manifest_corrected)
+    transcriptions = model.transcribe(audio=str(an4_val_manifest_corrected.absolute()), batch_size=batch_size)
+    ref_transcripts = [hyp.text for hyp in transcriptions]
+
+    all_hyps = [None for _ in range(len(manifest))]
+    decoding_computer: GreedyBatchedLabelLoopingComputerBase = model.decoding.decoding.decoding_computer
+    assert batch_size < len(
+        manifest
+    ), "Batch size should be less than the number of records in the manifest for continuous streaming test."
+    with torch.no_grad(), torch.inference_mode():
+        # get first 2 batches
+        encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+            manifest[:batch_size], model=model, device=device
+        )
+        encoder_output_next, encoder_output_len_next = get_batch_encoder_outputs_from_records(
+            manifest[batch_size : batch_size + batch_size], model=model, device=device
+        )
+        # we always work with encoder_output, getting next utterances from encoder_output_next
+        # so we need to pad encoder_output if it is shorter than encoder_output_next
+        if encoder_output.shape[2] < encoder_output_next.shape[2]:
+            encoder_output = F.pad(encoder_output, (0, encoder_output_next.shape[2] - encoder_output.shape[2]))
+        expanded_batch_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand(-1, chunk_size)
+        next_batch_i = 0
+        next_batch_global_i = batch_size
+        next_query_utterance_i = batch_size + batch_size
+        has_next = True  # if we have anything in next batch to decode
+        hyps: list[Hypothesis | None] = [None for _ in range(batch_size)]
+        hyps_global_indices = list(range(batch_size))
+        encoder_output_t = torch.zeros_like(encoder_output_len)
+        state = None  # decoding state
+        # while there is something to decode
+        while ((rest_len := encoder_output_len - encoder_output_t) > 0).any():
+            frame_indices = encoder_output_t[:, None] + torch.arange(chunk_size, device=device)[None, :]
+            frame_indices = torch.minimum(frame_indices, encoder_output_len[:, None] - 1)
+            current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+            current_len = torch.minimum(current_len, rest_len)
+            encoder_frames = encoder_output[expanded_batch_indices, :, frame_indices]
+            batched_hyps, state = decoding_computer(
+                x=encoder_frames,
+                out_len=current_len,
+                prev_batched_state=state,
+            )
+            hyps_continuations = batched_hyps_to_hypotheses(batched_hyps, batch_size=batch_size)
+            for i, (hyp, hyp_continuation) in enumerate(zip(hyps, hyps_continuations)):
+                if hyp is None:
+                    hyps[i] = hyp_continuation
+                else:
+                    hyp.merge_(hyp_continuation)
+            encoder_output_t += current_len
+            rest_len -= current_len
+
+            decoding_computer.reset_state_by_mask(state, rest_len == 0)
+            finished_decoding_indices = torch.nonzero(rest_len == 0, as_tuple=True)[0].cpu().tolist()
+            for idx in finished_decoding_indices:
+                hyp = hyps[idx]
+                if all_hyps[hyps_global_indices[idx]] is None:
+                    all_hyps[hyps_global_indices[idx]] = hyp
+                hyps[idx] = None  # reset to None
+                if has_next:
+                    # get next utterance to decode for finished hypothesis
+                    encoder_output[idx] = encoder_output_next[next_batch_i]
+                    encoder_output_len[idx] = encoder_output_len_next[next_batch_i]
+                    hyps_global_indices[idx] = next_batch_global_i
+                    encoder_output_t[idx] = 0
+                    next_batch_i += 1
+                    next_batch_global_i += 1
+                    # if next_batch_i is out of bounds, get next batch of encoder outputs
+                    if next_batch_i >= encoder_output_len_next.shape[0]:
+                        if next_query_utterance_i < len(manifest):
+                            encoder_output_next, encoder_output_len_next = get_batch_encoder_outputs_from_records(
+                                manifest[next_query_utterance_i : next_query_utterance_i + batch_size],
+                                model=model,
+                                device=device,
+                            )
+                            # pad if needed to allow futher assignment of encoder_output_next to encoder_output
+                            if encoder_output.shape[2] < encoder_output_next.shape[2]:
+                                encoder_output = F.pad(
+                                    encoder_output, (0, encoder_output_next.shape[2] - encoder_output.shape[2])
+                                )
+                            next_batch_i = 0
+                            next_query_utterance_i += batch_size
+                        else:
+                            has_next = False
+
+    streaming_transcripts = []
+    for hyp in all_hyps:
+        streaming_transcripts.append(model.tokenizer.ids_to_text(hyp.y_sequence.tolist()))
+    assert ref_transcripts == streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize(
+    "device,use_cuda_graph_decoder",
+    [(device, False) for device in DEVICES] + [(device, True) for device in DEVICES if device.type == "cuda"],
+)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_label_looping_continuous_streaming_partial_hypotheses(
+    tmp_path_factory,
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    use_cuda_graph_decoder: bool,
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    max_symbols: int,
+):
+    """Test streaming continuos decoding with partial hypotheses"""
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.to(device=device)
+
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "greedy_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+        decoding_cfg.greedy.max_symbols = max_symbols
+    model.change_decoding_strategy(decoding_cfg)
+
+    manifest = read_manifest(an4_val_manifest_corrected)
+    transcriptions = model.transcribe(audio=str(an4_val_manifest_corrected.absolute()), batch_size=batch_size)
+    ref_transcripts = [hyp.text for hyp in transcriptions]
+
+    all_hyps = [None for _ in range(len(manifest))]
+    rnnt_infer = model.decoding.decoding
+    assert batch_size < len(
+        manifest
+    ), "Batch size should be less than the number of records in the manifest for continuous streaming test."
+    with torch.no_grad(), torch.inference_mode():
+        # get first 2 batches
+        encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+            manifest[:batch_size], model=model, device=device
+        )
+        encoder_output_next, encoder_output_len_next = get_batch_encoder_outputs_from_records(
+            manifest[batch_size : batch_size + batch_size], model=model, device=device
+        )
+        # we always work with encoder_output, getting next utterances from encoder_output_next
+        # so we need to pad encoder_output if it is shorter than encoder_output_next
+        if encoder_output.shape[2] < encoder_output_next.shape[2]:
+            encoder_output = F.pad(encoder_output, (0, encoder_output_next.shape[2] - encoder_output.shape[2]))
+        expanded_batch_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand(-1, chunk_size)
+        # NB: we assume that encoder_output_len and encoder_output_len_next
+        # have no zero elements (no empty utterances), and we do not check this condition further
+        next_batch_i = 0
+        next_batch_global_i = batch_size
+        next_query_utterance_i = batch_size + batch_size
+        has_next = True  # if we have anything in next batch to decode
+        hyps: list[Hypothesis | None] = [None for _ in range(batch_size)]
+        hyps_global_indices = list(range(batch_size))
+        encoder_output_t = torch.zeros_like(encoder_output_len)
+        # while there is something to decode
+        while ((rest_len := encoder_output_len - encoder_output_t) > 0).any():
+            frame_indices = encoder_output_t[:, None] + torch.arange(chunk_size, device=device)[None, :]
+            frame_indices = torch.minimum(frame_indices, encoder_output_len[:, None] - 1)
+            current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+            current_len = torch.minimum(current_len, rest_len)
+            encoder_frames = encoder_output[expanded_batch_indices, :, frame_indices].transpose(1, 2)
+            (hyps,) = rnnt_infer(
+                encoder_output=encoder_frames,
+                encoded_lengths=current_len,
+                partial_hypotheses=hyps,
+            )
+            encoder_output_t += current_len
+            rest_len -= current_len
+            finished_decoding_indices = torch.nonzero(rest_len == 0, as_tuple=True)[0].cpu().tolist()
+            for idx in finished_decoding_indices:
+                hyp = hyps[idx]
+                all_hyps[hyps_global_indices[idx]] = hyp
+                # NB: we clean decoding state and set hyp to None only if we have next utterances to decode
+                # otherwise for each decoder invocation with 0 length it will recreate the hypothesis object,
+                # which is computationally expensive
+                # decoding current hyp with 0 length will not change the hypothesis
+                if has_next:
+                    hyp.clean_decoding_state_()
+                    hyps[idx] = None  # reset to None
+                    # get next utterance to decode for finished hypothesis
+                    encoder_output[idx] = encoder_output_next[next_batch_i]
+                    encoder_output_len[idx] = encoder_output_len_next[next_batch_i]
+                    hyps_global_indices[idx] = next_batch_global_i
+                    encoder_output_t[idx] = 0
+                    next_batch_i += 1
+                    next_batch_global_i += 1
+                    # if next_batch_i is out of bounds, get next batch of encoder outputs
+                    if next_batch_i >= encoder_output_len_next.shape[0]:
+                        if next_query_utterance_i < len(manifest):
+                            encoder_output_next, encoder_output_len_next = get_batch_encoder_outputs_from_records(
+                                manifest[next_query_utterance_i : next_query_utterance_i + batch_size],
+                                model=model,
+                                device=device,
+                            )
+                            # pad if needed to allow futher assignment of encoder_output_next to encoder_output
+                            if encoder_output.shape[2] < encoder_output_next.shape[2]:
+                                encoder_output = F.pad(
+                                    encoder_output, (0, encoder_output_next.shape[2] - encoder_output.shape[2])
+                                )
+                            next_batch_i = 0
+                            next_query_utterance_i += batch_size
+                        else:
+                            has_next = False
+    for hyp in hyps:
+        if hyp is not None:
+            hyp.clean_decoding_state_()
+
+    streaming_transcripts = []
+    for hyp in all_hyps:
+        streaming_transcripts.append(model.tokenizer.ids_to_text(hyp.y_sequence.tolist()))
+    assert ref_transcripts == streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize(
+    "device,use_cuda_graph_decoder",
+    [(device, False) for device in DEVICES] + [(device, True) for device in DEVICES if device.type == "cuda"],
+)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1])  # Small chunk size to trigger more state updates
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_label_looping_streaming_boosting_with_ref_transcripts(
+    tmp_path_factory,
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    use_cuda_graph_decoder: bool,
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    max_symbols: int,
+):
+    """
+    Metamorphic test: boosting with reference transcripts should yield identical results.
+
+    This test validates that when we boost with the exact transcripts that the model
+    would produce without boosting, the results remain the same. This is a metamorphic
+    property that should hold for correct implementations.
+
+    This test specifically validates the fix for TDT streaming boosting where
+    fusion states were incorrectly updated using `active_mask` instead of `found_labels_mask`.
+    """
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    # First, get reference transcripts without boosting
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "greedy_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+        decoding_cfg.greedy.max_symbols = max_symbols
+    model.change_decoding_strategy(decoding_cfg)
+
+    manifest = read_manifest(an4_val_manifest_corrected)
+    transcriptions = model.transcribe(audio=str(an4_val_manifest_corrected.absolute()), batch_size=batch_size)
+    ref_transcripts = [hyp.text for hyp in transcriptions]
+
+    # Now set up per-stream boosting with reference transcripts
+    decoding_cfg_boosted = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg_boosted.strategy = "greedy_batch"
+    with open_dict(decoding_cfg_boosted):
+        decoding_cfg_boosted.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+        decoding_cfg_boosted.greedy.max_symbols = max_symbols
+        decoding_cfg_boosted.greedy.enable_per_stream_biasing = True
+    model.change_decoding_strategy(decoding_cfg_boosted)
+
+    all_hyps = []
+    decoding_computer: GreedyBatchedLabelLoopingComputerBase = model.decoding.decoding.decoding_computer
+
+    with torch.no_grad(), torch.inference_mode():
+        for i in range(0, len(manifest), batch_size):
+            batch_records = manifest[i : i + batch_size]
+            batch_ref_transcripts = ref_transcripts[i : i + batch_size]
+
+            encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+                batch_records, model=model, device=device
+            )
+            local_batch_size = encoder_output_len.shape[0]
+
+            # Create biasing requests for each sample in the batch
+            biasing_requests = []
+            multi_biasing_ids = torch.full([local_batch_size], fill_value=-1, dtype=torch.long, device=device)
+
+            for batch_idx, ref_text in enumerate(batch_ref_transcripts):
+                if ref_text:  # Only boost non-empty transcripts
+                    request = BiasingRequestItemConfig(
+                        boosting_model_cfg=BoostingTreeModelConfig(key_phrases_list=[ref_text], unk_score=-100),
+                        boosting_model_alpha=10.0,
+                    )
+                    request.add_to_multi_model(
+                        tokenizer=model.tokenizer,
+                        biasing_multi_model=decoding_computer.biasing_multi_model,
+                    )
+                    if request.multi_model_id is not None:
+                        multi_biasing_ids[batch_idx] = request.multi_model_id
+                    biasing_requests.append(request)
+                else:
+                    biasing_requests.append(None)
+
+            # Decode encoder output by chunks, passing state between decoder invocations
+            state: Optional[BatchedLabelLoopingState] = None
+            batched_hyps: BatchedHyps | None = None
+            encoder_output = encoder_output.transpose(1, 2)
+
+            for t in range(0, encoder_output.shape[1], chunk_size):
+                rest_len = encoder_output_len - t
+                current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+                current_len = torch.minimum(current_len, rest_len)
+                current_len = torch.maximum(current_len, torch.zeros_like(current_len))
+
+                batched_hyps_chunk, state = decoding_computer(
+                    x=encoder_output[:, t : t + chunk_size],
+                    out_len=current_len,
+                    prev_batched_state=state,
+                    multi_biasing_ids=multi_biasing_ids,
+                )
+
+                if batched_hyps is None:
+                    batched_hyps = batched_hyps_chunk
+                else:
+                    batched_hyps.merge_(batched_hyps_chunk)
+
+            # Clean up biasing models
+            for request in biasing_requests:
+                if request is not None and request.multi_model_id is not None:
+                    decoding_computer.biasing_multi_model.remove_model(request.multi_model_id)
+                    request.multi_model_id = None
+
+            assert batched_hyps is not None
+            all_hyps.extend(batched_hyps_to_hypotheses(batched_hyps, batch_size=local_batch_size))
+
+    streaming_transcripts = []
+    for hyp in all_hyps:
+        streaming_transcripts.append(model.tokenizer.ids_to_text(hyp.y_sequence.tolist()))
+
+    # The key assertion: boosting with ref transcripts should yield same results
+    assert ref_transcripts == streaming_transcripts, (
+        f"Boosting with reference transcripts should yield identical results.\n"
+        f"This failure indicates a bug in fusion state handling during streaming decoding.\n"
+        f"Differences found:\n"
+        + "\n".join(
+            f"  [{i}] ref: {ref!r} != boosted: {boosted!r}"
+            for i, (ref, boosted) in enumerate(zip(ref_transcripts, streaming_transcripts))
+            if ref != boosted
+        )
+    )

tests/collections/asr/decoding/test_timestamps.py

@@ -0,0 +1,317 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os.path
+import re
+from functools import cached_property
+from typing import Any
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.collections.asr.parts.utils.timestamp_utils import get_segment_offsets, get_words_offsets
+
+
+class BaseTimestampsTest:
+    """
+    Base class for testing timestamps in decoders (CTC and RNNT).
+    This class defines common test methods that can be inherited by both
+    test_ctc_decoding.py and test_rnnt_decoding.py.
+    """
+
+    @cached_property
+    def bpe_tokenizer(self):
+        model_path = "/home/TestData/asr/stt_en_conformer_transducer_small.nemo"
+        if os.path.exists(model_path):
+            model = ASRModel.restore_from(model_path, map_location="cpu")
+        else:
+            model = ASRModel.from_pretrained("stt_en_conformer_transducer_small", map_location="cpu")
+        return model.tokenizer
+
+    @property
+    def char_offsets_chars(self):
+        char_offsets = [
+            {"char": "e", "start_offset": 0, "end_offset": 1},
+            {"char": " ", "start_offset": 2, "end_offset": 2},
+            {"char": "e", "start_offset": 3, "end_offset": 4},
+            {"char": " ", "start_offset": 5, "end_offset": 5},
+            {"char": ".", "start_offset": 6, "end_offset": 7},
+            {"char": " ", "start_offset": 8, "end_offset": 9},
+            {"char": "e", "start_offset": 10, "end_offset": 11},
+            {"char": " ", "start_offset": 12, "end_offset": 13},
+            {"char": "?", "start_offset": 14, "end_offset": 15},
+            {"char": " ", "start_offset": 16, "end_offset": 17},
+        ]
+        return char_offsets
+
+    @property
+    def word_offsets_chars_expected_output(self):
+        return [
+            {'word': 'e', 'start_offset': 0, 'end_offset': 1},
+            {'word': 'e.', 'start_offset': 3, 'end_offset': 7},
+            {'word': 'e?', 'start_offset': 10, 'end_offset': 15},
+        ]
+
+    @property
+    def word_offsets_chars_expected_output_other_delimiter(self):
+        return [
+            {'word': 'e e ', 'start_offset': 0, 'end_offset': 5},
+            {'word': ' e? ', 'start_offset': 8, 'end_offset': 17},
+        ]
+
+    @property
+    def segment_offsets_expected_output(self):
+        return [
+            {'segment': 'e e.', 'start_offset': 0, 'end_offset': 7},
+            {'segment': 'e?', 'start_offset': 10, 'end_offset': 15},
+        ]
+
+    @property
+    def segment_offsets_expected_output_gap(self):
+        return [
+            {'segment': 'e e. e?', 'start_offset': 0, 'end_offset': 15},
+        ]
+
+    @property
+    def char_offsets_wpe(self):
+        char_offsets = [
+            {"char": "nineteen", "start_offset": 0, "end_offset": 1},
+            {"char": "##th", "start_offset": 2, "end_offset": 2},
+            {"char": "re", "start_offset": 3, "end_offset": 4},
+            {"char": "seven", "start_offset": 5, "end_offset": 6},
+            {"char": "##ty", "start_offset": 6, "end_offset": 7},
+            {"char": "eighty", "start_offset": 8, "end_offset": 9},
+        ]
+
+        return char_offsets
+
+    @property
+    def word_offsets_wpe_expected_output(self):
+        return [
+            {'word': 'nineteenth', 'start_offset': 0, 'end_offset': 2},
+            {'word': 're', 'start_offset': 3, 'end_offset': 4},
+            {'word': 'seventy', 'start_offset': 5, 'end_offset': 7},
+            {'word': 'eighty', 'start_offset': 8, 'end_offset': 9},
+        ]
+
+    @property
+    def word_offsets_wpe_expected_output_other_delimiter(self):
+        return [
+            {'word': 'nineteenth', 'start_offset': 0, 'end_offset': 2},
+            {'word': 'seventy eighty', 'start_offset': 5, 'end_offset': 9},
+        ]
+
+    @property
+    def char_offsets_bpe(self):
+        char_offsets = [
+            {"char": "discuss", "start_offset": 0, "end_offset": 2},
+            {"char": "absolute", "start_offset": 2, "end_offset": 4},
+            {"char": "'", "start_offset": 5, "end_offset": 5},
+            {"char": "really", "start_offset": 5, "end_offset": 6},
+            {"char": "friend", "start_offset": 6, "end_offset": 7},
+            {"char": "ship", "start_offset": 8, "end_offset": 9},
+        ]
+
+        return char_offsets
+
+    @property
+    def encoded_char_offsets_bpe(self):
+        char_offsets = [
+            {"char": "▁discuss", "start_offset": 0, "end_offset": 2},
+            {"char": "▁absolute", "start_offset": 2, "end_offset": 4},
+            {"char": "'", "start_offset": 5, "end_offset": 5},
+            {"char": "▁really", "start_offset": 5, "end_offset": 6},
+            {"char": "▁friend", "start_offset": 6, "end_offset": 7},
+            {"char": "ship", "start_offset": 8, "end_offset": 9},
+        ]
+
+        return char_offsets
+
+    @property
+    def word_offsets_bpe_expected_output(self):
+        return [
+            {'word': "discuss", 'start_offset': 0, 'end_offset': 2},
+            {'word': "absolute'", 'start_offset': 2, 'end_offset': 5},
+            {'word': "really", 'start_offset': 5, 'end_offset': 6},
+            {'word': "friendship", 'start_offset': 6, 'end_offset': 9},
+        ]
+
+    @property
+    def word_offsets_bpe_expected_output_other_delimiter(self):
+        return [
+            {'word': "discuss absolute'", 'start_offset': 0, 'end_offset': 5},
+            {'word': "friendship", 'start_offset': 6, 'end_offset': 9},
+        ]
+
+    @staticmethod
+    def check_char_timestamps(hyp: Hypothesis, decoding: Any):
+        """Test character-level timestamps for both CTC and RNNT"""
+        assert hyp.timestamp is not None
+        assert isinstance(hyp.timestamp, dict)
+        assert 'timestep' in hyp.timestamp
+        assert 'char' in hyp.timestamp
+        assert 'word' in hyp.timestamp
+        assert 'segment' in hyp.timestamp
+
+        hypothesis_text = re.sub(r'\s+', ' ', hyp.text.strip())
+
+        words = hyp.text.split(decoding.word_seperator)
+        words = list(filter(lambda x: x != '', words))
+        assert len(hyp.timestamp['word']) == len(words)
+
+        words_from_timestamps = [ts['word'] for ts in hyp.timestamp['word']]
+        assert hypothesis_text == decoding.word_seperator.join(words_from_timestamps)
+
+        segments = []
+        segment = []
+
+        for word in words:
+            segment.append(word)
+            if word[-1] in decoding.segment_seperators:
+                segments.append(' '.join(segment))
+                segment = []
+
+        if segment:
+            segments.append(' '.join(segment))
+
+        assert len(hyp.timestamp['segment']) == len(segments)
+
+        segments_from_timestamps = [ts['segment'] for ts in hyp.timestamp['segment']]
+        assert hypothesis_text == decoding.word_seperator.join(segments_from_timestamps)
+
+    @staticmethod
+    def check_subword_timestamps(hyp: Hypothesis, decoding: Any):
+        """Test subword-level timestamps for both CTC and RNNT"""
+        assert hyp.timestamp is not None
+        assert isinstance(hyp.timestamp, dict)
+        assert 'timestep' in hyp.timestamp
+        assert 'char' in hyp.timestamp
+        assert 'word' in hyp.timestamp
+        assert 'segment' in hyp.timestamp
+
+        chars = list(hyp.text)
+        chars = list(filter(lambda x: x not in ['', ' ', '#'], chars))
+        all_chars = [list(decoding.tokenizer.tokens_to_text(data['char'])) for data in hyp.timestamp['char']]
+        all_chars = [char for subword in all_chars for char in subword]
+        all_chars = list(filter(lambda x: x not in ['', ' ', '#'], all_chars))
+        assert len(chars) == len(all_chars)
+
+        hypothesis_text = re.sub(r'\s+', ' ', hyp.text.strip())
+
+        words_from_timestamps = [ts['word'] for ts in hyp.timestamp['word']]
+        assert hypothesis_text == decoding.word_seperator.join(words_from_timestamps)
+
+        segments_count = sum([hyp.text.count(seperator) for seperator in decoding.segment_seperators])
+        if hyp.text[-1] not in decoding.segment_seperators:
+            segments_count += 1
+
+        if hyp.text in decoding.segment_seperators:
+            segments_count = 0
+
+        assert len(hyp.timestamp['segment']) == segments_count
+
+        segments_from_timestamps = [ts['segment'] for ts in hyp.timestamp['segment']]
+        assert hypothesis_text == decoding.word_seperator.join(segments_from_timestamps)
+
+    def test_word_offsets_chars(self):
+        word_offsets = get_words_offsets(
+            char_offsets=self.char_offsets_chars,
+            encoded_char_offsets=None,
+            word_delimiter_char=" ",
+            tokenizer_type="char",
+            supported_punctuation={'.', '!', '?'},
+            decode_tokens_to_str=self.decoding_char.decode_tokens_to_str,
+        )
+
+        assert word_offsets == self.word_offsets_chars_expected_output
+
+    def test_word_offsets_char_other_delimiter(self):
+        word_offsets = get_words_offsets(
+            char_offsets=self.char_offsets_chars,
+            encoded_char_offsets=None,
+            tokenizer_type="char",
+            word_delimiter_char=".",
+            supported_punctuation={'.', '!', '?'},
+            decode_tokens_to_str=self.decoding_char.decode_tokens_to_str,
+        )
+
+        assert word_offsets == self.word_offsets_chars_expected_output_other_delimiter
+
+    def test_word_offsets_subword_wpe(self):
+
+        word_offsets = get_words_offsets(
+            char_offsets=self.char_offsets_wpe,
+            encoded_char_offsets=None,
+            word_delimiter_char=" ",
+            tokenizer_type="wpe",
+            supported_punctuation={'.', '!', '?'},
+            decode_tokens_to_str=self.decoding_subword_wpe.decode_tokens_to_str,
+        )
+
+        assert word_offsets == self.word_offsets_wpe_expected_output
+
+    def test_word_offsets_subword_wpe_other_delimiter(self):
+
+        word_offsets = get_words_offsets(
+            char_offsets=self.char_offsets_wpe,
+            encoded_char_offsets=None,
+            word_delimiter_char="re",
+            tokenizer_type="wpe",
+            supported_punctuation={'.', '!', '?'},
+            decode_tokens_to_str=self.decoding_subword_wpe.decode_tokens_to_str,
+        )
+
+        assert word_offsets == self.word_offsets_wpe_expected_output_other_delimiter
+
+    def test_word_offsets_subword_bpe(self):
+
+        word_offsets = get_words_offsets(
+            char_offsets=self.char_offsets_bpe,
+            encoded_char_offsets=self.encoded_char_offsets_bpe,
+            word_delimiter_char=" ",
+            tokenizer_type="bpe",
+            supported_punctuation={'.', '!', '?'},
+            decode_tokens_to_str=self.decoding_subword_bpe.decode_tokens_to_str,
+        )
+
+        assert word_offsets == self.word_offsets_bpe_expected_output
+
+    def test_word_offsets_subword_bpe_other_delimiter(self):
+        word_offsets = get_words_offsets(
+            char_offsets=self.char_offsets_bpe,
+            encoded_char_offsets=self.encoded_char_offsets_bpe,
+            word_delimiter_char="really",
+            tokenizer_type="bpe",
+            supported_punctuation={'.', '!', '?'},
+            decode_tokens_to_str=self.decoding_subword_bpe.decode_tokens_to_str,
+        )
+
+        assert word_offsets == self.word_offsets_bpe_expected_output_other_delimiter
+
+    def test_segment_offsets_delimiter(self):
+        segment_offsets = get_segment_offsets(
+            word_offsets=self.word_offsets_chars_expected_output,
+            segment_delimiter_tokens=['.', '!', '?'],
+            supported_punctuation={'.', '!', '?'},
+        )
+
+        assert segment_offsets == self.segment_offsets_expected_output
+
+    def test_segment_offsets_gap(self):
+        segment_offsets = get_segment_offsets(
+            word_offsets=self.word_offsets_chars_expected_output,
+            segment_delimiter_tokens=[],
+            supported_punctuation={},
+            segment_gap_threshold=10,
+        )
+
+        assert segment_offsets == self.segment_offsets_expected_output_gap

tests/collections/asr/decoding/utils.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from contextlib import contextmanager
+
+import librosa
+import torch
+
+from nemo.collections.asr.models import ASRModel
+
+
+@contextmanager
+def preserve_decoding_cfg_and_cpu_device(model: ASRModel):
+    """
+    Context manager to preserve the decoding strategy and device of the model.
+    This is useful for tests that modify the model's decoding strategy or device
+    to avoid side effects or costly model reloading.
+    """
+    backup_decoding_cfg = copy.deepcopy(model.cfg.decoding)
+
+    try:
+        yield
+    finally:
+        model.to(device="cpu")
+        if model.cfg.decoding != backup_decoding_cfg:
+            model.change_decoding_strategy(backup_decoding_cfg)
+
+
+def load_audio(file_path, target_sr=16000) -> tuple[torch.Tensor, int]:
+    audio, sr = librosa.load(file_path, sr=target_sr)
+    return torch.tensor(audio, dtype=torch.float32), sr
+
+
+@contextmanager
+def avoid_sync_operations(device: torch.device):
+    try:
+        if device.type == "cuda":
+            torch.cuda.set_sync_debug_mode(2)  # fail if a blocking operation
+        yield
+    finally:
+        if device.type == "cuda":
+            torch.cuda.set_sync_debug_mode(0)  # default, blocking operations are allowed
+
+
+def make_preprocessor_deterministic(model: ASRModel):
+    model.preprocessor.featurizer.dither = 0.0
+    model.preprocessor.featurizer.pad_to = 0

tests/collections/asr/inference/test_audio_bufferer.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytest
+import torch
+
+from nemo.collections.asr.inference.streaming.buffering.audio_bufferer import AudioBufferer, BatchedAudioBufferer
+from nemo.collections.asr.inference.streaming.framing.mono_stream import MonoStream
+from nemo.collections.asr.inference.streaming.framing.multi_stream import MultiStream
+
+
+@pytest.fixture(scope="module")
+def test_audios():
+    return torch.ones(83200), torch.ones(118960)
+
+
+class TestAudioBufferer:
+
+    @pytest.mark.unit
+    def test_audio_bufferer(self, test_audios):
+        for audio in test_audios:
+            stream = MonoStream(16000, frame_size_in_secs=2.5, stream_id=0, pad_last_frame=False)
+            stream.load_audio(audio, options=None)
+
+            frame_bufferer = AudioBufferer(16000, buffer_size_in_secs=5.0)
+
+            for frame in iter(stream):
+                frame = frame[0]
+                frame_bufferer.update(frame)
+                buffer = frame_bufferer.get_buffer()
+
+                assert len(buffer) == frame_bufferer.buffer_size
+                assert torch.allclose(buffer[-frame.size :], frame.samples, atol=1e-5)
+
+
+class TestBatchedAudioBufferer:
+
+    @pytest.mark.unit
+    def test_batched_audio_bufferer(self, test_audios):
+
+        multi_stream = MultiStream(n_frames_per_stream=1)
+        for stream_id, audio in enumerate(test_audios):
+            stream = MonoStream(16000, 2.5, stream_id=stream_id, pad_last_frame=False)
+            stream.load_audio(audio, options=None)
+            multi_stream.add_stream(stream, stream_id=stream_id)
+
+        batched_audio_bufferer = BatchedAudioBufferer(16000, buffer_size_in_secs=5.0)
+
+        for frames in iter(multi_stream):
+            buffered_frames, left_paddings = batched_audio_bufferer.update(frames)
+            for idx, frame in enumerate(frames):
+                frame_buffer = buffered_frames[idx]
+                assert torch.allclose(frame_buffer[-frame.size :], frame.samples, atol=1e-5)

tests/collections/asr/inference/test_bpe_decoder.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.collections.asr.inference.model_wrappers.ctc_inference_wrapper import CTCInferenceWrapper
+from nemo.collections.asr.inference.utils.bpe_decoder import BPEDecoder
+from nemo.collections.asr.inference.utils.text_segment import TextSegment, Word
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig
+
+
+@pytest.fixture(scope="module")
+def bpe_decoder():
+    asr_model = CTCInferenceWrapper(
+        model_name="stt_en_conformer_ctc_small",
+        decoding_cfg=CTCDecodingConfig(),
+        device="cuda" if torch.cuda.is_available() else "cpu",
+    )
+    return BPEDecoder(
+        vocabulary=asr_model.get_vocabulary(),
+        tokenizer=asr_model.tokenizer,
+        confidence_aggregator=min,
+        asr_supported_puncts=asr_model.supported_punctuation(),
+        word_boundary_tolerance=0.0,  # Set to 0.0 for easy testing
+        token_duration_in_secs=asr_model.get_model_stride(in_secs=True),
+    )
+
+
+class TestBPEDecoder:
+
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "text",
+        [
+            "the quick brown fox jumps over the lazy dog",
+            "lorem ipsum dolor sit amet",
+            "this a test sentence",
+        ],
+    )
+    def test_group_tokens_into_words(self, bpe_decoder, text):
+        ground_truth_words = text.split()
+        tokens = bpe_decoder.tokenizer.text_to_ids(text)
+        n_tokens = len(tokens)
+        timestamps = [float(i) for i in range(n_tokens)]
+        confidences = [0.1] * n_tokens
+
+        words, need_merge = bpe_decoder.group_tokens_into_words(tokens, timestamps, confidences)
+        assert len(words) == len(ground_truth_words)
+        prev_word_end = -1
+        for word, ground_truth_word in zip(words, ground_truth_words):
+            assert isinstance(word, Word)
+            assert word.text == ground_truth_word
+            assert word.conf == 0.1
+            assert word.end > word.start and word.start >= prev_word_end
+            prev_word_end = word.end
+        assert need_merge == False
+
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "text",
+        [
+            "the quick brown fox jumps over the lazy dog",
+            "lorem ipsum dolor sit amet",
+            "this a test sentence",
+        ],
+    )
+    def test_group_tokens_into_segment(self, bpe_decoder, text):
+        tokens = bpe_decoder.tokenizer.text_to_ids(text)
+        n_tokens = len(tokens)
+        timestamps = [float(i) for i in range(n_tokens)]
+        confidences = [0.1] * n_tokens
+
+        segment, need_merge = bpe_decoder.group_tokens_into_segment(tokens, timestamps, confidences)
+        assert isinstance(segment, TextSegment)
+        assert need_merge == False
+        assert segment.text == text
+        assert segment.start == 0.0
+        assert segment.end == (n_tokens - 1) * bpe_decoder.token_duration_in_secs
+        assert segment.conf == 0.1

tests/collections/asr/inference/test_enums.py

@@ -0,0 +1,61 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from nemo.collections.asr.inference.utils.enums import (
+    ASRDecodingType,
+    ASROutputGranularity,
+    FeatureBufferPaddingMode,
+    PipelineType,
+    RequestType,
+)
+
+
+class TestEnums:
+
+    @pytest.mark.unit
+    def test_ASRDecodingType(self):
+        assert ASRDecodingType.from_str("ctc") == ASRDecodingType.CTC
+        assert ASRDecodingType.from_str("RNNT") == ASRDecodingType.RNNT
+        with pytest.raises(ValueError):
+            ASRDecodingType.from_str("invalid")
+
+    @pytest.mark.unit
+    def test_ASROutputGranularity(self):
+        assert ASROutputGranularity.from_str("word") == ASROutputGranularity.WORD
+        assert ASROutputGranularity.from_str("segment") == ASROutputGranularity.SEGMENT
+        with pytest.raises(ValueError):
+            ASROutputGranularity.from_str("invalid")
+
+    @pytest.mark.unit
+    def test_PipelineType(self):
+        assert PipelineType.from_str("buffered") == PipelineType.BUFFERED
+        assert PipelineType.from_str("cache_aware") == PipelineType.CACHE_AWARE
+        with pytest.raises(ValueError):
+            PipelineType.from_str("invalid")
+
+    @pytest.mark.unit
+    def test_RequestType(self):
+        assert RequestType.from_str("frame") == RequestType.FRAME
+        assert RequestType.from_str("feature_buffer") == RequestType.FEATURE_BUFFER
+        with pytest.raises(ValueError):
+            RequestType.from_str("invalid")
+
+    @pytest.mark.unit
+    def test_FeatureBufferPaddingMode(self):
+        assert FeatureBufferPaddingMode.from_str("left") == FeatureBufferPaddingMode.LEFT
+        assert FeatureBufferPaddingMode.from_str("right") == FeatureBufferPaddingMode.RIGHT
+        with pytest.raises(ValueError):
+            FeatureBufferPaddingMode.from_str("invalid")

tests/collections/asr/inference/test_framing.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.collections.asr.inference.streaming.framing.mono_stream import MonoStream
+from nemo.collections.asr.inference.streaming.framing.multi_stream import MultiStream
+
+
+@pytest.fixture(scope="module")
+def test_audios():
+    return torch.ones(83200), torch.ones(118960)
+
+
+class TestMonoWavStream:
+
+    @pytest.mark.unit
+    def test_mono_wav_stream_no_pad(self, test_audios):
+        for audio in test_audios:
+            stream = MonoStream(16000, 2.5, stream_id=0, pad_last_frame=False)
+            stream.load_audio(audio, options=None)
+            audio_len_in_samples = stream.samples.shape[0]
+            i = 0
+            total_samples = 0
+            for frame in iter(stream):
+                total_samples += len(frame[0].samples)
+                i += 1
+            assert total_samples == audio_len_in_samples
+            assert frame[0].is_last == True
+
+    @pytest.mark.unit
+    def test_mono_wav_stream_with_pad(self, test_audios):
+        for audio in test_audios:
+            stream = MonoStream(16000, 2.5, stream_id=0, pad_last_frame=True)
+            stream.load_audio(audio, options=None)
+            for frame in iter(stream):
+                last_frame_size = frame[0].size
+            assert last_frame_size == stream.frame_size
+
+
+class TestMultiStream:
+
+    @pytest.mark.unit
+    def test_multi_stream(self, test_audios):
+        multi_stream = MultiStream(n_frames_per_stream=1)
+        audio_len_in_samples = {}
+        for stream_id, audio in enumerate(test_audios):
+            stream = MonoStream(16000, 2.5, stream_id=stream_id, pad_last_frame=False)
+            stream.load_audio(audio, options=None)
+            multi_stream.add_stream(stream, stream_id=stream_id)
+            audio_len_in_samples[stream_id] = stream.samples.shape[0]
+
+        total_samples = {}
+        for frames in iter(multi_stream):
+            for frame in frames:
+                total_samples[frame.stream_id] = total_samples.get(frame.stream_id, 0) + frame.size
+
+        assert total_samples == audio_len_in_samples

tests/collections/asr/inference/test_greedy_decoder.py

@@ -0,0 +1,210 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.collections.asr.inference.streaming.decoders.greedy.greedy_ctc_decoder import CTCGreedyDecoder
+from nemo.collections.asr.inference.streaming.decoders.greedy.greedy_rnnt_decoder import (
+    ClippedRNNTGreedyDecoder,
+    RNNTGreedyDecoder,
+)
+
+
+class TestCTCGreedyDecoder:
+
+    @pytest.mark.unit
+    def test_ctc_greedy_decoder(self):
+
+        vocab = ["a", "b", "c", "d"]
+        decoder = CTCGreedyDecoder(vocabulary=vocab)
+
+        assert decoder.blank_id == len(vocab)
+        assert decoder.is_token_silent(len(vocab)) == True
+
+        for i in range(len(vocab)):
+            assert decoder.is_token_silent(i) == False
+
+        for i in range(len(vocab)):
+            assert decoder.is_token_start_of_word(i) == False
+
+        assert decoder.count_silent_tokens([0, 1, 2, 3, 4], 0, 5) == 1
+        assert decoder.count_silent_tokens([0, 1, 2, 3, 4], 0, 3) == 0
+        assert decoder.first_non_silent_token([1, 2, 3, 4], 0, 5) == 0
+
+        log_probs = torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.05], [0.4, 0.3, 0.2, 0.1, 0.05]])
+        assert decoder.get_labels(log_probs) == log_probs.argmax(dim=-1).tolist()
+
+    @pytest.mark.unit
+    def test_ctc_greedy_decoder_with_previous_token(self):
+        vocab = ["a", "b", "c", "d"]
+        decoder = CTCGreedyDecoder(vocabulary=vocab)
+
+        log_probs = torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.05], [0.1, 0.2, 0.3, 0.4, 0.05], [0.4, 0.3, 0.2, 0.1, 0.05]])
+        last_token_id = 3
+        output = decoder(log_probs, compute_confidence=False, previous=last_token_id)
+        assert output["tokens"] == [0]
+        assert output["timesteps"] == [2]
+
+        output = decoder(log_probs, compute_confidence=False, previous=None)
+        assert output["tokens"] == [3, 0]
+        assert output["timesteps"] == [0, 2]
+
+
+class TestRNNTGreedyDecoder:
+
+    @pytest.mark.unit
+    def test_rnnt_greedy_decoder(self):
+
+        vocab = ["a", "b", "c", "d"]
+        decoder = RNNTGreedyDecoder(vocab)
+
+        blank_id = len(vocab)
+        assert decoder.blank_id == blank_id
+        assert decoder.is_token_silent(blank_id) == True
+
+        for i in range(len(vocab)):
+            assert decoder.is_token_silent(i) == False
+
+        for i in range(len(vocab)):
+            assert decoder.is_token_start_of_word(i) == False
+
+        assert decoder.count_silent_tokens([0, 1, 2, 3, 4], 0, 5) == 1
+        assert decoder.count_silent_tokens([0, 1, 2, 3, 4], 0, 3) == 0
+        assert decoder.first_non_silent_token([1, 2, 3, 4], 0, 5) == 0
+
+    @pytest.mark.unit
+    def test_call_confidence_passthrough(self):
+        """Per-token confidences are propagated to the output, aligned with the decoded tokens."""
+        decoder = RNNTGreedyDecoder(["a", "b", "c", "d"])
+
+        output, _, new_offset = decoder(
+            global_timestamps=torch.tensor([0, 1, 2, 3]),
+            tokens=torch.tensor([0, 1, 2, 3]),
+            length=5,
+            offset=0,
+            confidences=torch.tensor([0.9, 0.8, 0.7, 0.6]),
+        )
+
+        assert output["tokens"] == [0, 1, 2, 3]
+        assert output["confidences"] == pytest.approx([0.9, 0.8, 0.7, 0.6])
+        assert new_offset == 4
+
+    @pytest.mark.unit
+    def test_call_confidence_trimmed_by_offset(self):
+        """Confidences are trimmed by `offset` together with tokens/timestamps."""
+        decoder = RNNTGreedyDecoder(["a", "b", "c", "d"])
+
+        output, _, _ = decoder(
+            global_timestamps=torch.tensor([0, 1, 2, 3]),
+            tokens=torch.tensor([0, 1, 2, 3]),
+            length=5,
+            offset=2,
+            confidences=[0.9, 0.8, 0.7, 0.6],
+        )
+
+        assert output["tokens"] == [2, 3]
+        assert output["confidences"] == pytest.approx([0.7, 0.6])
+
+    @pytest.mark.unit
+    def test_call_confidence_defaults_to_zero(self):
+        """Without confidences, zeros are returned for each decoded token (backward compatible)."""
+        decoder = RNNTGreedyDecoder(["a", "b", "c", "d"])
+
+        output, _, _ = decoder(
+            global_timestamps=torch.tensor([0, 1, 2, 3]),
+            tokens=torch.tensor([0, 1, 2, 3]),
+            length=5,
+            offset=0,
+            confidences=None,
+        )
+
+        assert output["tokens"] == [0, 1, 2, 3]
+        assert output["confidences"] == [0.0, 0.0, 0.0, 0.0]
+
+
+class TestClippedRNNTGreedyDecoder:
+
+    @pytest.mark.unit
+    def test_confidence_passthrough(self):
+        """Per-token confidences are clipped with the same mask as tokens/timesteps."""
+        vocab = ["a", "b", "c", "d"]
+        decoder = ClippedRNNTGreedyDecoder(vocabulary=vocab, tokens_per_frame=4, endpointer=None)
+
+        tokens = torch.tensor([0, 1, 2, 3])
+        timesteps = torch.tensor([0, 1, 2, 3])
+        confidences = torch.tensor([0.9, 0.8, 0.7, 0.6])
+
+        clipped_output, _, is_eou, _, _ = decoder(
+            global_timesteps=timesteps,
+            tokens=tokens,
+            clip_start=0,
+            clip_end=4,
+            alignment_length=4,
+            is_last=True,
+            is_start=True,
+            confidences=confidences,
+        )
+
+        assert is_eou is True
+        assert clipped_output["tokens"] == [0, 1, 2, 3]
+        assert clipped_output["confidences"] == pytest.approx([0.9, 0.8, 0.7, 0.6])
+        assert len(clipped_output["confidences"]) == len(clipped_output["tokens"])
+
+    @pytest.mark.unit
+    def test_confidence_clipping_follows_token_mask(self):
+        """Confidences outside the clip range are dropped together with their tokens."""
+        vocab = ["a", "b", "c", "d"]
+        decoder = ClippedRNNTGreedyDecoder(vocabulary=vocab, tokens_per_frame=4, endpointer=None)
+
+        tokens = torch.tensor([0, 1, 2, 3])
+        timesteps = torch.tensor([0, 1, 2, 3])
+        confidences = torch.tensor([0.9, 0.8, 0.7, 0.6])
+
+        clipped_output, _, _, _, _ = decoder(
+            global_timesteps=timesteps,
+            tokens=tokens,
+            clip_start=1,
+            clip_end=4,
+            alignment_length=4,
+            is_last=True,
+            is_start=True,
+            confidences=confidences,
+        )
+
+        assert clipped_output["tokens"] == [1, 2, 3]
+        assert clipped_output["confidences"] == pytest.approx([0.8, 0.7, 0.6])
+
+    @pytest.mark.unit
+    def test_confidence_defaults_to_zero(self):
+        """Without confidences, zeros are returned for each clipped token (backward compatible)."""
+        vocab = ["a", "b", "c", "d"]
+        decoder = ClippedRNNTGreedyDecoder(vocabulary=vocab, tokens_per_frame=4, endpointer=None)
+
+        tokens = torch.tensor([0, 1, 2, 3])
+        timesteps = torch.tensor([0, 1, 2, 3])
+
+        clipped_output, _, _, _, _ = decoder(
+            global_timesteps=timesteps,
+            tokens=tokens,
+            clip_start=0,
+            clip_end=4,
+            alignment_length=4,
+            is_last=True,
+            is_start=True,
+            confidences=None,
+        )
+
+        assert clipped_output["tokens"] == [0, 1, 2, 3]
+        assert clipped_output["confidences"] == [0.0, 0.0, 0.0, 0.0]

tests/collections/asr/inference/test_greedy_endpointing.py

@@ -0,0 +1,231 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.collections.asr.inference.streaming.endpointing.greedy.greedy_ctc_endpointing import CTCGreedyEndpointing
+from nemo.collections.asr.inference.streaming.endpointing.greedy.greedy_rnnt_endpointing import RNNTGreedyEndpointing
+from nemo.collections.asr.inference.utils.endpointing_utils import millisecond_to_frames
+
+
+class TestGreedyEndpointing:
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "inputs, expected",
+        [
+            ((100, 80), 2),
+            ((100, 100), 1),
+            ((100, 40), 3),
+        ],
+    )
+    def test_millisecond_to_frames(self, inputs, expected):
+        assert millisecond_to_frames(*inputs) == expected
+
+    @pytest.mark.unit
+    def test_endpointing_with_negative_stop_history_eou(self):
+        for endpointing_cls in [CTCGreedyEndpointing, RNNTGreedyEndpointing]:
+            greedy_endpointing = endpointing_cls(vocabulary=["a", "b", "c"], ms_per_timestep=100, stop_history_eou=-1)
+            if isinstance(greedy_endpointing, CTCGreedyEndpointing):
+                b = len(greedy_endpointing.greedy_ctc_decoder.vocabulary)
+            else:
+                b = len(greedy_endpointing.greedy_rnnt_decoder.vocabulary)
+            emissions = [0, 1, 2, b, b, b, b, b, b, b, b, b]
+
+            # False case, because stop_history_eou = -1
+            assert greedy_endpointing.detect_eou_given_emissions(emissions, 3) == (False, -1)
+
+    @pytest.mark.unit
+    def test_endpointing_with_positive_stop_history_eou(self):
+        for endpointing_cls in [CTCGreedyEndpointing, RNNTGreedyEndpointing]:
+            greedy_endpointing = endpointing_cls(
+                vocabulary=["a", "b", "c"], ms_per_timestep=20, stop_history_eou=100, residue_tokens_at_end=0
+            )
+            if isinstance(greedy_endpointing, CTCGreedyEndpointing):
+                b = len(greedy_endpointing.greedy_ctc_decoder.vocabulary)
+            else:
+                b = len(greedy_endpointing.greedy_rnnt_decoder.vocabulary)
+            emissions = [0, 1, 2, b, b, b, b, b, b, b, b, b]
+
+            for pivot_point in range(len(emissions)):
+                eou_detected, eou_detected_at = greedy_endpointing.detect_eou_given_emissions(emissions, pivot_point)
+                assert eou_detected == True
+
+    @pytest.mark.unit
+    def test_detect_eou_given_timestamps_empty_inputs(self):
+        for endpointing_cls in [CTCGreedyEndpointing, RNNTGreedyEndpointing]:
+            greedy_endpointing = endpointing_cls(
+                vocabulary=["a", "b", "c"], ms_per_timestep=80, stop_history_eou=100, residue_tokens_at_end=0
+            )
+
+            # Test with empty timesteps and tokens
+            timesteps = torch.tensor([])
+            tokens = torch.tensor([])
+            alignment_length = 10
+
+            eou_detected, eou_detected_at = greedy_endpointing.detect_eou_given_timestamps(
+                timesteps, tokens, alignment_length
+            )
+            assert eou_detected == False
+            assert eou_detected_at == -1
+
+    @pytest.mark.unit
+    def test_detect_eou_given_timestamps_disabled_stop_history(self):
+        for endpointing_cls in [CTCGreedyEndpointing, RNNTGreedyEndpointing]:
+            greedy_endpointing = endpointing_cls(
+                vocabulary=["a", "b", "c"],
+                ms_per_timestep=80,
+                stop_history_eou=-1,  # Disabled
+                residue_tokens_at_end=0,
+            )
+
+            timesteps = torch.tensor([0, 2, 4, 6])
+            tokens = torch.tensor([0, 1, 2, 3])
+            alignment_length = 10
+
+            eou_detected, eou_detected_at = greedy_endpointing.detect_eou_given_timestamps(
+                timesteps, tokens, alignment_length
+            )
+            assert eou_detected == False
+            assert eou_detected_at == -1
+
+    @pytest.mark.unit
+    def test_detect_eou_given_timestamps_trailing_silence(self):
+        for endpointing_cls in [CTCGreedyEndpointing, RNNTGreedyEndpointing]:
+            greedy_endpointing = endpointing_cls(
+                vocabulary=["a", "b", "c"], ms_per_timestep=20, stop_history_eou=80, residue_tokens_at_end=0
+            )
+
+            # Last token at position 5, alignment_length is 10
+            # Trailing silence = 10 - 4 - 1 = 5 frames > stop_history_eou (4)
+            timesteps = torch.tensor([0, 1, 2, 3, 4])
+            tokens = torch.tensor([0, 1, 2, 3, 4])
+            alignment_length = 10
+
+            eou_detected, eou_detected_at = greedy_endpointing.detect_eou_given_timestamps(
+                timesteps, tokens, alignment_length
+            )
+            assert eou_detected == True
+            # eou_detected_at = 4 + 1 + 4//2 = 7
+            assert eou_detected_at == 7
+
+    @pytest.mark.unit
+    def test_detect_eou_given_timestamps_no_trailing_silence(self):
+        for endpointing_cls in [CTCGreedyEndpointing, RNNTGreedyEndpointing]:
+            greedy_endpointing = endpointing_cls(
+                vocabulary=["a", "b", "c"], ms_per_timestep=20, stop_history_eou=80, residue_tokens_at_end=0
+            )
+
+            # Last token at position 8, alignment_length is 10
+            # Trailing silence = 10 - 8 - 1 = 1 frame < stop_history_eou (4)
+            timesteps = torch.tensor([0, 1, 2, 3, 8])
+            tokens = torch.tensor([0, 1, 2, 3, 4])
+            alignment_length = 10
+
+            eou_detected, eou_detected_at = greedy_endpointing.detect_eou_given_timestamps(
+                timesteps, tokens, alignment_length
+            )
+            assert eou_detected == False
+            assert eou_detected_at == -1
+
+    @pytest.mark.unit
+    def test_detect_eou_given_timestamps_gap_detection(self):
+        for endpointing_cls in [CTCGreedyEndpointing, RNNTGreedyEndpointing]:
+            greedy_endpointing = endpointing_cls(
+                vocabulary=["a", "b", "c"], ms_per_timestep=20, stop_history_eou=80, residue_tokens_at_end=0
+            )
+
+            # Large gap between tokens: 8 - 2 - 1 = 5 frames > stop_history_eou (4)
+            timesteps = torch.tensor([0, 2, 8, 9])
+            tokens = torch.tensor([0, 1, 2, 3])
+            alignment_length = 10
+
+            eou_detected, eou_detected_at = greedy_endpointing.detect_eou_given_timestamps(
+                timesteps, tokens, alignment_length
+            )
+            assert eou_detected == True
+            # eou_detected_at = 2 + 1 + 4//2 = 5
+            assert eou_detected_at == 5
+
+    @pytest.mark.unit
+    def test_rnnt_vad_endpointing_disabled(self):
+        rnnt_endpointing = RNNTGreedyEndpointing(
+            vocabulary=["a", "b", "c"],
+            ms_per_timestep=100,
+            effective_buffer_size_in_secs=None,  # VAD disabled
+            stop_history_eou=100,
+        )
+
+        # Test with VAD segments - should raise ValueError since VAD is disabled
+        vad_segments = torch.tensor([[0.0, 1.0], [1.5, 2.5]])
+
+        with pytest.raises(
+            ValueError, match="Effective buffer size in seconds is required for VAD-based EOU detection"
+        ):
+            rnnt_endpointing.detect_eou_vad(vad_segments)
+
+    @pytest.mark.unit
+    def test_rnnt_vad_endpointing_enabled_no_eou(self):
+        rnnt_endpointing = RNNTGreedyEndpointing(
+            vocabulary=["a", "b", "c"],
+            ms_per_timestep=100,
+            effective_buffer_size_in_secs=2.0,  # VAD enabled
+            stop_history_eou=100,
+        )
+
+        # Test with VAD segments that don't trigger EOU
+        vad_segments = torch.tensor([[0.0, 1.45], [1.5, 2.0]])
+        eou_detected, eou_detected_at = rnnt_endpointing.detect_eou_vad(vad_segments, stop_history_eou=100)
+
+        assert eou_detected == False
+        assert eou_detected_at == -1
+
+    @pytest.mark.unit
+    def test_rnnt_vad_endpointing_enabled_with_eou(self):
+        rnnt_endpointing = RNNTGreedyEndpointing(
+            vocabulary=["a", "b", "c"],
+            ms_per_timestep=100,
+            effective_buffer_size_in_secs=2.0,  # VAD enabled
+            stop_history_eou=100,
+        )
+
+        # Test with VAD segments that should trigger EOU
+        # Create segments with enough silence to trigger EOU
+        vad_segments = torch.tensor([[0.0, 0.5], [1.0, 2.0]])  # Gap of 0.5s between segments
+        eou_detected, eou_detected_at = rnnt_endpointing.detect_eou_vad(vad_segments, stop_history_eou=100)
+
+        # This should detect EOU if the silence gap is sufficient
+        # The exact behavior depends on the VAD logic implementation
+        assert eou_detected == True
+        assert eou_detected_at == 5
+
+    @pytest.mark.unit
+    def test_rnnt_vad_endpointing_enabled_with_eou_at_end(self):
+        rnnt_endpointing = RNNTGreedyEndpointing(
+            vocabulary=["a", "b", "c"],
+            ms_per_timestep=100,
+            effective_buffer_size_in_secs=2.0,  # VAD enabled
+            stop_history_eou=100,
+        )
+
+        # Test with VAD segments that should trigger EOU
+        # Create segments with enough silence to trigger EOU
+        vad_segments = torch.tensor([[0.0, 0.5], [1.0, 1.8]])  # Gap of 0.5s between segments
+        eou_detected, eou_detected_at = rnnt_endpointing.detect_eou_vad(vad_segments, stop_history_eou=100)
+
+        # This should detect EOU if the silence gap is sufficient
+        # The exact behavior depends on the VAD logic implementation
+        assert eou_detected == True
+        assert eou_detected_at == 18

tests/collections/asr/inference/test_incremental_audio_bufferer.py

@@ -0,0 +1,181 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytest
+import torch
+
+from nemo.collections.asr.inference.streaming.buffering.incremental_audio_bufferer import IncrementalAudioBufferer
+from nemo.collections.asr.inference.streaming.framing.mono_stream import MonoStream
+from nemo.collections.asr.inference.streaming.framing.request import Frame
+
+
+@pytest.fixture(scope="module")
+def test_audios():
+    return torch.ones(83200), torch.ones(118960)
+
+
+def _make_frame(samples: torch.Tensor, stream_id: int = 0, is_first: bool = False, is_last: bool = False) -> Frame:
+    return Frame(
+        samples=samples,
+        stream_id=stream_id,
+        is_first=is_first,
+        is_last=is_last,
+    )
+
+
+class TestIncrementalAudioBufferer:
+    """Tests for IncrementalAudioBufferer."""
+
+    @pytest.mark.unit
+    def test_constructor_valid_params(self):
+        """Constructor with valid params initializes buffer and capacity."""
+        sample_rate = 16000
+        buffer_size_in_secs = 5.0
+        chunk_size_in_secs = 2.5
+        overlap_size_in_secs = 2.5
+        buf = IncrementalAudioBufferer(
+            sample_rate=sample_rate,
+            buffer_size_in_secs=buffer_size_in_secs,
+            chunk_size_in_secs=chunk_size_in_secs,
+            overlap_size_in_secs=overlap_size_in_secs,
+        )
+        assert buf.sample_rate == sample_rate
+        assert buf.buffer_size == int(buffer_size_in_secs * sample_rate)
+        assert buf.chunk_size == int(chunk_size_in_secs * sample_rate)
+        assert buf.overlap_size == int(overlap_size_in_secs * sample_rate)
+        assert buf.sample_buffer.shape[0] == buf.buffer_size
+        assert buf.remaining_capacity == buf.buffer_size
+        assert buf.head == 0
+        assert not buf.is_full()
+
+    @pytest.mark.unit
+    def test_constructor_overlap_negative_raises(self):
+        """Overlap < 0 raises ValueError."""
+        with pytest.raises(ValueError, match="Overlap size.*must satisfy"):
+            IncrementalAudioBufferer(
+                sample_rate=16000,
+                buffer_size_in_secs=5.0,
+                chunk_size_in_secs=2.5,
+                overlap_size_in_secs=-0.1,
+            )
+
+    @pytest.mark.unit
+    def test_constructor_overlap_exceeds_buffer_raises(self):
+        """Overlap > buffer_size raises ValueError."""
+        with pytest.raises(ValueError, match="Overlap size.*must satisfy"):
+            IncrementalAudioBufferer(
+                sample_rate=16000,
+                buffer_size_in_secs=5.0,
+                chunk_size_in_secs=2.5,
+                overlap_size_in_secs=6.0,
+            )
+
+    @pytest.mark.unit
+    def test_constructor_buffer_not_divisible_by_chunk_raises(self):
+        """Buffer size not divisible by chunk size raises ValueError."""
+        with pytest.raises(ValueError, match="Buffer size.*must be divisible by chunk size"):
+            IncrementalAudioBufferer(
+                sample_rate=16000,
+                buffer_size_in_secs=5.0,
+                chunk_size_in_secs=1.7,
+                overlap_size_in_secs=1.7,
+            )
+
+    @pytest.mark.unit
+    def test_constructor_overlap_not_divisible_by_chunk_raises(self):
+        """Overlap not divisible by chunk size raises ValueError."""
+        with pytest.raises(ValueError, match="Overlap size.*must be divisible by chunk size"):
+            IncrementalAudioBufferer(
+                sample_rate=16000,
+                buffer_size_in_secs=5.0,
+                chunk_size_in_secs=2.5,
+                overlap_size_in_secs=2.0,
+            )
+
+    @pytest.mark.unit
+    def test_update_single_frame(self):
+        """Single frame update fills start of buffer and decreases remaining_capacity."""
+        buf = IncrementalAudioBufferer(
+            sample_rate=16000,
+            buffer_size_in_secs=5.0,
+            chunk_size_in_secs=2.5,
+            overlap_size_in_secs=2.5,
+        )
+        chunk_size = 40000  # 2.5 * 16000
+        samples = torch.arange(chunk_size, dtype=torch.float32)
+        frame = _make_frame(samples)
+        buf.update(frame)
+        assert buf.head == chunk_size
+        assert buf.remaining_capacity == buf.buffer_size - chunk_size
+        assert torch.allclose(buf.sample_buffer[:chunk_size], samples, atol=1e-5)
+        assert not buf.is_full()
+
+    @pytest.mark.unit
+    def test_update_multiple_frames_until_full(self):
+        """Multiple updates fill buffer; is_full() becomes True when capacity is 0."""
+        buf = IncrementalAudioBufferer(
+            sample_rate=16000,
+            buffer_size_in_secs=5.0,
+            chunk_size_in_secs=2.5,
+            overlap_size_in_secs=2.5,
+        )
+        chunk_size = 40000
+        for i in range(2):
+            samples = torch.full((chunk_size,), float(i), dtype=torch.float32)
+            frame = _make_frame(samples)
+            buf.update(frame)
+        assert buf.remaining_capacity == 0
+        assert buf.is_full()
+        assert buf.head == buf.buffer_size
+        assert torch.allclose(buf.sample_buffer[:chunk_size], torch.zeros(chunk_size), atol=1e-5)
+        assert torch.allclose(buf.sample_buffer[chunk_size:], torch.ones(chunk_size), atol=1e-5)
+
+    @pytest.mark.unit
+    def test_update_frame_exceeds_buffer_raises(self):
+        """Frame larger than buffer size raises RuntimeError."""
+        buf = IncrementalAudioBufferer(
+            sample_rate=16000,
+            buffer_size_in_secs=5.0,
+            chunk_size_in_secs=2.5,
+            overlap_size_in_secs=2.5,
+        )
+        oversized = torch.zeros(buf.buffer_size + 1)
+        frame = _make_frame(oversized)
+        with pytest.raises(RuntimeError, match="Frame size.*exceeds buffer size"):
+            buf.update(frame)
+
+    @pytest.mark.unit
+    def test_incremental_audio_bufferer_with_mono_stream(self, test_audios):
+        """Integration: feed frames from MonoStream; buffer contents and paddings are consistent."""
+        sample_rate = 16000
+        chunk_size_in_secs = 2.5
+        buffer_size_in_secs = 5.0
+        overlap_size_in_secs = 2.5
+        for audio in test_audios:
+            stream = MonoStream(sample_rate, frame_size_in_secs=chunk_size_in_secs, stream_id=0, pad_last_frame=False)
+            stream.load_audio(audio, options=None)
+            buf = IncrementalAudioBufferer(
+                sample_rate=sample_rate,
+                buffer_size_in_secs=buffer_size_in_secs,
+                chunk_size_in_secs=chunk_size_in_secs,
+                overlap_size_in_secs=overlap_size_in_secs,
+            )
+            for frame in iter(stream):
+                frame = frame[0]
+                buf.update(frame)
+                # Newest frame is at [head - frame.size : head]; after update it's at [head - frame.size : head]
+                start = buf.head - frame.size
+                assert torch.allclose(buf.sample_buffer[start : buf.head], frame.samples, atol=1e-5)
+                assert buf.remaining_capacity == max(0, buf.remaining_capacity)

tests/collections/asr/inference/test_itn.py

@@ -0,0 +1,179 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+nemo_text_processing = pytest.importorskip("nemo_text_processing", reason="Requires nemo_text_processing")
+
+from nemo.collections.asr.inference.itn.inverse_normalizer import AlignmentPreservingInverseNormalizer
+
+
+@pytest.fixture(scope="module")
+def en_itn_model():
+    return AlignmentPreservingInverseNormalizer(
+        lang="en", input_case=AlignmentPreservingInverseNormalizer.LOWER_CASED, cache_dir=None
+    )
+
+
+@pytest.fixture(scope="module")
+def de_itn_model():
+    return AlignmentPreservingInverseNormalizer(
+        lang="de", input_case=AlignmentPreservingInverseNormalizer.LOWER_CASED, cache_dir=None
+    )
+
+
+@pytest.fixture(scope="module")
+def es_itn_model():
+    return AlignmentPreservingInverseNormalizer(
+        lang="es", input_case=AlignmentPreservingInverseNormalizer.LOWER_CASED, cache_dir=None
+    )
+
+
+class TestAlignmentPreservingInverseNormalizer:
+
+    @pytest.mark.unit
+    def test_word_alignment_cardinal_en(self, en_itn_model):
+        text = "zzz minus twenty five thousand thirty seven zzz"
+        iwords, owords, alignment = en_itn_model.get_word_alignment(text, sep=" ")
+        assert iwords == ["zzz", "minus", "twenty", "five", "thousand", "thirty", "seven", "zzz"]
+        assert owords == ["zzz", "-25037", "zzz"]
+        assert alignment == [([0], [0], "name"), ([1, 2, 3, 4, 5, 6], [1], "cardinal"), ([7], [2], "name")]
+
+    @pytest.mark.unit
+    def test_word_alignment_time_en(self, en_itn_model):
+        text = "zzz eleven fifty five p m zzz"
+        iwords, owords, alignment = en_itn_model.get_word_alignment(text, sep=" ")
+        assert iwords == ["zzz", "eleven", "fifty", "five", "p", "m", "zzz"]
+        assert owords == ["zzz", "11:55", "p.m.", "zzz"]
+        assert alignment == [([0], [0], "name"), ([1, 2, 3, 4, 5], [1, 2], "time"), ([6], [3], "name")]
+
+    @pytest.mark.unit
+    def test_word_alignment_money_en(self, en_itn_model):
+        text = "zzz two hundred fifty dollars zzz"
+        iwords, owords, alignment = en_itn_model.get_word_alignment(text, sep=" ")
+        assert iwords == ["zzz", "two", "hundred", "fifty", "dollars", "zzz"]
+        assert owords == ["zzz", "$250", "zzz"]
+        assert alignment == [([0], [0], "name"), ([1, 2, 3, 4], [1], "money"), ([5], [2], "name")]
+
+    @pytest.mark.unit
+    def test_word_alignment_combo_en(self, en_itn_model):
+        text = "eleven twenty seven fifty seven october twenty fourth nineteen seventy"
+        iwords, owords, alignment = en_itn_model.get_word_alignment(text, sep=" ")
+        assert iwords == [
+            "eleven",
+            "twenty",
+            "seven",
+            "fifty",
+            "seven",
+            "october",
+            "twenty",
+            "fourth",
+            "nineteen",
+            "seventy",
+        ]
+        assert owords == ["1120", "07:57", "october", "24", "1970"]
+        assert alignment == [([0, 1], [0], "date"), ([2, 3, 4], [1], "time"), ([5, 6, 7, 8, 9], [2, 3, 4], "date")]
+
+    @pytest.mark.unit
+    def test_word_alignment_measure_en(self, en_itn_model):
+        text = "it is two hundred fifty meters long"
+        iwords, owords, alignment = en_itn_model.get_word_alignment(text, sep=" ")
+        assert iwords == ["it", "is", "two", "hundred", "fifty", "meters", "long"]
+        assert owords == ["it", "is", "250", "m", "long"]
+        assert alignment == [
+            ([0], [0], "name"),
+            ([1], [1], "name"),
+            ([2, 3, 4, 5], [2, 3], "measure"),
+            ([6], [4], "name"),
+        ]
+
+    @pytest.mark.unit
+    def test_word_alignment_sterling_en(self, en_itn_model):
+        text = "trade turnover of three million pounds sterling"
+        iwords, owords, alignment = en_itn_model.get_word_alignment(text, sep=" ")
+        assert iwords == ["trade", "turnover", "of", "three", "million", "pounds", "sterling"]
+        assert owords == ["trade", "turnover", "of", "£3", "million"]
+        assert alignment == [
+            ([0], [0], "name"),
+            ([1], [1], "name"),
+            ([2], [2], "name"),
+            ([3, 4, 5, 6], [3, 4], "money"),
+        ]
+
+    @pytest.mark.unit
+    def test_word_alignment_time_de(self, de_itn_model):
+        text = "zzz drei uhr zwanzig zzz"
+        iwords, owords, alignment = de_itn_model.get_word_alignment(text, sep=" ")
+        assert iwords == ["zzz", "drei", "uhr", "zwanzig", "zzz"]
+        assert owords == ['zzz', '03:20', 'Uhr', 'zzz']
+        assert alignment == [([0], [0], "name"), ([1, 2, 3], [1, 2], "time"), ([4], [3], "name")]
+
+    @pytest.mark.unit
+    def test_word_alignment_money_de(self, de_itn_model):
+        text = "zzz zwei hundert fünfzig dollar zzz"
+        iwords, owords, alignment = de_itn_model.get_word_alignment(text, sep=" ")
+        assert iwords == ["zzz", "zwei", "hundert", "fünfzig", "dollar", "zzz"]
+        assert owords == ["zzz", "$250", "zzz"]
+        assert alignment == [([0], [0], "name"), ([1, 2, 3, 4], [1], "money"), ([5], [2], "name")]
+
+    @pytest.mark.unit
+    def test_word_alignment_cardinal_de(self, de_itn_model):
+        text = "zzz minus fünfundzwanzigtausendsiebenunddreißig zzz"
+        iwords, owords, alignment = de_itn_model.get_word_alignment(text, sep=" ")
+        assert iwords == ["zzz", "minus", "fünfundzwanzigtausendsiebenunddreißig", "zzz"]
+        assert owords == ["zzz", "-25037", "zzz"]
+        assert alignment == [([0], [0], "name"), ([1, 2], [1], "cardinal"), ([3], [2], "name")]
+
+    @pytest.mark.unit
+    def test_word_alignment_measure_de(self, de_itn_model):
+        text = "es ist zweihundertfünfzig meter lang"
+        iwords, owords, alignment = de_itn_model.get_word_alignment(text, sep=" ")
+        assert iwords == ["es", "ist", "zweihundertfünfzig", "meter", "lang"]
+        assert owords == ["es", "ist", "250", "m", "lang"]
+        assert alignment == [([0], [0], "name"), ([1], [1], "name"), ([2, 3], [2, 3], "measure"), ([4], [4], "name")]
+
+    @pytest.mark.unit
+    def test_word_alignment_combo_es(self, es_itn_model):
+        text = "un mil intereses al diez por ciento a la semana estándar derecho diez por ciento"
+        iwords, owords, alignment = es_itn_model.get_word_alignment(text, sep=" ")
+        assert iwords == [
+            'un',
+            'mil',
+            'intereses',
+            'al',
+            'diez',
+            'por',
+            'ciento',
+            'a',
+            'la',
+            'semana',
+            'estándar',
+            'derecho',
+            'diez',
+            'por',
+            'ciento',
+        ]
+        assert owords == ['1000', 'intereses', 'al', '10', '%', 'a', 'la', 'semana', 'estándar', 'derecho', '10', '%']
+        assert alignment == [
+            ([0, 1], [0], 'cardinal'),
+            ([2], [1], 'name'),
+            ([3], [2], 'name'),
+            ([4, 5, 6], [3, 4], 'measure'),
+            ([7], [5], 'name'),
+            ([8], [6], 'name'),
+            ([9], [7], 'name'),
+            ([10], [8], 'name'),
+            ([11], [9], 'name'),
+            ([12, 13, 14], [10, 11], 'measure'),
+        ]

tests/collections/asr/inference/test_itn_utils.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from nemo.collections.asr.inference.utils.itn_utils import (
+    fallback_to_trivial_alignment,
+    find_tokens,
+    get_semiotic_class,
+    get_trivial_alignment,
+    split_text,
+)
+
+
+class TestItnUtils:
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "text, expected_words, expected_n",
+        [
+            ("hello world    how are you", ["hello", "world", "how", "are", "you"], 5),
+            ("hello", ["hello"], 1),
+            ("a hello world    b ccc d   e", ["a", "hello", "world", "b", "ccc", "d", "e"], 7),
+            (" a hello world    b ccc d   e", ["a", "hello", "world", "b", "ccc", "d", "e"], 7),
+            ("a hello world    b ccc d   e ", ["a", "hello", "world", "b", "ccc", "d", "e"], 7),
+            (" a hello world    b ccc d   e ", ["a", "hello", "world", "b", "ccc", "d", "e"], 7),
+            ("  a hello world    b ccc d   e  ", ["a", "hello", "world", "b", "ccc", "d", "e"], 7),
+        ],
+    )
+    def test_split_text(self, text, expected_words, expected_n):
+        words, n = split_text(text)
+        assert words == expected_words
+        assert n == expected_n
+
+    @pytest.mark.unit
+    def test_get_semiotic_class(self):
+        tokens = [{"tokens": {"name": "hello"}}]
+        semiotic_class = get_semiotic_class(tokens)
+        assert semiotic_class == "name"
+
+    @pytest.mark.unit
+    def test_find_tokens(self):
+        text = "tokens {name: hello}   tokens {name: world} tokens {name: how} tokens {name: are} tokens {name: you}"
+        tokens = find_tokens(text)
+        assert tokens == [
+            "tokens {name: hello}",
+            "tokens {name: world}",
+            "tokens {name: how}",
+            "tokens {name: are}",
+            "tokens {name: you}",
+        ]
+
+    @pytest.mark.unit
+    def test_get_trivial_alignment(self):
+        N = 5
+        i_shift = 1
+        o_shift = 2
+        alignment = get_trivial_alignment(N, i_shift, o_shift)
+        assert alignment == [
+            ([1], [2], "name"),
+            ([2], [3], "name"),
+            ([3], [4], "name"),
+            ([4], [5], "name"),
+            ([5], [6], "name"),
+        ]
+
+    @pytest.mark.unit
+    def test_fallback_to_trivial_alignment(self):
+        input_words = ["hello", "world", "how", "are", "you"]
+        input_words, output_words, word_alignment = fallback_to_trivial_alignment(input_words)
+        assert input_words == ["hello", "world", "how", "are", "you"]
+        assert output_words == ["hello", "world", "how", "are", "you"]
+        assert word_alignment == [
+            ([0], [0], "name"),
+            ([1], [1], "name"),
+            ([2], [2], "name"),
+            ([3], [3], "name"),
+            ([4], [4], "name"),
+        ]

tests/collections/asr/inference/test_lcs_merge.py

@@ -0,0 +1,254 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from nemo.collections.asr.inference.utils.lcs_merge import MergingStrategy, lcs_merge, longest_common_substring
+
+
+class TestLCSMerge:
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "buffer, data, expected_start1, expected_start2, expected_length",
+        [
+            ([1, 2, 3, 4, 5], [3, 4, 5, 6, 7], 2, 0, 3),
+            ([1, 2], [1], 0, 0, 1),
+            ([1], [1, 2], 0, 0, 1),
+            (
+                [1, 2, 3, 11, 12, 13, 4, 5, 6],
+                [1, 2, 3, 4, 5, 6, 11, 12, 13],
+                6,
+                3,
+                3,
+            ),
+            ([1, 2, 3, 11, 12, 13, 4, 5, 6, 7], [1, 2, 3, 4, 5, 6, 7, 11, 12, 13], 6, 3, 4),
+            ([1, 2, 3], [4, 5, 6], -1, -1, 0),
+            ([1, 2, 3, 1, 2, 3], [1, 2, 3], 3, 0, 3),
+            ([], [], -1, -1, 0),
+            ([1, 2, 3], [], -1, -1, 0),
+            ([1, 1, 1, 1, 1], [1, 1], 3, 0, 2),
+        ],
+    )
+    def test_longest_common_substring(self, buffer, data, expected_start1, expected_start2, expected_length):
+        start1, start2, length = longest_common_substring(buffer, data)
+        assert (start1, start2, length) == (expected_start1, expected_start2, expected_length)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "buffer, data, search_size, min_lcs_length, merging_strategy, expected_result",
+        [
+            ([1, 2, 3, 4, 5], [3, 4, 5, 6, 7], 5, 1, MergingStrategy.LCSUBSTR, [1, 2, 3, 4, 5, 6, 7]),
+            ([1, 2, 3, 4, 5], [6, 7, 8, 9, 10], 5, 1, MergingStrategy.LCSUBSTR, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
+            ([1, 2, 3, 4, 5], [3, 4, 5, 6, 7, 8, 9], 5, 1, MergingStrategy.LCSUBSTR, [1, 2, 3, 4, 5, 6, 7, 8, 9]),
+            ([1, 2, 3, 4, 9], [3, 4, 5, 6, 7], 5, 1, MergingStrategy.LCS, [1, 2, 3, 4, 5, 6, 7]),
+            ([1, 2, 3, 4, 5], [3, 4, 5, 6, 7], 1, 2, MergingStrategy.LCSUBSTR, [1, 2, 3, 4, 5, 3, 4, 5, 6, 7]),
+        ],
+    )
+    def test_lcs_merge(self, buffer, data, search_size, min_lcs_length, merging_strategy, expected_result):
+        result = lcs_merge(
+            buffer,
+            data,
+            search_size=search_size,
+            min_lcs_length=min_lcs_length,
+            merging_strategy=merging_strategy,
+            sep_id=None,
+        )
+        assert result == expected_result
+
+    @pytest.mark.unit
+    def test_lcs_merge_empty_buffer(self):
+        """Test that empty buffer returns just the data."""
+        result = lcs_merge([], [1, 2, 3], search_size=5, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR)
+        assert result == [1, 2, 3]
+
+    @pytest.mark.unit
+    def test_lcs_merge_empty_data(self):
+        """Test that empty data returns the buffer unchanged."""
+        result = lcs_merge([1, 2, 3], [], search_size=5, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR)
+        assert result == [1, 2, 3]
+
+    @pytest.mark.unit
+    def test_lcs_merge_both_empty(self):
+        """Test merging two empty lists."""
+        result = lcs_merge([], [], search_size=5, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR)
+        assert result == []
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("search_size", [0, -1, -10])
+    def test_lcs_merge_invalid_search_size(self, search_size):
+        """Test that search_size <= 0 results in simple concatenation with separator."""
+        buffer = [1, 2, 3]
+        data = [4, 5, 6]
+        result = lcs_merge(
+            buffer, data, search_size=search_size, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR
+        )
+        assert result == [1, 2, 3, 4, 5, 6]
+
+    @pytest.mark.unit
+    def test_lcs_merge_with_sep_id_no_overlap(self):
+        """Test separator is inserted when no LCS is found."""
+        buffer = [1, 2, 3]
+        data = [7, 8, 9]
+        sep_id = [100]
+        result = lcs_merge(
+            buffer, data, search_size=3, sep_id=sep_id, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR
+        )
+        assert result == [1, 2, 3, 100, 7, 8, 9]
+
+    @pytest.mark.unit
+    def test_lcs_merge_with_multi_token_sep_id(self):
+        """Test separator with multiple tokens is inserted correctly."""
+        buffer = [1, 2, 3]
+        data = [7, 8, 9]
+        sep_id = [100, 101, 102]
+        result = lcs_merge(
+            buffer, data, search_size=3, sep_id=sep_id, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR
+        )
+        assert result == [1, 2, 3, 100, 101, 102, 7, 8, 9]
+
+    @pytest.mark.unit
+    def test_lcs_merge_with_sep_id_when_overlap_exists(self):
+        """Test separator is NOT inserted when LCS is found."""
+        buffer = [1, 2, 3, 4, 5]
+        data = [4, 5, 6, 7]
+        sep_id = [100]
+        result = lcs_merge(
+            buffer, data, search_size=5, sep_id=sep_id, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR
+        )
+        assert result == [1, 2, 3, 4, 5, 6, 7]
+        assert 100 not in result
+
+    @pytest.mark.unit
+    def test_lcs_merge_search_size_larger_than_buffer(self):
+        """Test that search_size larger than buffer still works correctly."""
+        buffer = [1, 2, 3]
+        data = [2, 3, 4, 5]
+        result = lcs_merge(buffer, data, search_size=100, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR)
+        assert result == [1, 2, 3, 4, 5]
+
+    @pytest.mark.unit
+    def test_lcs_merge_search_size_limits_overlap_detection(self):
+        """Test that search_size limits the overlap detection window."""
+        buffer = [1, 2, 3, 4, 5, 6, 7, 8]
+        data = [2, 3, 9, 10]  # overlap [2,3] is outside search window of last 3 elements
+        result = lcs_merge(buffer, data, search_size=3, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR)
+        # No overlap in last 3 elements [6,7,8], so data is appended
+        assert result == [1, 2, 3, 4, 5, 6, 7, 8, 2, 3, 9, 10]
+
+    @pytest.mark.unit
+    def test_lcs_merge_min_lcs_length_threshold(self):
+        """Test that LCS shorter than min_lcs_length causes concatenation."""
+        buffer = [1, 2, 3, 4, 5]
+        data = [5, 6, 7]  # only 1 element overlap
+        result = lcs_merge(buffer, data, search_size=5, min_lcs_length=2, merging_strategy=MergingStrategy.LCSUBSTR)
+        # LCS length is 1, which is < min_lcs_length=2, so concatenate
+        assert result == [1, 2, 3, 4, 5, 5, 6, 7]
+
+    @pytest.mark.unit
+    def test_lcs_merge_min_lcs_length_exact_match(self):
+        """Test that LCS equal to min_lcs_length triggers merge."""
+        buffer = [1, 2, 3, 4, 5]
+        data = [4, 5, 6, 7]  # 2 element overlap
+        result = lcs_merge(buffer, data, search_size=5, min_lcs_length=2, merging_strategy=MergingStrategy.LCSUBSTR)
+        assert result == [1, 2, 3, 4, 5, 6, 7]
+
+    @pytest.mark.unit
+    def test_lcs_merge_data_is_subset_of_buffer_end(self):
+        """Test when data is entirely contained in the end of buffer."""
+        buffer = [1, 2, 3, 4, 5]
+        data = [4, 5]
+        result = lcs_merge(buffer, data, search_size=5, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR)
+        assert result == [1, 2, 3, 4, 5]
+
+    @pytest.mark.unit
+    def test_lcs_merge_complete_overlap(self):
+        """Test when data starts exactly where buffer search window begins."""
+        buffer = [1, 2, 3, 4, 5]
+        data = [3, 4, 5, 6, 7, 8]
+        result = lcs_merge(buffer, data, search_size=3, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR)
+        assert result == [1, 2, 3, 4, 5, 6, 7, 8]
+
+    @pytest.mark.unit
+    def test_lcs_merge_lcs_strategy_with_gaps(self):
+        """Test LCS strategy handles non-contiguous common subsequences."""
+        buffer = [1, 2, 100, 3, 4]  # has 100 inserted
+        data = [2, 3, 4, 5, 6]  # continuous 2, 3, 4
+        result = lcs_merge(buffer, data, search_size=5, min_lcs_length=1, merging_strategy=MergingStrategy.LCS)
+        # LCS finds [2, 3, 4] even with gap
+        assert result == [1, 2, 100, 3, 4, 5, 6]
+
+    @pytest.mark.unit
+    def test_lcs_merge_single_element_overlap(self):
+        """Test merging with single element overlap."""
+        buffer = [1, 2, 3]
+        data = [3, 4, 5]
+        result = lcs_merge(buffer, data, search_size=3, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR)
+        assert result == [1, 2, 3, 4, 5]
+
+    @pytest.mark.unit
+    def test_lcs_merge_repeated_elements(self):
+        """Test merging lists with repeated elements."""
+        buffer = [1, 1, 1, 2, 2, 2]
+        data = [2, 2, 2, 3, 3, 3]
+        result = lcs_merge(buffer, data, search_size=6, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR)
+        assert result == [1, 1, 1, 2, 2, 2, 3, 3, 3]
+
+    @pytest.mark.unit
+    def test_lcs_merge_long_sequences(self):
+        """Test merging longer sequences for performance sanity."""
+        buffer = list(range(100))
+        data = list(range(90, 150))  # overlap from 90-99
+        result = lcs_merge(buffer, data, search_size=20, min_lcs_length=1, merging_strategy=MergingStrategy.LCSUBSTR)
+        assert result == list(range(150))
+
+
+class TestLongestCommonSubstringEdgeCases:
+    """Additional edge case tests for longest_common_substring function."""
+
+    @pytest.mark.unit
+    def test_single_element_match(self):
+        """Test with single element lists that match."""
+        start1, start2, length = longest_common_substring([5], [5])
+        assert (start1, start2, length) == (0, 0, 1)
+
+    @pytest.mark.unit
+    def test_single_element_no_match(self):
+        """Test with single element lists that don't match."""
+        start1, start2, length = longest_common_substring([5], [6])
+        assert (start1, start2, length) == (-1, -1, 0)
+
+    @pytest.mark.unit
+    def test_match_at_buffer_start(self):
+        """Test when LCS is at the start of buffer."""
+        start1, start2, length = longest_common_substring([1, 2, 3, 9, 9], [1, 2, 3, 8, 8])
+        assert (start1, start2, length) == (0, 0, 3)
+
+    @pytest.mark.unit
+    def test_match_at_data_end(self):
+        """Test when LCS is at the end of data."""
+        start1, start2, length = longest_common_substring([7, 8, 9], [1, 2, 7, 8, 9])
+        assert (start1, start2, length) == (0, 2, 3)
+
+    @pytest.mark.unit
+    def test_entire_data_is_substring(self):
+        """Test when entire data is a substring of buffer."""
+        start1, start2, length = longest_common_substring([1, 2, 3, 4, 5, 6], [3, 4, 5])
+        assert (start1, start2, length) == (2, 0, 3)
+
+    @pytest.mark.unit
+    def test_entire_buffer_is_substring(self):
+        """Test when entire buffer is a substring of data."""
+        start1, start2, length = longest_common_substring([3, 4, 5], [1, 2, 3, 4, 5, 6])
+        assert (start1, start2, length) == (0, 2, 3)

tests/collections/asr/inference/test_pipeline_eval.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.inference.utils.pipeline_eval import calculate_asr_laal
+from nemo.collections.asr.parts.utils.eval_utils import compute_laal
+
+# EoU enabled (stop_history_eou >= 0) so ASR LAAL is computed.
+CFG = OmegaConf.create({"metrics": {"asr": {"gt_text_attr_name": "text"}}, "endpointing": {"stop_history_eou": 800}})
+
+
+class TestCalculateAsrLaal:
+    @pytest.mark.unit
+    def test_matches_direct_compute_laal(self):
+        # Two finalized steps: "hello world" committed at 2.0 s, "foo" at 5.0 s of audio elapsed.
+        # Each word inherits its step's delay (ms); LAAL is computed against the reference word count.
+        durations = {"a.wav": 10.0}  # seconds -> 10000 ms
+        manifest = [{"audio_filepath": "a.wav", "text": "hello world foo bar"}]  # 4 reference words
+        output = {0: {"audio_filepath": "a.wav", "asr_segments": [("hello world", 2.0), (" foo", 5.0)]}}
+
+        expected = compute_laal([2000.0, 2000.0, 5000.0], 10000.0, 4)
+        got = calculate_asr_laal(output, durations, manifest, CFG)
+        assert got == pytest.approx(expected)
+
+    @pytest.mark.unit
+    def test_delay_capped_at_duration(self):
+        # A delay beyond the audio duration is clamped to the duration.
+        durations = {"a.wav": 3.0}  # 3000 ms
+        manifest = [{"audio_filepath": "a.wav", "text": "hello world"}]
+        output = {0: {"audio_filepath": "a.wav", "asr_segments": [("hello world", 99.0)]}}
+
+        expected = compute_laal([3000.0, 3000.0], 3000.0, 2)
+        assert calculate_asr_laal(output, durations, manifest, CFG) == pytest.approx(expected)
+
+    @pytest.mark.unit
+    def test_no_manifest_returns_none(self):
+        output = {0: {"audio_filepath": "a.wav", "asr_segments": [("hi", 1.0)]}}
+        assert calculate_asr_laal(output, {"a.wav": 1.0}, None, CFG) is None
+
+    @pytest.mark.unit
+    def test_no_reference_returns_none(self):
+        # Stream's audio is absent from the manifest -> nothing to score.
+        output = {0: {"audio_filepath": "missing.wav", "asr_segments": [("hi", 1.0)]}}
+        manifest = [{"audio_filepath": "other.wav", "text": "hello"}]
+        assert calculate_asr_laal(output, {"missing.wav": 1.0}, manifest, CFG) is None
+
+    @pytest.mark.unit
+    def test_eou_disabled_returns_none(self):
+        # EoU disabled (stop_history_eou < 0) -> one segment per utterance, no latency signal -> skipped.
+        cfg = OmegaConf.create(
+            {"metrics": {"asr": {"gt_text_attr_name": "text"}}, "endpointing": {"stop_history_eou": -1}}
+        )
+        manifest = [{"audio_filepath": "a.wav", "text": "hello world"}]
+        output = {0: {"audio_filepath": "a.wav", "asr_segments": [("hello world", 2.0)]}}
+        assert calculate_asr_laal(output, {"a.wav": 10.0}, manifest, cfg) is None

tests/collections/asr/inference/test_pipeline_utils.py

@@ -0,0 +1,78 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import re
+
+import pytest
+import torch
+
+from nemo.collections.asr.inference.utils.pipeline_utils import (
+    check_existance_of_required_attributes,
+    drop_trailing_features,
+    get_leading_punctuation_regex_pattern,
+    get_repeated_punctuation_regex_pattern,
+)
+
+
+class TestPipelineUtils:
+
+    @pytest.mark.unit
+    def test_drop_trailing_features(self):
+        x = torch.randn(10, 10, 20)
+        expected_feature_buffer_len = 15
+        x_dropped = drop_trailing_features(x, expected_feature_buffer_len)
+        assert x_dropped.shape == (10, 10, 15)
+        assert x_dropped.allclose(x[:, :, :15])
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "text, expected_text",
+        [
+            ("", ""),
+            (" ", " "),
+            ("simple text", "simple text"),
+            ("just a 2nd . Yeah, I hope", "just a 2nd. Yeah, I hope"),
+            ("Hello , world ! How are you ?", "Hello, world! How are you?"),
+            ("The quick, brown fox jumps ? over the lazy ! dog.", "The quick, brown fox jumps? over the lazy! dog."),
+        ],
+    )
+    def test_remove_leading_punctuation_spaces(self, text, expected_text):
+        puncts = {"!", "?", ".", ","}
+        pattern = get_leading_punctuation_regex_pattern(puncts)
+        assert re.sub(pattern, r'\1', text) == expected_text
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "text, expected_text",
+        [
+            ("", ""),
+            (" ", " "),
+            ("simple text", "simple text"),
+            ("Hello, world!! How are you???", "Hello, world! How are you?"),
+            ("The quick,, brown fox jumps? over the lazy! dog..", "The quick, brown fox jumps? over the lazy! dog."),
+        ],
+    )
+    def test_remove_repeated_punctuation(self, text, expected_text):
+        puncts = {"!", "?", ".", ","}
+        pattern = get_repeated_punctuation_regex_pattern(puncts)
+        assert re.sub(pattern, r'\1', text) == expected_text
+
+    @pytest.mark.unit
+    def test_check_existance_of_required_attributes(self):
+        class TestClass:
+            pass
+
+        with pytest.raises(ValueError):
+            check_existance_of_required_attributes(TestClass, ['test_attr'])

tests/collections/asr/inference/test_state_management_utils.py

@@ -0,0 +1,109 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytest
+
+from nemo.collections.asr.inference.utils.state_management_utils import (
+    detect_overlap,
+    find_max_overlap,
+    merge_segment_tail,
+    merge_timesteps,
+    merge_word_tail,
+)
+from nemo.collections.asr.inference.utils.text_segment import TextSegment, Word
+
+
+class TestStateManagementUtils:
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "timesteps1, timesteps2, expected_merged_timesteps",
+        [
+            ([0, 1, 2, 3], [4, 5, 6, 7], [0, 1, 2, 3, 4, 5, 6, 7]),
+            ([0, 1, 2, 3], [], [0, 1, 2, 3]),
+            ([], [4, 5, 6, 7], [4, 5, 6, 7]),
+            ([-1, 0, 1, 2], [0, 1, 2, 3], [0, 1, 2, 3, 4, 5, 6, 7]),
+            ([-3, 1, 2, 3], [], [0, 4, 5, 6]),
+            ([], [-3, 1, 2, 3], [0, 4, 5, 6]),
+        ],
+    )
+    def test_merge_timesteps(self, timesteps1, timesteps2, expected_merged_timesteps):
+        merged_timesteps = merge_timesteps(timesteps1, timesteps2)
+        assert merged_timesteps == expected_merged_timesteps
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "state_tokens, new_tokens, limit, expected_max_overlap",
+        [
+            ([0, 1, 2, 3], [2, 3, 4, 5], 4, 2),
+            ([0, 2, 3, 4], [2, 3, 4, 5], 4, 3),
+            ([0, 0, 0, 1], [2, 3, 4, 5], 4, 0),
+        ],
+    )
+    def test_find_max_overlap(self, state_tokens, new_tokens, limit, expected_max_overlap):
+        max_overlap = find_max_overlap(state_tokens, new_tokens, limit)
+        assert max_overlap == expected_max_overlap
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "state_tokens, state_timesteps, new_tokens, new_timesteps, expected_overlap",
+        [
+            ([0, 1, 2, 3], [0.0, 1.0, 2.0, 3.0], [2, 3, 4, 5], [2.0, 3.0, 4.0, 5.0], 2),
+            ([0, 1, 2, 3], [0.0, 1.0, 2.0, 3.0], [2, 3, 4, 5], [1.0, 2.0, 4.0, 5.0], 2),
+            ([0, 1, 2, 3], [0.0, 1.0, 2.0, 3.0], [2, 3, 4, 5], [5.0, 7.0, 8.0, 9.0], 0),
+        ],
+    )
+    def test_detect_overlap(self, state_tokens, state_timesteps, new_tokens, new_timesteps, expected_overlap):
+        overlap = detect_overlap(state_tokens, state_timesteps, new_tokens, new_timesteps)
+        assert overlap == expected_overlap
+
+    @pytest.mark.unit
+    def test_merge_word_tail_without_pnc(self):
+        word_head = Word(text="meaning", start=0.0, end=1.0, conf=0.5)
+        word_tail = Word(text="ful", start=1.0, end=2.0, conf=0.6)
+        head, _ = merge_word_tail(word_head, word_tail, conf_aggregator=min)
+
+        assert head.text == "meaningful"
+        assert head.start == 0.0
+        assert head.end == 2.0
+        assert head.conf == 0.5
+
+    @pytest.mark.unit
+    def test_merge_word_tail_with_pnc(self):
+
+        word_head = Word(text="meaning", start=0.0, end=1.0, conf=0.5)
+        word_tail = Word(text="s", start=1.0, end=2.0, conf=0.6)
+        pnc_head = Word(text="Meaning?", start=0.0, end=1.0, conf=0.5)
+        new_head, new_pnc_head = merge_word_tail(word_head, word_tail, conf_aggregator=min, pnc_word_head=pnc_head)
+
+        assert new_head.text == "meanings"
+        assert new_head.start == 0.0
+        assert new_head.end == 2.0
+        assert new_head.conf == 0.5
+        assert new_pnc_head.text == "Meanings?"
+        assert new_pnc_head.start == 0.0
+        assert new_pnc_head.end == 2.0
+        assert new_pnc_head.conf == 0.5
+
+    @pytest.mark.unit
+    def test_merge_segment_tail(self):
+        seg1 = TextSegment(text="Good morn", start=0.0, end=1.0, conf=0.5)
+        seg2 = TextSegment(text="ing", start=1.0, end=2.0, conf=0.6)
+        merged_seg = merge_segment_tail(seg1, seg2, conf_aggregator=min)
+
+        assert merged_seg.text == "Good morning"
+        assert merged_seg.start == 0.0
+        assert merged_seg.end == 2.0
+        assert merged_seg.conf == 0.5

tests/collections/asr/inference/test_text_segment.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytest
+
+from nemo.collections.asr.inference.utils.text_segment import TextSegment, Word, join_segments
+
+
+class TestTextSegment:
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("text, expected_text", [("Hello!", "hello"), ("HeLLo!", "hello")])
+    def test_normalize_text_inplace(self, text, expected_text):
+        for cls in [Word, TextSegment]:
+            text_segment = cls(text, 0, 1, 0.5)
+            text_segment.normalize_text_inplace(punct_marks='!', sep=' ')
+            assert text_segment.text == expected_text
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("text, expected_text", [("Hello!", "hello"), ("HeLLo!", "hello")])
+    def test_with_normalized_text(self, text, expected_text):
+        for cls in [Word, TextSegment]:
+            text_segment = cls(text, 0, 1, 0.5)
+            text_segment_copy = text_segment.with_normalized_text(punct_marks='!', sep=' ')
+            assert text_segment_copy.text == expected_text
+            assert text_segment.text == text
+
+    @pytest.mark.unit
+    def test_join_segments(self):
+        for cls in [Word, TextSegment]:
+            segments = [
+                [cls('hello', 0, 1, 0.5), cls('world', 1, 2, 0.5)],
+                [cls('how', 2, 3, 0.5), cls('are', 3, 4, 0.5), cls('you', 4, 5, 0.5)],
+            ]
+            transcriptions = join_segments(segments, sep=' ')
+            assert transcriptions == ['hello world', 'how are you']
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("text, expected_text", [("hello", "Hello"), ("World!", "World!")])
+    def test_capitalize(self, text, expected_text):
+        for cls in [Word, TextSegment]:
+            text_segment = cls(text, 0, 1, 0.5)
+            text_segment.capitalize()
+            assert text_segment.text == expected_text

tests/collections/asr/k2/test_graph_transducer.py

@@ -0,0 +1,336 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+from typing import List
+
+import numpy as np
+import pytest
+import torch
+
+from nemo.collections.asr.parts.k2.rnnt_logprobs import rnnt_logprobs_torch
+from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_numpy import RNNTLoss as RNNTLoss_Numpy
+from nemo.core.utils.optional_libs import K2_AVAILABLE, TRITON_AVAILABLE
+
+if K2_AVAILABLE:
+    import k2
+
+    from nemo.collections.asr.parts.k2.graph_transducer import GraphRnntLoss
+
+if TRITON_AVAILABLE:
+    from nemo.collections.asr.parts.k2.rnnt_logprobs_triton import rnnt_logprobs_triton
+
+
+EPS_SM_INPUT = 1e-6
+EPS_L_INPUT = 1e-4
+
+DEVICES = ['cpu']
+
+if K2_AVAILABLE and torch.cuda.is_available() and k2.with_cuda:
+    DEVICES.append('cuda')
+
+
+@pytest.mark.skipif(not K2_AVAILABLE, reason="k2 is not installed, skipping Graph-RNNT tests.")
+class TestGraphRnnt:
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("blank_first", [True, False])
+    @pytest.mark.parametrize("num_frames", [1, 3, 6])
+    @pytest.mark.parametrize("vocab_size", [3])
+    def test_temporal_schema(self, device, blank_first, num_frames, vocab_size):
+        blank_id = 0 if blank_first else vocab_size - 1
+        loss = GraphRnntLoss(blank=blank_id)
+        temporal_schema = loss.get_temporal_schema(
+            num_frames=num_frames, vocab_size=vocab_size, device=torch.device(device)
+        )
+
+        etalon_schema_fst: List[List[int]] = []
+        for time_i in range(num_frames):
+            for label_i in range(vocab_size):
+                if label_i == blank_id:
+                    # transition to the next state
+                    etalon_schema_fst.append([time_i, time_i + 1, label_i, time_i, 0])
+                else:
+                    # self-loop
+                    etalon_schema_fst.append([time_i, time_i, label_i, time_i, 0])
+        etalon_schema_fst.append([num_frames, num_frames + 1, -1, -1, 0])  # transition to final state
+        etalon_schema_fst.append([num_frames + 1])  # final state
+        etalon_schema_fst = sorted(etalon_schema_fst)  # required for k2.Fsa.from_str
+        etalon_schema_fst_str = "\n".join([" ".join(map(str, line)) for line in etalon_schema_fst])
+        etalon_temporal_schema = k2.Fsa.from_str(etalon_schema_fst_str, num_aux_labels=1)
+
+        assert temporal_schema.num_arcs == etalon_temporal_schema.num_arcs
+        assert temporal_schema.shape == etalon_temporal_schema.shape  # (num_states, None)
+        assert k2.is_rand_equivalent(
+            temporal_schema, etalon_temporal_schema, log_semiring=True, treat_epsilons_specially=False
+        ), "Temporal schema mismatch"
+        assert k2.is_rand_equivalent(
+            temporal_schema.invert(),
+            etalon_temporal_schema.invert(),
+            log_semiring=True,
+            treat_epsilons_specially=False,
+        ), "Temporal schema output labels mismatch"
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("blank_first", [True, False])
+    def test_unit_schema(self, device, blank_first):
+        vocab_size = 3
+        blank_id = 0 if blank_first else vocab_size - 1
+        if blank_first:
+            labels = [1, 1, 2, 1]
+        else:
+            labels = [1, 1, 0, 1]
+        loss = GraphRnntLoss(blank=blank_id)
+        unit_schema = loss.get_unit_schema(
+            units_tensor=torch.tensor(labels, device=torch.device(device)), vocab_size=vocab_size
+        )
+
+        etalon_schema_fst: List[List[int]] = []
+        for label_i, label in enumerate(labels):
+            etalon_schema_fst.append([label_i, label_i + 1, label, label, label_i, 0])  # forward: label
+            etalon_schema_fst.append([label_i, label_i, blank_id, blank_id, label_i, 0])  # self-loop: blank
+        etalon_schema_fst.append([len(labels), len(labels), blank_id, blank_id, len(labels), 0])
+        etalon_schema_fst.append([len(labels), len(labels) + 1, -1, -1, -1, 0])  # transition to final state
+        etalon_schema_fst.append([len(labels) + 1])  # final state
+        etalon_schema_fst = sorted(etalon_schema_fst)  # required for k2.Fsa.from_str
+        etalon_schema_fst_str = "\n".join([" ".join(map(str, line)) for line in etalon_schema_fst])
+        etalon_unit_schema = k2.Fsa.from_str(etalon_schema_fst_str, aux_label_names=["aux_labels", "unit_positions"])
+
+        assert unit_schema.num_arcs == etalon_unit_schema.num_arcs
+        assert unit_schema.shape == etalon_unit_schema.shape  # (num_states, None)
+        assert k2.is_rand_equivalent(
+            unit_schema, etalon_unit_schema, log_semiring=True, treat_epsilons_specially=False
+        ), "Unit schema input labels mismatch"
+        assert k2.is_rand_equivalent(
+            unit_schema.invert(), etalon_unit_schema.invert(), log_semiring=True, treat_epsilons_specially=False
+        ), "Unit schema output labels mismatch"
+
+        # swap aux_labels and unit positions to test unit_positions
+        unit_schema.aux_labels, unit_schema.unit_positions = unit_schema.unit_positions, unit_schema.aux_labels
+        etalon_unit_schema.aux_labels, etalon_unit_schema.unit_positions = (
+            etalon_unit_schema.unit_positions,
+            etalon_unit_schema.aux_labels,
+        )
+        assert k2.is_rand_equivalent(
+            unit_schema.invert(), etalon_unit_schema.invert(), log_semiring=True, treat_epsilons_specially=False
+        ), "Unit schema unit positions mismatch"
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("blank_first", [True, False])
+    def test_grid_schema(self, device, blank_first):
+        vocab_size = 3
+        blank_id = 0 if blank_first else vocab_size - 1
+        if blank_first:
+            labels = [1, 1, 2, 1]
+        else:
+            labels = [1, 1, 0, 1]
+        text_length = len(labels)
+        num_frames = 5
+        loss = GraphRnntLoss(blank=blank_id)
+        grid_schema = loss.get_grid(
+            units_tensor=torch.tensor(labels, device=torch.device(device)),
+            num_frames=num_frames,
+            vocab_size=vocab_size,
+        )
+
+        etalon_schema_fst: List[List[int]] = []
+        for frame_i in range(num_frames):
+            for label_i in range(text_length + 1):
+                state = frame_i * (text_length + 1) + label_i
+                if label_i < text_length:
+                    next_state_label = state + 1
+                    # next unit
+                    etalon_schema_fst.append([state, next_state_label, labels[label_i], frame_i, label_i, 0])
+                if frame_i < num_frames - 1:
+                    next_state_frame = (frame_i + 1) * (text_length + 1) + label_i
+                    # next time frame (blank)
+                    etalon_schema_fst.append([state, next_state_frame, blank_id, frame_i, label_i, 0])
+
+        last_grid_state = num_frames * (text_length + 1) - 1
+        etalon_schema_fst.append([last_grid_state, last_grid_state + 1, blank_id, num_frames - 1, text_length, 0])
+        etalon_schema_fst.append(
+            [last_grid_state + 1, last_grid_state + 2, -1, -1, -1, 0]
+        )  # transition to final state
+        etalon_schema_fst.append([last_grid_state + 2])  # final state
+        etalon_schema_fst = sorted(etalon_schema_fst)  # required for k2.Fsa.from_str
+        etalon_schema_fst_str = "\n".join([" ".join(map(str, line)) for line in etalon_schema_fst])
+        etalon_grid_schema = k2.Fsa.from_str(etalon_schema_fst_str, aux_label_names=["aux_labels", "unit_positions"])
+
+        assert grid_schema.num_arcs == etalon_grid_schema.num_arcs
+        assert grid_schema.shape == etalon_grid_schema.shape  # (num_states, None)
+        assert k2.is_rand_equivalent(
+            grid_schema, etalon_grid_schema, log_semiring=True, treat_epsilons_specially=False
+        ), "Grid schema input labels mismatch"
+        assert k2.is_rand_equivalent(
+            grid_schema.invert(), etalon_grid_schema.invert(), log_semiring=True, treat_epsilons_specially=False
+        ), "Grid schema output labels mismatch"
+
+        # swap aux_labels and unit positions to test unit_positions
+        grid_schema.aux_labels, grid_schema.unit_positions = grid_schema.unit_positions, grid_schema.aux_labels
+        etalon_grid_schema.aux_labels, etalon_grid_schema.unit_positions = (
+            etalon_grid_schema.unit_positions,
+            etalon_grid_schema.aux_labels,
+        )
+        assert k2.is_rand_equivalent(
+            grid_schema.invert(), etalon_grid_schema.invert(), log_semiring=True, treat_epsilons_specially=False
+        ), "Grid schema unit positions mismatch"
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("connect_composed", [True, False])
+    @pytest.mark.parametrize("blank_first", [True, False])
+    def test_small_compose_transducer(
+        self, device, connect_composed, blank_first, rnnt_test_helper, rnn_loss_sample_data
+    ):
+        if blank_first:
+            sample_data = rnn_loss_sample_data.get_sample_small()
+        else:
+            sample_data = rnn_loss_sample_data.get_sample_small_blank_last()
+        graph_rnnt = GraphRnntLoss(
+            blank=sample_data.blank_id, connect_composed=connect_composed, use_grid_implementation=False
+        )
+        graph_cost, graph_grads = rnnt_test_helper.wrap_and_call(
+            graph_rnnt, sample_data.logits, sample_data.targets, device
+        )
+        assert np.allclose(graph_cost, sample_data.expected_cost.numpy(), rtol=EPS_SM_INPUT), "costs mismatch."
+        assert np.allclose(graph_grads, sample_data.expected_grads.numpy(), atol=1e-6), "gradient mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_small_grid_transducer(self, device, rnnt_test_helper, rnn_loss_sample_data):
+        sample_data = rnn_loss_sample_data.get_sample_small()
+        graph_rnnt = GraphRnntLoss(blank=0, use_grid_implementation=True)
+        graph_cost, graph_grads = rnnt_test_helper.wrap_and_call(
+            graph_rnnt, sample_data.logits, sample_data.targets, device
+        )
+        assert np.allclose(graph_cost, sample_data.expected_cost.numpy(), rtol=EPS_SM_INPUT), "costs mismatch."
+        assert np.allclose(graph_grads, sample_data.expected_grads.numpy(), atol=1e-6), "gradient mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("use_triton", [True, False])
+    def test_medium_grid_transducer(self, device, use_triton: bool, rnnt_test_helper, rnn_loss_sample_data):
+        if use_triton and device == "cpu":
+            pytest.skip("Triton does not support CPU yet")
+        sample_data = rnn_loss_sample_data.get_sample_medium()
+        graph_rnnt = GraphRnntLoss(blank=0, use_grid_implementation=True, use_triton=use_triton)
+        graph_cost, graph_grads = rnnt_test_helper.wrap_and_call(
+            graph_rnnt, sample_data.logits, sample_data.targets, device
+        )
+        assert np.allclose(graph_cost, sample_data.expected_cost.numpy(), rtol=EPS_SM_INPUT), "costs mismatch."
+        assert np.allclose(graph_grads, sample_data.expected_grads.numpy(), atol=1e-6), "gradient mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("use_triton", [True, False])
+    def test_medium_random_var_size(self, device, use_triton: bool, rnnt_test_helper, rnn_loss_sample_data):
+        if use_triton and device == "cpu":
+            pytest.skip("Triton does not support CPU yet")
+        sample_data = rnn_loss_sample_data.get_sample_medium_random_var_size(blank_first=True)
+        graph_rnnt = GraphRnntLoss(blank=0, use_grid_implementation=True, use_triton=use_triton)
+        graph_cost, graph_grads = rnnt_test_helper.wrap_and_call(
+            graph_rnnt,
+            sample_data.logits.detach(),
+            sample_data.targets,
+            device,
+            input_lengths=sample_data.input_lengths,
+            target_lengths=sample_data.target_lengths,
+        )
+        etalon_rnnt = RNNTLoss_Numpy(blank=0)
+        etalon_cost, etalon_grads = rnnt_test_helper.wrap_and_call(
+            etalon_rnnt,
+            sample_data.logits.detach(),
+            sample_data.targets,
+            device,
+            input_lengths=sample_data.input_lengths,
+            target_lengths=sample_data.target_lengths,
+        )
+        assert np.allclose(graph_cost.sum(), etalon_cost, rtol=EPS_SM_INPUT), "costs mismatch."
+        assert np.allclose(graph_grads, etalon_grads, atol=1e-4), "gradient mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("blank_first", [True, False])
+    def test_small_random_grid_compose_equivalent(self, device: torch.device, blank_first: bool, rnn_loss_sample_data):
+        sample_data = rnn_loss_sample_data.get_sample_small_random(blank_first, device=device)
+        criterion = GraphRnntLoss(blank=sample_data.blank_id, connect_composed=True, use_grid_implementation=False)
+        text_tensor = sample_data.targets[0]
+        num_frames = sample_data.logits.shape[1]
+        graph_grid = criterion.get_grid(text_tensor, num_frames, sample_data.vocab_size)
+        graph_composed = criterion.get_composed_lattice(text_tensor, num_frames, sample_data.vocab_size)
+        assert k2.is_rand_equivalent(
+            graph_grid, graph_composed, log_semiring=True, treat_epsilons_specially=False
+        ), "Grid and composed graphs are not equivalent."
+
+
+@pytest.mark.skipif(not TRITON_AVAILABLE, reason="Triton is not installed, skipping RNNT Log Probs tests")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is unavailable")
+class TestRnntLogProbs:
+    @pytest.mark.parametrize(
+        "batch_size,num_frames,num_text_units,vocab_size",
+        [
+            (1, 4, 2, 4),
+            (2, 3, 2, 5),
+            (2, 16, 31, 17),
+            (16, 129, 65, 2048),
+        ],
+    )
+    @pytest.mark.parametrize(
+        "float_dtype",
+        [torch.float32] + ([torch.bfloat16] if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else []),
+    )
+    def test_rnnt_logprobs_random(
+        self, batch_size: int, num_frames: int, num_text_units: int, vocab_size: int, float_dtype: torch.dtype
+    ):
+        """
+        Test Triton-based implementation using etalon Torch-based implementation for RNN-T log-probs.
+        """
+        device = torch.device("cuda")
+        torch.manual_seed(777)
+
+        targets = torch.tensor(
+            [[random.randrange(0, vocab_size - 1) for i in range(num_text_units)] for j in range(batch_size)],
+            device=device,
+            dtype=torch.long,
+        )
+
+        logits = torch.rand(
+            [batch_size, num_frames, num_text_units + 1, vocab_size + 1],
+            dtype=float_dtype,
+            device=device,
+            requires_grad=True,
+        )
+
+        # Triton-based implementation works in float32 precision for accuracy purposes, should compare with float32
+        target_scores_etalon, blank_scores_etalon = rnnt_logprobs_torch(
+            logits=logits.to(torch.float32), targets=targets, blank_id=vocab_size
+        )
+        logits2 = logits.clone().detach()
+        logits2.requires_grad_(True)
+        target_scores, blank_scores = rnnt_logprobs_triton(logits=logits2, targets=targets, blank_id=vocab_size)
+        target_scores[..., -1:] = 0.0
+        target_scores_etalon[..., -1:] = 0.0
+        assert torch.allclose(blank_scores, blank_scores_etalon, atol=1e-5)
+        assert torch.allclose(target_scores, target_scores_etalon, atol=1e-5)
+
+        # test backward
+        target_scales = torch.rand_like(target_scores, requires_grad=False)
+        blank_scales = torch.rand_like(blank_scores, requires_grad=False)
+        loss_etalon = (target_scales * target_scores_etalon + blank_scales * blank_scores_etalon).sum()
+        loss = (target_scales * target_scores + blank_scales * blank_scores).sum()
+        loss_etalon.backward()
+        loss.backward()
+        assert torch.allclose(logits.grad, logits2.grad, atol=1e-5)

tests/collections/asr/k2/test_w_transducer.py

@@ -0,0 +1,260 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+import numpy as np
+import pytest
+import torch
+
+try:
+    from nemo.collections.asr.parts.k2.w_transducer import GraphWTransducerLoss
+    from nemo.core.utils.k2_guard import k2
+except (ImportError, ModuleNotFoundError):
+    pytest.skip("k2 is not installed, skipping Graph-W-Transducer tests.", allow_module_level=True)
+
+DEVICES = ['cpu']
+
+if torch.cuda.is_available() and k2.with_cuda:
+    DEVICES.append('cuda')
+
+
+class TestGraphWTransducerLoss:
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("blank_first", [True, False])
+    @pytest.mark.parametrize("num_frames", [1, 3, 6])
+    @pytest.mark.parametrize("vocab_size", [3])
+    @pytest.mark.parametrize("last_blank_mode", ["force_final", "allow_ignore"])
+    def test_temporal_schema(self, device, blank_first, num_frames, vocab_size, last_blank_mode):
+        blank_id = 0 if blank_first else vocab_size - 1
+        loss = GraphWTransducerLoss(blank=blank_id, last_blank_mode=last_blank_mode)
+        temporal_schema = loss.get_temporal_schema(
+            num_frames=num_frames, vocab_size=vocab_size, device=torch.device(device)
+        )
+
+        etalon_schema_fst: List[List[int]] = []
+        for time_i in range(num_frames):
+            for label_i in range(vocab_size):
+                if label_i == blank_id:
+                    # transition to the next state
+                    etalon_schema_fst.append([time_i, time_i + 1, label_i, time_i, 0])
+                else:
+                    # self-loop
+                    etalon_schema_fst.append([time_i, time_i, label_i, time_i, 0])
+
+        # eps transitions from the first state
+        eps_from_first_state = vocab_size
+        for time_i in range(1, num_frames):
+            etalon_schema_fst.append([0, time_i, eps_from_first_state, 0, 0])
+
+        # eps transitions to the last state
+        eps_to_last_state = vocab_size + 1
+        last_state_eps = num_frames - 1 if last_blank_mode == "force_final" else num_frames
+        for time_i in range(0, num_frames - 1):
+            etalon_schema_fst.append([time_i, last_state_eps, eps_to_last_state, time_i, 0])
+
+        # transition to the final state
+        etalon_schema_fst.append([num_frames, num_frames + 1, -1, -1, 0])
+        # final state
+        etalon_schema_fst.append([num_frames + 1])
+
+        etalon_schema_fst = sorted(etalon_schema_fst)  # required for k2.Fsa.from_str
+        etalon_schema_fst_str = "\n".join([" ".join(map(str, line)) for line in etalon_schema_fst])
+        etalon_temporal_schema = k2.Fsa.from_str(etalon_schema_fst_str, num_aux_labels=1)
+
+        assert temporal_schema.num_arcs == etalon_temporal_schema.num_arcs
+        assert temporal_schema.shape == etalon_temporal_schema.shape  # (num_states, None)
+        assert k2.is_rand_equivalent(
+            temporal_schema, etalon_temporal_schema, log_semiring=True, treat_epsilons_specially=False
+        ), "Temporal schema mismatch"
+        assert k2.is_rand_equivalent(
+            temporal_schema.invert(),
+            etalon_temporal_schema.invert(),
+            log_semiring=False,
+            treat_epsilons_specially=False,
+        ), "Temporal schema output labels mismatch"
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("blank_first", [True, False])
+    def test_unit_schema(self, device, blank_first):
+        vocab_size = 3
+        blank_id = 0 if blank_first else vocab_size - 1
+        if blank_first:
+            labels = [1, 1, 2, 1]
+        else:
+            labels = [1, 1, 0, 1]
+        loss = GraphWTransducerLoss(blank=blank_id)
+        unit_schema = loss.get_unit_schema(
+            units_tensor=torch.tensor(labels, device=torch.device(device)), vocab_size=vocab_size
+        )
+
+        etalon_schema_fst: List[List[int]] = []
+        for label_i, label in enumerate(labels):
+            etalon_schema_fst.append([label_i, label_i + 1, label, label, label_i, 0])  # forward: label
+            etalon_schema_fst.append([label_i, label_i, blank_id, blank_id, label_i, 0])  # self-loop: blank
+        etalon_schema_fst.append([len(labels), len(labels), blank_id, blank_id, len(labels), 0])
+        # eps-transitions
+        etalon_schema_fst.append([0, 0, vocab_size, vocab_size, 0, 0])
+        etalon_schema_fst.append([len(labels), len(labels), vocab_size + 1, vocab_size + 1, len(labels), 0])
+
+        etalon_schema_fst.append([len(labels), len(labels) + 1, -1, -1, -1, 0])  # transition to final state
+        etalon_schema_fst.append([len(labels) + 1])  # final state
+        etalon_schema_fst = sorted(etalon_schema_fst)  # required for k2.Fsa.from_str
+        etalon_schema_fst_str = "\n".join([" ".join(map(str, line)) for line in etalon_schema_fst])
+        etalon_unit_schema = k2.Fsa.from_str(etalon_schema_fst_str, aux_label_names=["aux_labels", "unit_positions"])
+
+        assert unit_schema.num_arcs == etalon_unit_schema.num_arcs
+        assert unit_schema.shape == etalon_unit_schema.shape  # (num_states, None)
+        assert k2.is_rand_equivalent(
+            unit_schema, etalon_unit_schema, log_semiring=True, treat_epsilons_specially=False
+        ), "Unit schema input labels mismatch"
+        assert k2.is_rand_equivalent(
+            unit_schema.invert(), etalon_unit_schema.invert(), log_semiring=True, treat_epsilons_specially=False
+        ), "Unit schema output labels mismatch"
+
+        # swap aux_labels and unit positions to test unit_positions
+        unit_schema.aux_labels, unit_schema.unit_positions = unit_schema.unit_positions, unit_schema.aux_labels
+        etalon_unit_schema.aux_labels, etalon_unit_schema.unit_positions = (
+            etalon_unit_schema.unit_positions,
+            etalon_unit_schema.aux_labels,
+        )
+        assert k2.is_rand_equivalent(
+            unit_schema.invert(), etalon_unit_schema.invert(), log_semiring=True, treat_epsilons_specially=False
+        ), "Unit schema unit positions mismatch"
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("blank_first", [True, False])
+    @pytest.mark.parametrize("last_blank_mode", ["force_final", "allow_ignore"])
+    def test_grid_schema(self, device, blank_first, last_blank_mode):
+        vocab_size = 3
+        blank_id = 0 if blank_first else vocab_size - 1
+        if blank_first:
+            labels = [1, 1, 2, 1]
+        else:
+            labels = [1, 1, 0, 1]
+        text_length = len(labels)
+        num_frames = 5
+        loss = GraphWTransducerLoss(blank=blank_id, last_blank_mode=last_blank_mode)
+        grid_schema = loss.get_grid(
+            units_tensor=torch.tensor(labels, device=torch.device(device)),
+            num_frames=num_frames,
+            vocab_size=vocab_size,
+        )
+
+        etalon_schema_fst: List[List[int]] = []
+        for frame_i in range(num_frames):
+            for label_i in range(text_length + 1):
+                state = frame_i * (text_length + 1) + label_i
+                if label_i < text_length:
+                    next_state_label = state + 1
+                    # next unit
+                    etalon_schema_fst.append([state, next_state_label, labels[label_i], frame_i, label_i, 0])
+                if frame_i < num_frames - 1:
+                    next_state_frame = (frame_i + 1) * (text_length + 1) + label_i
+                    # next time frame (blank)
+                    etalon_schema_fst.append([state, next_state_frame, blank_id, frame_i, label_i, 0])
+
+        # start eps-transition
+        for frame_i in range(1, num_frames):
+            etalon_schema_fst.append([0, frame_i * (text_length + 1), vocab_size, 0, 0, 0])
+
+        last_grid_state = num_frames * (text_length + 1) - 1
+
+        # end eps-transitions
+        if last_blank_mode == "force_final":
+            last_eps_state = last_grid_state
+        else:
+            assert last_blank_mode == "allow_ignore"
+            last_eps_state = last_grid_state + 1
+
+        for frame_i in range(num_frames - 1):
+            etalon_schema_fst.append(
+                [(frame_i + 1) * (text_length + 1) - 1, last_eps_state, vocab_size + 1, frame_i, text_length, 0]
+            )
+
+        etalon_schema_fst.append([last_grid_state, last_grid_state + 1, blank_id, num_frames - 1, text_length, 0])
+        etalon_schema_fst.append(
+            [last_grid_state + 1, last_grid_state + 2, -1, -1, -1, 0]
+        )  # transition to final state
+        etalon_schema_fst.append([last_grid_state + 2])  # final state
+        etalon_schema_fst = sorted(etalon_schema_fst)  # required for k2.Fsa.from_str
+        etalon_schema_fst_str = "\n".join([" ".join(map(str, line)) for line in etalon_schema_fst])
+        etalon_grid_schema = k2.Fsa.from_str(etalon_schema_fst_str, aux_label_names=["aux_labels", "unit_positions"])
+
+        assert grid_schema.num_arcs == etalon_grid_schema.num_arcs
+        assert grid_schema.shape == etalon_grid_schema.shape  # (num_states, None)
+        assert k2.is_rand_equivalent(
+            grid_schema, etalon_grid_schema, log_semiring=True, treat_epsilons_specially=False
+        ), "Grid schema input labels mismatch"
+        assert k2.is_rand_equivalent(
+            grid_schema.invert(), etalon_grid_schema.invert(), log_semiring=True, treat_epsilons_specially=False
+        ), "Grid schema output labels mismatch"
+
+        # swap aux_labels and unit positions to test unit_positions
+        grid_schema.aux_labels, grid_schema.unit_positions = grid_schema.unit_positions, grid_schema.aux_labels
+        etalon_grid_schema.aux_labels, etalon_grid_schema.unit_positions = (
+            etalon_grid_schema.unit_positions,
+            etalon_grid_schema.aux_labels,
+        )
+        assert k2.is_rand_equivalent(
+            grid_schema.invert(), etalon_grid_schema.invert(), log_semiring=True, treat_epsilons_specially=False
+        ), "Grid schema unit positions mismatch"
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("blank_first", [True, False])
+    @pytest.mark.parametrize("last_blank_mode", ["allow_ignore", "force_final"])
+    def test_small_random_grid_compose_equivalent(
+        self, device: torch.device, blank_first: bool, last_blank_mode, rnn_loss_sample_data
+    ):
+        sample_data = rnn_loss_sample_data.get_sample_small_random(blank_first, device=device)
+        criterion = GraphWTransducerLoss(
+            blank=sample_data.blank_id,
+            last_blank_mode=last_blank_mode,
+            connect_composed=True,
+            use_grid_implementation=False,
+        )
+        text_tensor = sample_data.targets[0]
+        num_frames = sample_data.logits.shape[1]
+        graph_grid = criterion.get_grid(text_tensor, num_frames, sample_data.vocab_size)
+        graph_composed = criterion.get_composed_lattice(text_tensor, num_frames, sample_data.vocab_size)
+        assert k2.is_rand_equivalent(
+            graph_grid, graph_composed, log_semiring=True, treat_epsilons_specially=False
+        ), "Grid and composed graphs are not equivalent."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("last_blank_mode", ["allow_ignore", "force_final"])
+    @pytest.mark.parametrize("use_grid_implementation", [True, False])
+    def test_small_grid_transducer_inf_penalty(
+        self, device, last_blank_mode, use_grid_implementation, rnnt_test_helper, rnn_loss_sample_data
+    ):
+        """
+        With -inf eps penalty W-Transducer loss should be equivalent to RNN-T loss.
+        """
+        sample_data = rnn_loss_sample_data.get_sample_small()
+        graph_rnnt = GraphWTransducerLoss(
+            blank=0,
+            eps_weight=-100.0,
+            last_blank_mode=last_blank_mode,
+            use_grid_implementation=use_grid_implementation,
+        )
+        graph_cost, graph_grads = rnnt_test_helper.wrap_and_call(
+            graph_rnnt, sample_data.logits, sample_data.targets, device
+        )
+        assert np.allclose(graph_cost, sample_data.expected_cost.numpy(), rtol=1e-6), "costs mismatch."
+        assert np.allclose(graph_grads, sample_data.expected_grads.numpy(), atol=1e-6), "gradient mismatch."

tests/collections/asr/mixins/adapters/test_asr_adapter_mixin.py

@@ -0,0 +1,763 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import pytest
+import torch
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+from nemo.collections.asr.models import ASRModel, EncDecCTCModel, EncDecMultiTaskModel, EncDecRNNTModel
+from nemo.collections.asr.parts.submodules.adapters import (
+    multi_head_attention_adapter_module,
+    transformer_multi_head_attention_adapter_module,
+)
+from nemo.collections.asr.parts.utils import adapter_utils
+from nemo.collections.common.parts import adapter_modules
+from nemo.core.classes.mixins.access_mixins import AccessMixin
+from nemo.core.classes.mixins.adapter_mixins import AdapterModuleMixin, get_registered_adapter
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+from nemo.utils import model_utils
+
+NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cpu_is_supported(
+    __NUMBA_MINIMUM_VERSION__
+) or numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+
+
+@pytest.fixture()
+def model():
+    preprocessor = {'_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor'}
+    encoder = {
+        '_target_': 'nemo.collections.asr.modules.ConvASREncoderAdapter',
+        'feat_in': 64,
+        'activation': 'relu',
+        'conv_mask': True,
+        'jasper': [
+            {
+                'filters': 50,
+                'repeat': 1,
+                'kernel': [1],
+                'stride': [1],
+                'dilation': [1],
+                'dropout': 0.0,
+                'residual': False,
+                'separable': True,
+                'se': True,
+                'se_context_size': -1,
+            }
+        ],
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+        'feat_in': 50,
+        'num_classes': 28,
+        'vocabulary': [
+            ' ',
+            'a',
+            'b',
+            'c',
+            'd',
+            'e',
+            'f',
+            'g',
+            'h',
+            'i',
+            'j',
+            'k',
+            'l',
+            'm',
+            'n',
+            'o',
+            'p',
+            'q',
+            'r',
+            's',
+            't',
+            'u',
+            'v',
+            'w',
+            'x',
+            'y',
+            'z',
+            "'",
+        ],
+    }
+    modelConfig = DictConfig(
+        {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
+    )
+
+    model_instance = EncDecCTCModel(cfg=modelConfig)
+    return model_instance
+
+
+@pytest.fixture()
+def conformer_ctc_adapter():
+    preprocessor = {'_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor'}
+    encoder = {
+        '_target_': 'nemo.collections.asr.modules.ConformerEncoderAdapter',
+        'feat_in': 64,
+        'feat_out': -1,
+        'n_layers': 2,
+        'd_model': 128,
+        'subsampling': 'striding',
+        'subsampling_factor': 4,
+        'self_attention_model': 'rel_pos',
+        'n_heads': 4,
+        'conv_kernel_size': 31,
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+        'feat_in': 128,
+        'num_classes': 28,
+        'vocabulary': [
+            ' ',
+            'a',
+            'b',
+            'c',
+            'd',
+            'e',
+            'f',
+            'g',
+            'h',
+            'i',
+            'j',
+            'k',
+            'l',
+            'm',
+            'n',
+            'o',
+            'p',
+            'q',
+            'r',
+            's',
+            't',
+            'u',
+            'v',
+            'w',
+            'x',
+            'y',
+            'z',
+            "'",
+        ],
+    }
+    modelConfig = DictConfig(
+        {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
+    )
+
+    model_instance = EncDecCTCModel(cfg=modelConfig)
+    return model_instance
+
+
+@pytest.fixture()
+def rnnt_model():
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+
+    # fmt: off
+    labels = [' ', 'a', 'b', 'c', 'd', 'e', 'f',
+              'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o',
+              'p', 'q', 'r', 's', 't', 'u', 'v', 'w',
+              'x', 'y', 'z', "'",
+              ]
+    # fmt: on
+
+    model_defaults = {'enc_hidden': 96, 'pred_hidden': 64}
+
+    # Test case where Encoder (default) is not adapter compatible
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': model_defaults['enc_hidden'],
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                }
+            ],
+        },
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.RNNTDecoder',
+        'prednet': {'pred_hidden': model_defaults['pred_hidden'], 'pred_rnn_layers': 1},
+    }
+
+    joint = {
+        '_target_': 'nemo.collections.asr.modules.RNNTJoint',
+        'jointnet': {'joint_hidden': 32, 'activation': 'relu'},
+    }
+
+    decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 10}}
+
+    loss = {'loss_name': 'default', 'warprnnt_numba_kwargs': {'fastemit_lambda': 0.001}}
+
+    modelConfig = DictConfig(
+        {
+            'labels': ListConfig(labels),
+            'preprocessor': DictConfig(preprocessor),
+            'model_defaults': DictConfig(model_defaults),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'joint': DictConfig(joint),
+            'decoding': DictConfig(decoding),
+            'loss': DictConfig(loss),
+        }
+    )
+
+    model_instance = EncDecRNNTModel(cfg=modelConfig)
+    return model_instance
+
+
+@pytest.fixture()
+def multitask_model(test_data_dir):
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+
+    # fmt: off
+    tokenizer = {
+        'dir': None,
+        'type': 'agg',
+        'langs': {
+            'spl_tokens': {
+                'dir': os.path.join(test_data_dir, 'asr', 'tokenizers', 'canary'),
+                'type': 'bpe',
+            },
+            'en': {
+                'dir': os.path.join(test_data_dir, 'asr', 'tokenizers', 'an4_spe_128'),
+                'type': 'bpe',
+            }
+        },
+        'custom_tokenizer': {
+            '_target_': 'nemo.collections.common.tokenizers.canary_tokenizer.CanaryTokenizer',
+            'tokenizers': None,
+        }
+    }
+    # fmt: on
+
+    model_defaults = {"asr_enc_hidden": 128, "lm_enc_hidden": 128, "lm_dec_hidden": 128}
+
+    # Test case where Encoder (default) is not adapter compatible
+    encoder = {
+        '_target_': 'nemo.collections.asr.modules.ConformerEncoder',
+        'feat_in': 64,
+        'feat_out': -1,
+        'n_layers': 2,
+        'd_model': 128,
+        'subsampling': 'striding',
+        'subsampling_factor': 4,
+        'self_attention_model': 'rel_pos',
+        'n_heads': 4,
+        'conv_kernel_size': 31,
+    }
+
+    transf_encoder = {
+        "_target_": "nemo.collections.asr.modules.transformer.transformer_encoders.TransformerEncoder",
+        "num_layers": 1,
+        "hidden_size": "${model_defaults.lm_enc_hidden}",
+        "inner_size": int(4 * model_defaults['lm_enc_hidden']),
+        "num_attention_heads": 8,
+        "ffn_dropout": 0.1,
+        "attn_score_dropout": 0.1,
+        "attn_layer_dropout": 0.1,
+        "mask_future": False,
+        "pre_ln": True,
+        "pre_ln_final_layer_norm": True,
+    }
+
+    transf_decoder = {
+        "_target_": "nemo.collections.asr.modules.transformer.get_nemo_transformer",
+        "model_name": None,
+        "pretrained": False,
+        "encoder": None,
+        "pre_ln_final_layer_norm": True,
+        "config_dict": {
+            "max_sequence_length": 512,
+            "num_token_types": 0,
+            "embedding_dropout": 0.1,
+            "learn_positional_encodings": False,
+            "hidden_size": "${model_defaults.lm_dec_hidden}",
+            "inner_size": "${multiply:${model_defaults.lm_dec_hidden}, 4}",
+            "num_layers": 2,
+            "num_attention_heads": 8,
+            "ffn_dropout": 0.1,
+            "attn_score_dropout": 0.1,
+            "attn_layer_dropout": 0.1,
+            "hidden_act": "relu",
+            "pre_ln": True,
+            "vocab_size": None,  # Will be set by the model at runtime
+            "adapter": True,  # Add support for adapter class
+        },
+    }
+
+    head = {
+        "_target_": "nemo.collections.asr.parts.submodules.token_classifier.TokenClassifier",
+        "num_layers": 1,
+        "activation": "relu",
+        "log_softmax": True,
+        "hidden_size": "${transf_decoder.config_dict.hidden_size}",
+        "num_classes": None,  # Will be set by the model at runtime
+        "dropout": 0.0,
+        "use_transformer_init": True,
+    }
+
+    decoding = {'strategy': 'beam', 'beam': {'beam_size': 1, 'len_pen': 0.0, 'max_generation_delta': 50}}
+
+    loss = {
+        "_target_": "nemo.collections.common.losses.smoothed_cross_entropy.SmoothedCrossEntropyLoss",
+        "label_smoothing": 0.0,
+        "pad_id": None,
+    }
+
+    modelConfig = DictConfig(
+        {
+            'sample_rate': 16000,
+            'prompt_format': 'canary',
+            'preprocessor': DictConfig(preprocessor),
+            'model_defaults': DictConfig(model_defaults),
+            'tokenizer': DictConfig(tokenizer),
+            'encoder': DictConfig(encoder),
+            'transf_encoder': DictConfig(transf_encoder),
+            'transf_decoder': DictConfig(transf_decoder),
+            'head': DictConfig(head),
+            'decoding': DictConfig(decoding),
+            'loss': DictConfig(loss),
+        }
+    )
+
+    model_instance = EncDecMultiTaskModel(cfg=modelConfig)
+
+    # Execute the model class swap logic
+    model_instance.replace_adapter_compatible_modules()
+    return model_instance
+
+
+def get_adapter_cfg(in_features=50, dim=100, norm_pos='pre', atype='linear', **kwargs):
+    valid_types = ['linear', 'mha', 'relmha', 'transf_mha']
+    if atype not in valid_types:
+        raise ValueError(f"Invalid type. Valid types = {atype}")
+
+    if atype == 'linear':
+        cfg = adapter_modules.LinearAdapterConfig(in_features=in_features, dim=dim, norm_position=norm_pos)
+    elif atype == 'mha':
+        cfg = multi_head_attention_adapter_module.MultiHeadAttentionAdapterConfig(
+            n_head=kwargs.get('n_head', 1),
+            n_feat=in_features,
+            proj_dim=kwargs.get('proj_dim', None),
+        )
+    elif atype == 'transf_mha':
+        cfg = transformer_multi_head_attention_adapter_module.TransformerMultiHeadAttentionAdapterConfig(
+            num_attention_heads=kwargs.get('n_head', 1),
+            hidden_size=in_features,
+            proj_dim=kwargs.get('proj_dim', None),
+        )
+    else:  # atype == 'relmha'
+        cfg = multi_head_attention_adapter_module.RelPositionMultiHeadAttentionAdapterConfig(
+            n_head=kwargs.get('n_head', 1), n_feat=in_features
+        )
+
+    print(cfg._target_)
+
+    cfg = OmegaConf.structured(cfg)
+    return cfg
+
+
+class TestASRAdapterMixin:
+    @pytest.mark.unit
+    def test_class_paths_are_correct(self):
+        # Resolve all object names in module
+        obj_keys = list(dir(adapter_utils))
+        for key in obj_keys:
+            if 'CLASSPATH' in key:
+                classpath = getattr(adapter_utils, key)
+                # This will raise import error if it fails
+                _ = model_utils.import_class_by_path(classpath)
+
+                # Try getting thmulti_head_attention_adapter_module.pye config of the class
+                config_path = classpath + "Config"
+                _ = model_utils.import_class_by_path(config_path)
+
+    @pytest.mark.unit
+    def test_asr_model_constructor(self, model):
+        original_num_params = model.num_weights
+
+        model.add_adapter(name='adapter_0', cfg=get_adapter_cfg())
+        new_num_params = model.num_weights
+        assert new_num_params > original_num_params
+
+    @pytest.mark.unit
+    def test_asr_model_constructor_mha_adapter(self, model):
+        with pytest.raises(ValueError):
+            model.add_adapter(name='adapter_0', cfg=get_adapter_cfg(atype='mha'))
+
+    @pytest.mark.unit
+    def test_conformer_constructor_mha_adapter(self, conformer_ctc_adapter):
+        original_num_params = conformer_ctc_adapter.num_weights
+
+        conformer_ctc_adapter.add_adapter(name='adapter_0', cfg=get_adapter_cfg(atype='relmha'))
+        new_num_params = conformer_ctc_adapter.num_weights
+        assert new_num_params > original_num_params
+
+    @pytest.mark.unit
+    def test_asr_model_constructor_encoder_module(self, model):
+        original_num_params = model.num_weights
+
+        model.add_adapter(name='encoder:adapter_0', cfg=get_adapter_cfg())
+        new_num_params = model.num_weights
+        assert new_num_params > original_num_params
+
+    @pytest.mark.unit
+    def test_asr_model_constructor_decoder_module(self, model):
+        original_num_params = model.num_weights
+
+        model.add_adapter(name='decoder:adapter_0', cfg=get_adapter_cfg())
+        new_num_params = model.num_weights
+        assert new_num_params > original_num_params
+        assert model.decoder.is_adapter_available()
+        assert model.decoder.get_enabled_adapters()[0] == 'adapter_0'
+
+    @pytest.mark.unit
+    def test_asr_model_constructor_joint_module_ctc_skip(self, model):
+        original_num_params = model.num_weights
+
+        # this step should exit without adding adapters and without errors
+        with pytest.raises(ValueError):
+            model.add_adapter(name='joint:adapter_0', cfg=get_adapter_cfg())
+        new_num_params = model.num_weights
+        assert new_num_params == original_num_params
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_asr_model_constructor_joint_module_rnnt(self, rnnt_model):
+        original_num_params = rnnt_model.num_weights
+
+        rnnt_model.add_adapter(name='joint:adapter_0', cfg=get_adapter_cfg())
+        new_num_params = rnnt_model.num_weights
+        assert new_num_params > original_num_params
+        assert rnnt_model.joint.is_adapter_available()
+        assert rnnt_model.joint.get_enabled_adapters()[0] == 'adapter_0'
+
+    @pytest.mark.unit
+    def test_asr_multiple_adapter(self, model):
+        original_num_params = model.num_weights
+
+        model.add_adapter(name='adapter_0', cfg=get_adapter_cfg())
+        new_num_params = model.num_weights
+        assert new_num_params > original_num_params
+
+        original_num_params = new_num_params
+        model.add_adapter(name='adapter_1', cfg=get_adapter_cfg())
+        new_num_params = model.num_weights
+        assert new_num_params > original_num_params
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('name', ['adapter_0', 'encoder:adapter_0', 'decoder:adapter_0'])
+    def test_asr_forward_linear_pre(self, model, name):
+        model.eval()
+        torch.random.manual_seed(0)
+        input_signal = torch.randn(2, 512)
+        input_signal_length = torch.tensor([512, 512], dtype=torch.int32)
+
+        origial_output = model(input_signal=input_signal, input_signal_length=input_signal_length)[0]
+
+        model.add_adapter(name=name, cfg=get_adapter_cfg())
+        new_output = model(input_signal=input_signal, input_signal_length=input_signal_length)[0]
+
+        assert torch.mean(torch.abs(origial_output - new_output)) < 1e-5
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('name', ['adapter_0', 'encoder:adapter_0', 'decoder:adapter_0'])
+    def test_asr_forward_linear_post(self, model, name):
+        model.eval()
+        torch.random.manual_seed(0)
+        input_signal = torch.randn(2, 512)
+        input_signal_length = torch.tensor([512, 512], dtype=torch.int32)
+
+        origial_output = model(input_signal=input_signal, input_signal_length=input_signal_length)[0]
+
+        model.add_adapter(name=name, cfg=get_adapter_cfg(norm_pos='post'))
+        new_output = model(input_signal=input_signal, input_signal_length=input_signal_length)[0]
+
+        assert torch.mean(torch.abs(origial_output - new_output)) < 1e-5
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('name', ['adapter_0', 'encoder:adapter_0'])
+    def test_conformer_forward_mha(self, conformer_ctc_adapter, name):
+        conformer_ctc_adapter.eval()
+        torch.random.manual_seed(0)
+        input_signal = torch.randn(2, 512)
+        input_signal_length = torch.tensor([512, 512], dtype=torch.int32)
+
+        origial_output = conformer_ctc_adapter(input_signal=input_signal, input_signal_length=input_signal_length)[0]
+
+        conformer_ctc_adapter.add_adapter(name=name, cfg=get_adapter_cfg(in_features=128, atype='mha'))
+        new_output = conformer_ctc_adapter(input_signal=input_signal, input_signal_length=input_signal_length)[0]
+
+        assert torch.mean(torch.abs(origial_output - new_output)) < 1e-5
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('adapter_type', ['linear', 'attn'])
+    @pytest.mark.parametrize(
+        'name', ['adapter_0', 'encoder:adapter_0', 'transf_encoder:adapter_0', 'transf_decoder:adapter_0']
+    )
+    def test_canary_forward_mha(self, multitask_model, name, adapter_type):
+        multitask_model.eval()
+        torch.random.manual_seed(0)
+        input_signal = torch.randn(2, 512)
+        input_signal_length = torch.tensor([512, 512], dtype=torch.int32)
+        transcript = torch.randint(0, multitask_model.tokenizer.vocab_size, size=(2, 10))
+        transcript_len = torch.tensor([10, 9], dtype=torch.int32)
+
+        origial_output = multitask_model(
+            input_signal=input_signal,
+            input_signal_length=input_signal_length,
+            transcript=transcript,
+            transcript_length=transcript_len,
+        )
+        og_logprob = origial_output[0]
+        og_enc_out = origial_output[2]
+
+        if adapter_type == 'attn':
+            adapter_type = 'transf_mha' if 'transf' in name else 'mha'
+
+        multitask_model.add_adapter(name=name, cfg=get_adapter_cfg(in_features=128, atype=adapter_type, proj_dim=4))
+
+        new_output = multitask_model(
+            input_signal=input_signal,
+            input_signal_length=input_signal_length,
+            transcript=transcript,
+            transcript_length=transcript_len,
+        )
+
+        new_logprob = new_output[0]
+        new_enc_out = new_output[2]
+
+        assert torch.mean(torch.abs(og_logprob - new_logprob)) < 1e-5
+        assert torch.mean(torch.abs(og_enc_out - new_enc_out)) < 1e-5
+
+        if 'linear' in adapter_type:
+            mod_name = name.split(":")[-1]
+            for mod in multitask_model.modules():
+                if isinstance(mod, AdapterModuleMixin):
+                    amodule = mod.get_adapter_module(mod_name)
+                    if amodule is not None:
+                        assert isinstance(amodule, adapter_modules.LinearAdapter)
+
+        # Try to use incorrect adapter
+        with pytest.raises(ValueError):
+            multitask_model.add_adapter(
+                name="transf_encoder:adapter_1", cfg=get_adapter_cfg(in_features=128, atype='mha')
+            )
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('name', ['transf_decoder:adapter_0'])
+    def test_canary_forward_mha_decoder_fails_without_support(self, multitask_model, name):
+        multitask_model.eval()
+        torch.random.manual_seed(0)
+
+        # Change internal class of transf_decoder module
+        adapter_class = multitask_model.transf_decoder.__class__
+        multitask_model.transf_decoder.__class__ = get_registered_adapter(adapter_class).base_class
+
+        with pytest.raises(AttributeError):
+            adapter_type = 'transf_mha' if 'transf' in name else 'mha'
+            multitask_model.add_adapter(name=name, cfg=get_adapter_cfg(in_features=128, atype=adapter_type))
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('name1', ['adapter_0', 'encoder:adapter_0', 'decoder:adapter_0'])
+    @pytest.mark.parametrize('name2', ['adapter_1', 'encoder:adapter_1', 'decoder:adapter_1'])
+    def test_asr_multi_adapter_forward(self, model, name1, name2):
+        model.eval()
+        torch.random.manual_seed(0)
+        input_signal = torch.randn(2, 512)
+        input_signal_length = torch.tensor([512, 512], dtype=torch.int32)
+
+        origial_output = model(input_signal=input_signal, input_signal_length=input_signal_length)[0]
+
+        model.add_adapter(name=name1, cfg=get_adapter_cfg())
+        model.add_adapter(name=name2, cfg=get_adapter_cfg())
+        new_output = model(input_signal=input_signal, input_signal_length=input_signal_length)[0]
+
+        resolved_name1 = model.resolve_adapter_module_name_(name1)[-1]
+        resolved_name2 = model.resolve_adapter_module_name_(name2)[-1]
+        assert model.get_enabled_adapters() == [resolved_name1, resolved_name2]
+        assert torch.mean(torch.abs(origial_output - new_output)) < 1e-5
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.parametrize('name1', ['decoder:adapter_0', 'joint:adapter_0'])
+    @pytest.mark.parametrize('name2', ['decoder:adapter_1', 'joint:adapter_1'])
+    @pytest.mark.unit
+    def test_asr_multi_adapter_forward(self, rnnt_model, name1, name2):
+        rnnt_model.eval()
+        torch.random.manual_seed(0)
+        input_signal = torch.randn(2, 512)
+        input_signal_length = torch.tensor([512, 512], dtype=torch.int32)
+
+        origial_output = rnnt_model(input_signal=input_signal, input_signal_length=input_signal_length)[0]
+
+        rnnt_model.add_adapter(name=name1, cfg=get_adapter_cfg())
+        rnnt_model.add_adapter(name=name2, cfg=get_adapter_cfg())
+        new_output = rnnt_model(input_signal=input_signal, input_signal_length=input_signal_length)[0]
+
+        resolved_name1 = rnnt_model.resolve_adapter_module_name_(name1)[-1]
+        resolved_name2 = rnnt_model.resolve_adapter_module_name_(name2)[-1]
+        assert rnnt_model.get_enabled_adapters() == [resolved_name1, resolved_name2]
+        assert torch.mean(torch.abs(origial_output - new_output)) < 1e-5
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('name1', ['adapter_0', 'encoder:adapter_0', 'decoder:adapter_0'])
+    @pytest.mark.parametrize('name2', ['adapter_1', 'encoder:adapter_1', 'decoder:adapter_1'])
+    def test_asr_multi_adapter_partial_forward(self, model, name1, name2):
+        model.eval()
+        torch.random.manual_seed(0)
+        input_signal = torch.randn(2, 512)
+        input_signal_length = torch.tensor([512, 512], dtype=torch.int32)
+
+        origial_output = model(input_signal=input_signal, input_signal_length=input_signal_length)[0]
+
+        model.add_adapter(name=name1, cfg=get_adapter_cfg())
+        model.add_adapter(name=name2, cfg=get_adapter_cfg())
+
+        model.set_enabled_adapters(name=name1, enabled=False)
+        new_output = model(input_signal=input_signal, input_signal_length=input_signal_length)[0]
+
+        resolved_name2 = model.resolve_adapter_module_name_(name2)[-1]
+        assert model.get_enabled_adapters() == [resolved_name2]
+        assert torch.mean(torch.abs(origial_output - new_output)) < 1e-5
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('name', ['adapter_0', 'encoder:adapter_0', 'decoder:adapter_0'])
+    def test_asr_forward_unfrozen_adapters(self, model, name):
+        model.eval()
+        original_num_params = model.num_weights
+
+        dim = 10
+        model.add_adapter(name=name, cfg=get_adapter_cfg(dim=dim))
+        model.freeze()
+        model.unfreeze_enabled_adapters()
+
+        assert original_num_params == 5443
+
+        original_params = 0
+        adapter_params = 0
+        for name, param in model.named_parameters():
+            if 'adapter' not in name:
+                assert param.requires_grad is False
+                original_params += param.numel()
+            else:
+                assert param.requires_grad is True
+                adapter_params += param.numel()
+
+        for mname, module in model.named_modules():
+            if isinstance(module, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d, torch.nn.BatchNorm3d)):
+                assert module.track_running_stats is False
+
+        assert original_params > adapter_params
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_constructor_pretrained(self):
+        # Check to/from config_dict:
+        cfg = ASRModel.from_pretrained('stt_en_citrinet_256', map_location='cpu', return_config=True)
+        adapter_metadata = get_registered_adapter(cfg.encoder._target_)
+        if adapter_metadata is not None:
+            cfg.encoder._target_ = adapter_metadata.adapter_class_path
+        model = ASRModel.from_pretrained('stt_en_citrinet_256', override_config_path=cfg)
+
+        assert isinstance(model, AdapterModuleMixin)
+        assert hasattr(model, 'encoder')
+        assert isinstance(model.encoder, AdapterModuleMixin)
+
+        model.add_adapter('adapter_0', cfg=get_adapter_cfg(in_features=cfg.encoder.jasper[0].filters, dim=5))
+        assert model.is_adapter_available()
+
+        model.freeze()
+        model.unfreeze_enabled_adapters()
+        assert model.num_weights < 1e5
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_constructor_pretrained_rnnt(self):
+        # Check to/from config_dict:
+        cfg = ASRModel.from_pretrained('stt_en_fastconformer_transducer_large', map_location='cpu', return_config=True)
+        adapter_metadata = get_registered_adapter(cfg.encoder._target_)
+        if adapter_metadata is not None:
+            cfg.encoder._target_ = adapter_metadata.adapter_class_path
+        model = ASRModel.from_pretrained('stt_en_fastconformer_transducer_large', override_config_path=cfg)
+
+        assert isinstance(model, AdapterModuleMixin)
+        assert hasattr(model, 'encoder')
+        assert isinstance(model.encoder, AdapterModuleMixin)
+        assert hasattr(model, 'decoder')
+        assert isinstance(model.decoder, AdapterModuleMixin)
+        assert hasattr(model, 'joint')
+        assert isinstance(model.joint, AdapterModuleMixin)
+
+        model.add_adapter('adapter_0', cfg=get_adapter_cfg(in_features=cfg.encoder.d_model, dim=5))
+        model.add_adapter('decoder:adapter_1', cfg=get_adapter_cfg(in_features=cfg.decoder.prednet.pred_hidden, dim=5))
+        model.add_adapter('joint:adapter_2', cfg=get_adapter_cfg(in_features=cfg.joint.jointnet.joint_hidden, dim=5))
+        assert model.is_adapter_available()
+
+        model.freeze()
+        model.unfreeze_enabled_adapters()
+        assert model.num_weights < 2e5
+
+    @pytest.mark.unit
+    def test_asr_model_adapter_loss(self, model):
+        original_num_params = model.num_weights
+        x = torch.randn(2, 512)
+        x_len = torch.tensor([400, 512], dtype=torch.int32)
+
+        adapter_cfg = get_adapter_cfg()  # type: adapter_modules.LinearAdapterConfig
+        adapter_cfg.adapter_strategy.l2_lambda = 0.01
+
+        model.add_adapter(name='adapter_0', cfg=adapter_cfg)
+        new_num_params = model.num_weights
+        assert new_num_params > original_num_params
+
+        model.train()  # set training mode to true
+
+        with torch.no_grad():
+            AccessMixin.reset_registry(model)
+            AccessMixin.update_access_cfg({'save_encoder_tensors': False}, model.model_guid)
+            _ = model(input_signal=x, input_signal_length=x_len)
+
+            # extract losses
+            auxiliary_losses = AccessMixin.get_module_registry(model)
+
+            loss = list(auxiliary_losses.values())[0]
+            assert 'adapter_loss' in loss
+            assert loss['adapter_loss'][0] == torch.tensor(0.0)  # initially adapter is 0 init, no loss required.
+
+            AccessMixin.reset_registry(model)

tests/collections/asr/mixins/adapters/test_asr_adapter_modules.py

@@ -0,0 +1,331 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.collections.asr.parts.submodules import adapters as adapter_modules
+from nemo.core.classes.mixins import adapter_mixin_strategies
+from nemo.utils import config_utils
+
+
+def _create_masks(att_mask, max_audio_length, padding_length):
+    # pad_mask is the masking to be used to ignore paddings
+    pad_mask = torch.arange(0, max_audio_length).expand(padding_length.size(0), -1) < padding_length.unsqueeze(-1)
+
+    # pad_mask_for_att_mask is the mask which helps to ignore paddings
+    pad_mask_for_att_mask = pad_mask.unsqueeze(1).repeat([1, max_audio_length, 1])
+    pad_mask_for_att_mask = torch.logical_and(pad_mask_for_att_mask, pad_mask_for_att_mask.transpose(1, 2))
+    # att_mask is the masking to be used by the MHA layers to ignore the tokens not supposed to be visible
+    att_mask = att_mask[:, :max_audio_length, :max_audio_length]
+    # paddings should also get ignored, so pad_mask_for_att_mask is used to ignore their corresponding scores
+    att_mask = torch.logical_and(pad_mask_for_att_mask, att_mask.to(pad_mask_for_att_mask.device))
+
+    pad_mask = ~pad_mask
+    att_mask = ~att_mask
+
+    return pad_mask, att_mask
+
+
+def get_mask(lengths: torch.Tensor):
+    max_seq_len = lengths.max()
+    att_mask = torch.ones(1, max_seq_len, max_seq_len, dtype=torch.bool)
+
+    pad_mask, att_mask = _create_masks(att_mask, max_seq_len, lengths)
+    return pad_mask, att_mask
+
+
+class TestASRAdapterModules:
+    @pytest.mark.unit
+    def test_mha_adapter_config(self):
+        IGNORED_ARGS = ['_target_']
+
+        result = config_utils.assert_dataclass_signature_match(
+            adapter_modules.MultiHeadAttentionAdapter,
+            adapter_modules.MultiHeadAttentionAdapterConfig,
+            ignore_args=IGNORED_ARGS,
+        )
+
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_relpos_mha_adapter_config(self):
+        IGNORED_ARGS = ['_target_']
+
+        result = config_utils.assert_dataclass_signature_match(
+            adapter_modules.RelPositionMultiHeadAttentionAdapter,
+            adapter_modules.RelPositionMultiHeadAttentionAdapterConfig,
+            ignore_args=IGNORED_ARGS,
+        )
+
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_abs_pos_encoding_adapter_config(self):
+        IGNORED_ARGS = ['_target_']
+
+        result = config_utils.assert_dataclass_signature_match(
+            adapter_modules.PositionalEncodingAdapter,
+            adapter_modules.PositionalEncodingAdapterConfig,
+            ignore_args=IGNORED_ARGS,
+        )
+
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_rel_pos_encoding_adapter_config(self):
+        IGNORED_ARGS = ['_target_']
+
+        result = config_utils.assert_dataclass_signature_match(
+            adapter_modules.RelPositionalEncodingAdapter,
+            adapter_modules.RelPositionalEncodingAdapterConfig,
+            ignore_args=IGNORED_ARGS,
+        )
+
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_transformer_mha_adapter_config(self):
+        IGNORED_ARGS = ['_target_']
+
+        result = config_utils.assert_dataclass_signature_match(
+            adapter_modules.TransformerMultiHeadAttentionAdapter,
+            adapter_modules.TransformerMultiHeadAttentionAdapterConfig,
+            ignore_args=IGNORED_ARGS,
+        )
+
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('n_head', [1, 2, 10])
+    @pytest.mark.parametrize('proj_dim', [None, -1])
+    def test_mha_adapter_init(self, n_head, proj_dim):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 32, 50)
+        lengths = torch.randint(1, x.size(1), size=(x.size(0),))
+        lengths[torch.randint(0, x.size(0), size=(1,))[0]] = x.size(1)
+
+        adapter = adapter_modules.MultiHeadAttentionAdapter(
+            n_head=n_head, n_feat=50, dropout_rate=0.0, proj_dim=proj_dim
+        )
+
+        pad_mask, att_mask = get_mask(lengths)
+
+        with torch.no_grad():
+            assert adapter.linear_out.weight.sum() == 0
+            if hasattr(adapter.linear_out, 'bias') and adapter.linear_out.bias is not None:
+                assert adapter.linear_out.bias.sum() == 0
+
+            out = adapter(x, x, x, att_mask)
+            assert out.sum().abs() <= 1e-8
+            assert out.shape == x.shape
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('n_head', [1, 2, 10])
+    @pytest.mark.parametrize('proj_dim', [None, -1])
+    def test_relmha_adapter_init(self, n_head, proj_dim):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 32, 50)
+        lengths = torch.randint(1, x.size(1), size=(x.size(0),))
+        lengths[torch.randint(0, x.size(0), size=(1,))[0]] = x.size(1)
+
+        adapter = adapter_modules.RelPositionMultiHeadAttentionAdapter(
+            n_head=n_head, n_feat=50, dropout_rate=0.0, proj_dim=proj_dim
+        )
+        relpos_enc = adapter_modules.RelPositionalEncodingAdapter(d_model=50)
+
+        pad_mask, att_mask = get_mask(lengths)
+        relpos_enc.extend_pe(lengths.max(), device='cpu', dtype=torch.float32)
+
+        with torch.no_grad():
+            assert adapter.linear_out.weight.sum() == 0
+            if hasattr(adapter.linear_out, 'bias') and adapter.linear_out.bias is not None:
+                assert adapter.linear_out.bias.sum() == 0
+
+            _, pos_emb = relpos_enc(x)
+            out = adapter(x, x, x, att_mask, pos_emb)
+            assert out.sum().abs() <= 1e-8
+            assert out.shape == x.shape
+
+    @pytest.mark.unit
+    def test_relmha_adapter_with_torch_sdpa(self):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 32, 50)
+        lengths = torch.randint(1, x.size(1), size=(x.size(0),))
+        lengths[torch.randint(0, x.size(0), size=(1,))[0]] = x.size(1)
+
+        adapter_torch_sdpa = adapter_modules.RelPositionMultiHeadAttentionAdapter(
+            n_head=2, n_feat=50, dropout_rate=0.0, proj_dim=-1, use_pytorch_sdpa=True
+        )
+        adapter = adapter_modules.RelPositionMultiHeadAttentionAdapter(
+            n_head=2, n_feat=50, dropout_rate=0.0, proj_dim=-1, use_pytorch_sdpa=False
+        )
+        # to dont reset linear_out parameters to zero
+        adapter.linear_out = torch.nn.Linear(adapter.linear_out.in_features, adapter.linear_out.out_features)
+        for original_param, sdpa_param in zip(adapter.parameters(), adapter_torch_sdpa.parameters()):
+            sdpa_param.data.copy_(original_param.data)
+        relpos_enc = adapter_modules.RelPositionalEncodingAdapter(d_model=50)
+
+        pad_mask, att_mask = get_mask(lengths)
+        relpos_enc.extend_pe(lengths.max(), device='cpu', dtype=torch.float32)
+
+        with torch.no_grad():
+            _, pos_emb = relpos_enc(x)
+            out = adapter(x, x, x, att_mask, pos_emb)
+            out_sdpa = adapter_torch_sdpa(x, x, x, att_mask, pos_emb)
+            assert torch.allclose(out_sdpa, out, atol=1e-5)
+
+    @pytest.mark.unit
+    def test_mha_adapter_with_torch_sdpa(self):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 32, 50)
+        lengths = torch.randint(1, x.size(1), size=(x.size(0),))
+        lengths[torch.randint(0, x.size(0), size=(1,))[0]] = x.size(1)
+
+        adapter_torch_sdpa = adapter_modules.MultiHeadAttentionAdapter(
+            n_head=2, n_feat=50, dropout_rate=0.0, proj_dim=-1, use_pytorch_sdpa=True
+        )
+        adapter = adapter_modules.MultiHeadAttentionAdapter(
+            n_head=2, n_feat=50, dropout_rate=0.0, proj_dim=-1, use_pytorch_sdpa=False
+        )
+        # to dont reset linear_out parameters to zero
+        adapter.linear_out = torch.nn.Linear(adapter.linear_out.in_features, adapter.linear_out.out_features)
+
+        for original_param, sdpa_param in zip(adapter.parameters(), adapter_torch_sdpa.parameters()):
+            sdpa_param.data.copy_(original_param.data)
+
+        pad_mask, att_mask = get_mask(lengths)
+        with torch.no_grad():
+            out = adapter(x, x, x, att_mask)
+            out_sdpa = adapter_torch_sdpa(x, x, x, att_mask)
+            assert torch.allclose(out_sdpa, out, atol=1e-5)
+
+    @pytest.mark.unit
+    def test_abspos_encoding_init(self):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 32, 50)
+        lengths = torch.randint(1, x.size(1), size=(x.size(0),))
+        lengths[torch.randint(0, x.size(0), size=(1,))[0]] = x.size(1)
+
+        relpos_enc = adapter_modules.PositionalEncodingAdapter(d_model=50)
+
+        relpos_enc.extend_pe(lengths.max(), device='cpu', dtype=torch.float32)
+
+        with torch.no_grad():
+            out, pos_emb = relpos_enc(x)
+            assert (out - pos_emb - x).sum().abs() <= 1e-5
+            assert out.shape == x.shape
+
+    @pytest.mark.unit
+    def test_relpos_encoding_init(self):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 32, 50)
+        lengths = torch.randint(1, x.size(1), size=(x.size(0),))
+        lengths[torch.randint(0, x.size(0), size=(1,))[0]] = x.size(1)
+
+        relpos_enc = adapter_modules.RelPositionalEncodingAdapter(d_model=50)
+
+        relpos_enc.extend_pe(lengths.max(), device='cpu', dtype=torch.float32)
+
+        with torch.no_grad():
+            out, pos_emb = relpos_enc(x)
+            assert (out - x).sum().abs() <= 1e-8
+            assert out.shape == x.shape
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('n_head', [1, 2, 10])
+    @pytest.mark.parametrize('proj_dim', [None, -1])
+    def test_transformer_mha_adapter_init(self, n_head, proj_dim):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 32, 50)
+        lengths = torch.randint(1, x.size(1), size=(x.size(0),))
+        lengths[torch.randint(0, x.size(0), size=(1,))[0]] = x.size(1)
+
+        adapter = adapter_modules.TransformerMultiHeadAttentionAdapter(
+            num_attention_heads=n_head, hidden_size=50, attn_layer_dropout=0.0, proj_dim=proj_dim
+        )
+
+        pad_mask, att_mask = get_mask(lengths)
+        att_mask = att_mask.unsqueeze(1)
+
+        with torch.no_grad():
+            assert adapter.out_projection.weight.sum() == 0
+            if hasattr(adapter.out_projection, 'bias') and adapter.out_projection.bias is not None:
+                assert adapter.out_projection.bias.sum() == 0
+
+            out = adapter(x, x, x, att_mask)
+            assert out.sum().abs() <= 1e-8
+            assert out.shape == x.shape
+
+    @pytest.mark.unit
+    def test_mha_adapter_strategy(self):
+        adapter = adapter_modules.MultiHeadAttentionAdapter(n_head=1, n_feat=50, dropout_rate=0.0)
+        assert hasattr(adapter, 'adapter_strategy')
+        assert adapter.adapter_strategy is not None
+        # assert default strategy is set
+        assert isinstance(adapter.adapter_strategy, adapter_modules.MHAResidualAddAdapterStrategy)
+
+    @pytest.mark.unit
+    def test_relpos_mha_adapter_strategy(self):
+        adapter = adapter_modules.RelPositionMultiHeadAttentionAdapter(n_head=1, n_feat=50, dropout_rate=0.0)
+        assert hasattr(adapter, 'adapter_strategy')
+        assert adapter.adapter_strategy is not None
+        # assert default strategy is set
+        assert isinstance(adapter.adapter_strategy, adapter_modules.MHAResidualAddAdapterStrategy)
+
+    @pytest.mark.unit
+    def test_abspos_encoding_adapter_strategy(self):
+        adapter = adapter_modules.PositionalEncodingAdapter(d_model=50)
+        assert hasattr(adapter, 'adapter_strategy')
+        assert adapter.adapter_strategy is not None
+        # assert default strategy is set
+        assert isinstance(adapter.adapter_strategy, adapter_mixin_strategies.ReturnResultAdapterStrategy)
+
+    @pytest.mark.unit
+    def test_relpos_encoding_adapter_strategy(self):
+        adapter = adapter_modules.RelPositionalEncodingAdapter(d_model=50)
+        assert hasattr(adapter, 'adapter_strategy')
+        assert adapter.adapter_strategy is not None
+        # assert default strategy is set
+        assert isinstance(adapter.adapter_strategy, adapter_mixin_strategies.ReturnResultAdapterStrategy)
+
+    @pytest.mark.unit
+    def test_transformer_mha_adapter_strategy(self):
+        adapter = adapter_modules.TransformerMultiHeadAttentionAdapter(
+            num_attention_heads=1, hidden_size=50, attn_layer_dropout=0.0
+        )
+        assert hasattr(adapter, 'adapter_strategy')
+        assert adapter.adapter_strategy is not None
+        # assert default strategy is set
+        assert isinstance(adapter.adapter_strategy, adapter_modules.MHAResidualAddAdapterStrategy)

tests/collections/asr/mixins/test_transcription.py

@@ -0,0 +1,583 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import json
+import os
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+import pytest
+import torch
+from omegaconf import open_dict
+from torch.utils.data import DataLoader, Dataset
+
+from nemo.collections.asr.data.audio_to_text import _speech_collate_fn
+from nemo.collections.asr.models.aed_multitask_models import MultiTaskTranscriptionConfig
+from nemo.collections.asr.parts.mixins import TranscribeConfig, TranscriptionMixin
+from nemo.collections.asr.parts.mixins.transcription import GenericTranscriptionType
+from nemo.collections.asr.parts.submodules.multitask_decoding import MultiTaskDecodingConfig
+from nemo.collections.asr.parts.utils import Hypothesis
+
+
+class DummyModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = torch.nn.Linear(1, 1)
+
+        self.execution_count = 0
+        self.flag_begin = False
+        self.flag_end = False
+
+    def forward(self, x):
+        # Input: [1, 1] Output = [1, 1
+        out = self.encoder(x)
+        return out
+
+
+class DummyDatasetAudioOnly(Dataset):
+    def __init__(self, audio_files: List[str], config: Dict):
+        self.audio_files = audio_files
+        self.config = config
+
+    def __getitem__(self, index):
+        data = self.audio_files[index]
+        data = torch.tensor([float(data)]).view(1)
+        return data
+
+    def __len__(self):
+        return len(self.audio_files)
+
+
+class DummyDataset(Dataset):
+    def __init__(self, audio_tensors: List[str], config: Dict = None):
+        self.audio_tensors = audio_tensors
+        self.config = config
+
+    def __getitem__(self, index):
+        data = self.audio_tensors[index]
+        samples = torch.tensor(data)
+        # Calculate seq length
+        seq_len = torch.tensor(samples.shape[0], dtype=torch.long)
+
+        # Dummy text tokens
+        text_tokens = torch.tensor([0], dtype=torch.long)
+        text_tokens_len = torch.tensor(1, dtype=torch.long)
+
+        return (samples, seq_len, text_tokens, text_tokens_len)
+
+    def __len__(self):
+        return len(self.audio_tensors)
+
+
+@pytest.fixture()
+def audio_files(test_data_dir):
+    """
+    Returns a list of audio files for testing.
+    """
+    import soundfile as sf
+
+    audio_file1 = os.path.join(test_data_dir, "asr", "train", "an4", "wav", "an46-mmap-b.wav")
+    audio_file2 = os.path.join(test_data_dir, "asr", "train", "an4", "wav", "an104-mrcb-b.wav")
+
+    audio1, _ = sf.read(audio_file1, dtype='float32')
+    audio2, _ = sf.read(audio_file2, dtype='float32')
+
+    return audio1, audio2
+
+
+class TranscribableDummy(DummyModel, TranscriptionMixin):
+    def _transcribe_on_begin(self, audio, trcfg: TranscribeConfig):
+        super()._transcribe_on_begin(audio, trcfg)
+        self.flag_begin = True
+
+    def _transcribe_input_manifest_processing(self, audio_files: List[str], temp_dir: str, trcfg: TranscribeConfig):
+        # Create a dummy manifest
+        manifest_path = os.path.join(temp_dir, 'dummy_manifest.json')
+        with open(manifest_path, 'w', encoding='utf-8') as fp:
+            for audio_file in audio_files:
+                entry = {'audio_filepath': audio_file, 'duration': 100000, 'text': ''}
+                fp.write(json.dumps(entry) + '\n')
+
+        ds_config = {
+            'paths2audio_files': audio_files,
+            'batch_size': trcfg.batch_size,
+            'temp_dir': temp_dir,
+            'num_workers': trcfg.num_workers,
+            'channel_selector': trcfg.channel_selector,
+        }
+
+        return ds_config
+
+    def _setup_transcribe_dataloader(self, config: Dict) -> DataLoader:
+        dataset = DummyDatasetAudioOnly(config['paths2audio_files'], config)
+
+        return DataLoader(
+            dataset=dataset,
+            batch_size=config['batch_size'],
+            num_workers=config['num_workers'],
+            pin_memory=False,
+            drop_last=False,
+        )
+
+    def _transcribe_forward(self, batch: Any, trcfg: TranscribeConfig):
+        output = self(batch)
+        return output
+
+    def _transcribe_output_processing(self, outputs, trcfg: TranscribeConfig) -> GenericTranscriptionType:
+        self.execution_count += 1
+
+        result = []
+        for output in outputs:
+            result.append(float(output.item()))
+
+        if hasattr(trcfg, 'output_type') and trcfg.output_type == 'dict':
+            results = {'output': result}
+            return results
+
+        if hasattr(trcfg, 'output_type') and trcfg.output_type == 'dict2':
+            results = [{'output': res} for res in result]
+            return results
+
+        if hasattr(trcfg, 'output_type') and trcfg.output_type == 'tuple':
+            result = tuple(result)
+            return result
+
+        # Pass list of results by default
+        return result
+
+    def _transcribe_on_end(self, trcfg: TranscribeConfig):
+        super()._transcribe_on_end(trcfg)
+        self.flag_end = True
+
+
+@pytest.fixture()
+def dummy_model():
+    return TranscribableDummy()
+
+
+class TestTranscriptionMixin:
+    @pytest.mark.unit
+    def test_constructor_non_instance(self):
+        model = DummyModel()
+        assert not isinstance(model, TranscriptionMixin)
+        assert not hasattr(model, 'transcribe')
+
+    @pytest.mark.unit
+    def test_transcribe(self, dummy_model):
+        dummy_model = dummy_model.eval()
+        dummy_model.encoder.weight.data.fill_(1.0)
+        dummy_model.encoder.bias.data.fill_(0.0)
+
+        audio = ['1.0', '2.0', '3.0']
+        outputs = dummy_model.transcribe(audio, batch_size=1)
+        assert len(outputs) == 3
+        assert outputs[0] == 1.0
+        assert outputs[1] == 2.0
+        assert outputs[2] == 3.0
+
+    @pytest.mark.unit
+    def test_transcribe_generator(self, dummy_model):
+        dummy_model = dummy_model.eval()
+        dummy_model.encoder.weight.data.fill_(1.0)
+        dummy_model.encoder.bias.data.fill_(0.0)
+
+        audio = ['1.0', '2.0', '3.0']
+
+        transribe_config = TranscribeConfig(batch_size=1)
+        generator = dummy_model.transcribe_generator(audio, override_config=transribe_config)
+
+        outputs = []
+        index = 1
+        for result in generator:
+            outputs.extend(result)
+            assert len(result) == 1
+            assert len(outputs) == index
+            index += 1
+
+        assert len(outputs) == 3
+        assert outputs[0] == 1.0
+        assert outputs[1] == 2.0
+        assert outputs[2] == 3.0
+
+    @pytest.mark.unit
+    def test_transcribe_generator_explicit_stop_check(self, dummy_model):
+        dummy_model = dummy_model.eval()
+        dummy_model.encoder.weight.data.fill_(1.0)
+        dummy_model.encoder.bias.data.fill_(0.0)
+
+        audio = ['1.0', '2.0', '3.0']
+
+        transribe_config = TranscribeConfig(batch_size=1)
+        generator = dummy_model.transcribe_generator(audio, override_config=transribe_config)
+
+        outputs = []
+        index = 1
+        while True:
+            try:
+                result = next(generator)
+            except StopIteration:
+                break
+            outputs.extend(result)
+            assert len(result) == 1
+            assert len(outputs) == index
+            index += 1
+
+        assert len(outputs) == 3
+        assert outputs[0] == 1.0
+        assert outputs[1] == 2.0
+        assert outputs[2] == 3.0
+
+    @pytest.mark.unit
+    def test_transcribe_check_flags(self, dummy_model):
+        dummy_model = dummy_model.eval()
+
+        audio = ['1.0', '2.0', '3.0']
+        dummy_model.transcribe(audio, batch_size=1)
+        assert dummy_model.flag_begin
+        assert dummy_model.flag_end
+
+    @pytest.mark.unit
+    def test_transribe_override_config_incorrect(self, dummy_model):
+        # Not subclassing TranscribeConfig
+        @dataclass
+        class OverrideConfig:
+            batch_size: int = 1
+            output_type: str = 'dict'
+
+        dummy_model = dummy_model.eval()
+
+        audio = [1.0, 2.0, 3.0]
+        override_cfg = OverrideConfig(batch_size=1, output_type='dict')
+        with pytest.raises(ValueError):
+            _ = dummy_model.transcribe(audio, override_config=override_cfg)
+
+    @pytest.mark.unit
+    def test_transribe_override_config_correct(self, dummy_model):
+        @dataclass
+        class OverrideConfig(TranscribeConfig):
+            output_type: str = 'dict'
+            verbose: bool = False
+
+        dummy_model = dummy_model.eval()
+        dummy_model.encoder.weight.data.fill_(1.0)
+        dummy_model.encoder.bias.data.fill_(0.0)
+
+        audio = ['1.0', '2.0', '3.0']
+        override_cfg = OverrideConfig(batch_size=1, output_type='dict')
+        outputs = dummy_model.transcribe(audio, override_config=override_cfg)
+
+        assert isinstance(outputs, dict)
+        assert len(outputs) == 1
+        assert dummy_model.execution_count == 3
+        assert outputs['output'][0] == 1.0
+        assert outputs['output'][1] == 2.0
+        assert outputs['output'][2] == 3.0
+
+        # Reset execution count
+        dummy_model.execution_count = 0
+
+        override_cfg = OverrideConfig(batch_size=1, output_type='dict2')
+        outputs = dummy_model.transcribe(audio, override_config=override_cfg)
+
+        # Output now is list of dict of value each
+        assert isinstance(outputs, list)
+        assert len(outputs) == 3
+        assert dummy_model.execution_count == 3
+        assert outputs[0]['output'] == 1.0
+        assert outputs[1]['output'] == 2.0
+        assert outputs[2]['output'] == 3.0
+
+        # Reset execution count
+        dummy_model.execution_count = 0
+
+        # Test tuple
+        override_cfg = OverrideConfig(batch_size=1, output_type='tuple')
+        outputs = dummy_model.transcribe(audio, override_config=override_cfg)
+
+        assert isinstance(outputs, tuple)
+        assert len(outputs) == 1
+        assert dummy_model.execution_count == 3
+        assert outputs[0][0] == 1.0
+        assert outputs[0][1] == 2.0
+        assert outputs[0][2] == 3.0
+
+    pytest.mark.with_downloads()
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_transcribe_return_hypothesis(self, test_data_dir, fast_conformer_ctc_model):
+        audio_file = os.path.join(test_data_dir, "asr", "train", "an4", "wav", "an46-mmap-b.wav")
+
+        # Audio file test
+        outputs = fast_conformer_ctc_model.transcribe(audio_file, batch_size=1, return_hypotheses=True)
+        assert len(outputs) == 1
+        assert isinstance(outputs[0], Hypothesis)
+
+        hyp = outputs[0]
+        assert isinstance(hyp.text, str)
+        assert isinstance(hyp.y_sequence, torch.Tensor)
+        assert isinstance(hyp.alignments, torch.Tensor)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_transcribe_tensor(self, audio_files, fast_conformer_ctc_model):
+
+        audio, _ = audio_files
+        # Numpy array test
+        outputs = fast_conformer_ctc_model.transcribe(audio, batch_size=1)
+        assert len(outputs) == 1
+        assert isinstance(outputs[0], Hypothesis)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_transcribe_multiple_tensor(self, audio_files, fast_conformer_ctc_model):
+
+        audio, audio_2 = audio_files
+        # Mix second audio to torch.tensor()
+        audio_2 = torch.tensor(audio_2)
+
+        # Numpy array test
+        outputs = fast_conformer_ctc_model.transcribe([audio, audio_2], batch_size=2)
+        assert len(outputs) == 2
+        assert isinstance(outputs[0], Hypothesis)
+        assert isinstance(outputs[1], Hypothesis)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_transcribe_dataloader(self, audio_files, fast_conformer_ctc_model):
+
+        audio, audio2 = audio_files
+
+        dataset = DummyDataset([audio, audio2])
+        collate_fn = lambda x: _speech_collate_fn(x, pad_id=0)
+        dataloader = DataLoader(dataset, batch_size=2, shuffle=False, num_workers=0, collate_fn=collate_fn)
+
+        # DataLoader test
+        outputs = fast_conformer_ctc_model.transcribe(dataloader, batch_size=1)
+        assert len(outputs) == 2
+        assert isinstance(outputs[0], Hypothesis)
+        assert isinstance(outputs[1], Hypothesis)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_transcribe_return_nbest_rnnt(self, audio_files, fast_conformer_transducer_model):
+        fast_conformer_transducer_model.eval()
+        audio1, audio2 = audio_files
+
+        orig_decoding_config = copy.deepcopy(fast_conformer_transducer_model.cfg.decoding)
+
+        decoding_config = copy.deepcopy(fast_conformer_transducer_model.cfg.decoding)
+        with open_dict(decoding_config):
+            decoding_config["strategy"] = "malsd_batch"
+            decoding_config["beam"]["beam_size"] = 4
+            decoding_config["beam"]["return_best_hypothesis"] = False
+            decoding_config["beam"]["allow_cuda_graphs"] = False
+        fast_conformer_transducer_model.change_decoding_strategy(decoding_config)
+
+        outputs = fast_conformer_transducer_model.transcribe([audio1, audio2], batch_size=1, timestamps=False)
+
+        assert len(outputs) == 2
+        assert all(len(output) >= 1 for output in outputs)
+        assert all(isinstance(output, list) for output in outputs)
+        assert all(isinstance(hyp, Hypothesis) for output in outputs for hyp in output)
+
+        # Reset the decoding strategy to original
+        fast_conformer_transducer_model.change_decoding_strategy(orig_decoding_config)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_transcribe_return_nbest_canary(self, audio_files, canary_1b_flash):
+        canary_1b_flash.eval()
+        audio1, audio2 = audio_files
+
+        orig_decoding_config = copy.deepcopy(canary_1b_flash.cfg.decoding)
+
+        decoding_config = copy.deepcopy(canary_1b_flash.cfg.decoding)
+        with open_dict(decoding_config):
+            decoding_config["beam"]["beam_size"] = 4
+            decoding_config["beam"]["return_best_hypothesis"] = False
+        canary_1b_flash.change_decoding_strategy(decoding_config)
+
+        outputs = canary_1b_flash.transcribe([audio1, audio2], batch_size=1, timestamps=False)
+
+        assert len(outputs) == 2
+        assert all(len(output) >= 1 for output in outputs)
+        assert all(isinstance(output, list) for output in outputs)
+        assert all(isinstance(hyp, Hypothesis) for output in outputs for hyp in output)
+
+        # Reset the decoding strategy to original
+        canary_1b_flash.change_decoding_strategy(orig_decoding_config)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_timestamps_with_transcribe(self, audio_files, fast_conformer_ctc_model):
+        audio1, audio2 = audio_files
+
+        output = fast_conformer_ctc_model.transcribe([audio1, audio2], timestamps=True)
+
+        # check len of output
+        assert len(output) == 2
+
+        # check hypothesis object
+        assert isinstance(output[0], Hypothesis)
+        # check transcript
+        assert output[0].text == 'stop'
+        assert output[1].text == 'start'
+
+        # check timestamp
+        assert output[0].timestamp['segment'][0]['start'] == pytest.approx(0.4)
+        assert output[0].timestamp['segment'][0]['end'] == pytest.approx(0.48)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_timestamps_with_transcribe_hybrid(self, audio_files, fast_conformer_hybrid_model):
+        audio1, audio2 = audio_files
+
+        output = fast_conformer_hybrid_model.transcribe([audio1, audio2], timestamps=True)
+
+        # check len of output
+        assert len(output) == 2
+
+        # check hypothesis object
+        assert isinstance(output[0], Hypothesis)
+        # check transcript
+        assert output[0].text == 'Stop?'
+        assert output[1].text == 'Start.'
+
+        # check timestamp
+        assert output[0].timestamp['segment'][0]['start'] == pytest.approx(0.48)
+        assert output[0].timestamp['segment'][0]['end'] == pytest.approx(0.72)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_timestamps_with_transcribe_hybrid_ctc_head(self, audio_files, fast_conformer_hybrid_model):
+        audio1, audio2 = audio_files
+
+        fast_conformer_hybrid_model.change_decoding_strategy(decoder_type='ctc')
+
+        output = fast_conformer_hybrid_model.transcribe([audio1, audio2], timestamps=True)
+
+        # check len of output
+        assert len(output) == 2
+
+        # check hypothesis object
+        assert isinstance(output[0], Hypothesis)
+        # check transcript
+        assert output[0].text in ['Stop', 'Stop?']
+        assert output[1].text in ['Start', 'Start.']
+
+        # check timestamp
+        assert output[0].timestamp['segment'][0]['start'] == pytest.approx(0.4)
+        assert output[0].timestamp['segment'][0]['end'] == pytest.approx(0.72)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_timestamps_with_transcribe_canary_flash(self, audio_files, canary_1b_flash):
+        audio1, audio2 = audio_files
+
+        output = canary_1b_flash.transcribe([audio1, audio2], timestamps=True)
+
+        # check len of output
+        assert len(output) == 2
+
+        # check hypothesis object
+        assert isinstance(output[0], Hypothesis)
+        # check transcript
+        assert output[0].text == 'Stop'
+        assert output[1].text == 'start'
+
+        # check timestamp
+        assert output[0].timestamp['segment'][0]['start'] == pytest.approx(0.32)
+        assert output[0].timestamp['segment'][0]['end'] == pytest.approx(0.72)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_transcribe_return_nbest_hybrid_rnnt_ctc_prompt(self, audio_files, hybrid_rnnt_ctc_bpe_model_with_prompt):
+        """Test n-best hypothesis return for hybrid RNNT-CTC BPE model with prompts."""
+        hybrid_rnnt_ctc_bpe_model_with_prompt.eval()
+        audio1, audio2 = audio_files
+
+        orig_decoding_config = copy.deepcopy(hybrid_rnnt_ctc_bpe_model_with_prompt.cfg.decoding)
+
+        decoding_config = copy.deepcopy(hybrid_rnnt_ctc_bpe_model_with_prompt.cfg.decoding)
+        with open_dict(decoding_config):
+            decoding_config["strategy"] = "beam"
+            decoding_config["beam"]["beam_size"] = 4
+            decoding_config["beam"]["return_best_hypothesis"] = False
+
+        hybrid_rnnt_ctc_bpe_model_with_prompt.change_decoding_strategy(decoding_config)
+
+        # Transcribe audio with prompt parameters
+        hypotheses = hybrid_rnnt_ctc_bpe_model_with_prompt.transcribe(
+            [audio1, audio2], batch_size=2, return_hypotheses=True, target_lang="en-US"
+        )
+
+        # Check results
+        assert len(hypotheses) == 2
+        assert isinstance(hypotheses[0], list)  # n-best list
+        assert len(hypotheses[0]) > 0  # at least one hypothesis
+
+        # Restore original decoding config
+        hybrid_rnnt_ctc_bpe_model_with_prompt.change_decoding_strategy(orig_decoding_config)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_timestamps_with_transcribe_hybrid_prompt(self, audio_files, hybrid_rnnt_ctc_bpe_model_with_prompt):
+        audio1, audio2 = audio_files
+
+        output = hybrid_rnnt_ctc_bpe_model_with_prompt.transcribe(
+            [audio1, audio2], timestamps=True, target_lang="en-US"
+        )
+
+        # check len of output
+        assert len(output) == 2
+
+        # check hypothesis object
+        assert isinstance(output[0], Hypothesis)
+        # check transcript
+        assert output[0].text == 'Stop'
+        assert output[1].text == 'Start'
+
+        # check timestamp
+        assert output[0].timestamp['segment'][0]['start'] == pytest.approx(0.16)
+        assert output[0].timestamp['segment'][0]['end'] == pytest.approx(0.56)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_transcribe_returns_xattn(self, audio_files, canary_1b_v2):
+        canary_1b_v2.eval()
+        audio1, audio2 = audio_files
+
+        orig_decoding_config = copy.deepcopy(canary_1b_v2.cfg.decoding)
+
+        decoding_config = MultiTaskDecodingConfig()
+        decoding_config.return_xattn_scores = True
+        canary_1b_v2.change_decoding_strategy(decoding_config)
+
+        config = MultiTaskTranscriptionConfig(
+            batch_size=4,
+            return_hypotheses=True,
+            num_workers=0,
+            verbose=False,
+            prompt={'source_lang': 'en', 'target_lang': 'en'},
+            enable_chunking=False,
+        )
+
+        output = canary_1b_v2.transcribe([audio1, audio2], override_config=config)
+        assert output[0].xatt_scores is not None
+        assert output[1].xatt_scores is not None
+
+        # Reset the decoding strategy to original
+        canary_1b_v2.change_decoding_strategy(orig_decoding_config)

tests/collections/asr/numba/rnnt_loss/test_rnnt_pytorch.py

@@ -0,0 +1,629 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+
+import numpy as np
+import pytest
+import torch
+
+from nemo.collections.asr.losses.rnnt import MultiblankRNNTLossPytorch, RNNTLossPytorch, TDTLossPytorch
+from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_numpy import RNNTLoss as RNNTLoss_Numpy
+from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_pytorch import (
+    MultiblankRNNTLossNumba,
+    RNNTLossNumba,
+    TDTLossNumba,
+)
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+
+DEVICES = ['cpu']
+
+if torch.cuda.is_available():
+    DEVICES.append('cuda')
+
+CUDA_ONLY_DEVICE = ['cuda']
+
+DTYPES = [np.float32]
+if numba_utils.is_numba_cuda_fp16_supported():
+    DTYPES.append(np.float16)
+
+
+def wrap_and_call(fn, acts, labels, device):
+    if not torch.is_tensor(acts):
+        acts = torch.tensor(acts)
+
+    if 'cuda' in device:
+        acts = acts.cuda()
+
+    if not acts.requires_grad:
+        acts.requires_grad = True
+
+    lengths = [acts.shape[1]] * acts.shape[0]
+    label_lengths = [len(l) for l in labels]
+    labels = torch.LongTensor(labels)
+    lengths = torch.LongTensor(lengths)
+    label_lengths = torch.LongTensor(label_lengths)
+    if 'cuda' in device:
+        labels = labels.cuda()
+        lengths = lengths.cuda()
+        label_lengths = label_lengths.cuda()
+
+    costs = fn(acts, labels, lengths, label_lengths)
+    cost = torch.sum(costs)
+    cost.backward()
+
+    if 'cuda' in device:
+        torch.cuda.synchronize()
+
+    if acts.grad is not None:
+        grad = acts.grad.data.cpu().numpy()
+    else:
+        grad = None
+
+    return costs.data.cpu().numpy(), grad
+
+
+class TestRNNTLossPytorch:
+    @pytest.mark.unit
+    @pytest.mark.parametrize('device', DEVICES)
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_case_small(self, device, dtype):
+        if device == 'cuda':
+            numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        acts = np.array(
+            [
+                [
+                    [[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.2, 0.8, 0.1]],
+                    [[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.2, 0.1, 0.1], [0.7, 0.1, 0.2, 0.1, 0.1]],
+                ]
+            ]
+        ).astype(dtype)
+        labels = [[1, 2]]
+
+        cost_threshold = 1e-8 if dtype == np.float32 else 5e-4
+        grad_threshold = 1e-8 if dtype == np.float32 else 1e-4
+        rtol = 1e-5 if dtype == np.float32 else 1e-3
+
+        fn_pt = RNNTLossNumba(blank=0, reduction='sum')
+        pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
+
+        fn_np = RNNTLoss_Numpy()
+        np_cost, np_grads = wrap_and_call(fn_np, acts, labels, device)
+
+        fn_ag = RNNTLossPytorch(blank=0, reduction='sum')  # ag for automatic gradient computation
+        ag_cost, ag_grads = wrap_and_call(fn_ag, acts, labels, device)
+
+        expected_cost = 4.495666
+        expected_grads = np.array(
+            [
+                [
+                    [
+                        [-0.13116688, -0.3999269, 0.17703125, 0.17703125, 0.17703125],
+                        [-0.18572757, 0.12247056, -0.18168412, 0.12247056, 0.12247056],
+                        [-0.32091254, 0.06269141, 0.06928472, 0.12624499, 0.06269141],
+                    ],
+                    [
+                        [0.05456069, -0.21824276, 0.05456069, 0.05456069, 0.05456069],
+                        [0.12073959, 0.12073959, -0.48295835, 0.12073959, 0.12073959],
+                        [-0.6925882, 0.16871116, 0.18645467, 0.16871116, 0.16871116],
+                    ],
+                ]
+            ]
+        )
+
+        assert np.allclose(pt_cost, expected_cost, atol=cost_threshold, rtol=1e-6), "small_test costs mismatch."
+        assert np.allclose(pt_grads, expected_grads, atol=grad_threshold, rtol=rtol), "small_test gradient mismatch."
+
+        assert np.allclose(pt_cost, np_cost, atol=cost_threshold, rtol=rtol), "small_test costs mismatch."
+        assert np.allclose(pt_grads, np_grads, atol=grad_threshold, rtol=rtol), "small_test gradient mismatch."
+
+        assert np.allclose(ag_cost, np_cost, atol=cost_threshold, rtol=rtol), "small_test costs mismatch."
+        assert np.allclose(ag_grads, np_grads, atol=cost_threshold, rtol=rtol), "small_test gradient mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('device', DEVICES)
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_case_small_random(self, device, dtype):
+        if device == 'cuda':
+            numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        cost_threshold = 1e-8 if dtype == np.float32 else 5e-4
+        grad_threshold = 1e-8 if dtype == np.float32 else 1e-4
+        rtol = 1e-5 if dtype == np.float32 else 1e-3
+
+        rng = np.random.RandomState(0)
+        acts = rng.randn(1, 4, 3, 3).astype(dtype)
+        labels = [[1, 2]]
+
+        fn_pt = RNNTLossNumba(blank=0, reduction='sum')
+        pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
+
+        fn_np = RNNTLoss_Numpy()
+        np_cost, np_grads = wrap_and_call(fn_np, acts, labels, device)
+
+        fn_ag = RNNTLossPytorch(blank=0, reduction='sum')  # ag for automatic gradient computation
+        ag_cost, ag_grads = wrap_and_call(fn_ag, acts, labels, device)
+
+        assert np.allclose(pt_cost, np_cost, atol=cost_threshold, rtol=rtol), "small_random_test costs mismatch."
+        assert np.allclose(pt_grads, np_grads, atol=grad_threshold, rtol=rtol), "small_random_test gradient mismatch."
+
+        assert np.allclose(pt_cost, ag_cost, atol=cost_threshold, rtol=rtol), "small_random_test costs mismatch."
+        assert np.allclose(pt_grads, ag_grads, atol=grad_threshold, rtol=rtol), "small_random_test gradient mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('device', DEVICES)
+    @pytest.mark.parametrize('dtype', DTYPES)
+    @pytest.mark.parametrize('fastemit_lambda', [1.0, 0.01, 0.00001])
+    def test_case_small_random_fastemit_reg(self, device, dtype, fastemit_lambda):
+        if device == 'cuda':
+            numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        rng = np.random.RandomState(0)
+        acts = rng.randn(1, 4, 3, 3)
+        labels = [[1, 2]]
+
+        fn_pt = RNNTLossNumba(blank=0, reduction='sum', fastemit_lambda=fastemit_lambda)
+        pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
+
+        fn_np = RNNTLoss_Numpy(fastemit_lambda=fastemit_lambda)
+        np_cost, np_grads = wrap_and_call(fn_np, acts, labels, device)
+
+        assert np.allclose(pt_cost, np_cost, rtol=1e-6), "small_random_test costs mismatch."
+        assert np.allclose(pt_grads, np_grads, rtol=1e-5), "small_random_test gradient mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('device', DEVICES)
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_case_big_tensor(self, device, dtype):
+        if device == 'cuda':
+            numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        # minibatch x T x U x alphabet_size
+        activations = [
+            [
+                [
+                    [0.06535690384862791, 0.7875301411923206, 0.08159176605666074],
+                    [0.5297155426466327, 0.7506749639230854, 0.7541348379087998],
+                    [0.6097641124736383, 0.8681404965673826, 0.6225318186056529],
+                ],
+                [
+                    [0.6685222872103057, 0.8580392805336061, 0.16453892311765583],
+                    [0.989779515236694, 0.944298460961015, 0.6031678586829663],
+                    [0.9467833543605416, 0.666202507295747, 0.28688179752461884],
+                ],
+                [
+                    [0.09418426230195986, 0.3666735970751962, 0.736168049462793],
+                    [0.1666804425271342, 0.7141542198635192, 0.3993997272216727],
+                    [0.5359823524146038, 0.29182076440286386, 0.6126422611507932],
+                ],
+                [
+                    [0.3242405528768486, 0.8007644367291621, 0.5241057606558068],
+                    [0.779194617063042, 0.18331417220174862, 0.113745182072432],
+                    [0.24022162381327106, 0.3394695622533106, 0.1341595066017014],
+                ],
+            ],
+            [
+                [
+                    [0.5055615569388828, 0.051597282072282646, 0.6402903936686337],
+                    [0.43073311517251, 0.8294731834714112, 0.1774668847323424],
+                    [0.3207001991262245, 0.04288308912457006, 0.30280282975568984],
+                ],
+                [
+                    [0.6751777088333762, 0.569537369330242, 0.5584738347504452],
+                    [0.08313242153985256, 0.06016544344162322, 0.10795752845152584],
+                    [0.7486153608562472, 0.943918041459349, 0.4863558118797222],
+                ],
+                [
+                    [0.4181986264486809, 0.6524078485043804, 0.024242983423721887],
+                    [0.13458171554507403, 0.3663418070512402, 0.2958297395361563],
+                    [0.9236695822497084, 0.6899291482654177, 0.7418981733448822],
+                ],
+                [
+                    [0.25000547599982104, 0.6034295486281007, 0.9872887878887768],
+                    [0.5926057265215715, 0.8846724004467684, 0.5434495396894328],
+                    [0.6607698886038497, 0.3771277082495921, 0.3580209022231813],
+                ],
+            ],
+        ]
+
+        expected_costs = [4.2806528590890736, 3.9384369822503591]
+        expected_grads = [
+            [
+                [
+                    [-1.86843902e-01, -6.25548810e-02, 2.49398798e-01],
+                    [-2.03376666e-01, 2.02399328e-01, 9.77333169e-04],
+                    [-1.41016081e-01, 7.91234672e-02, 6.18926100e-02],
+                ],
+                [
+                    [-1.15517676e-02, -8.12802389e-02, 9.28319991e-02],
+                    [-1.54257029e-01, 2.29432687e-01, -7.51756504e-02],
+                    [-2.46593088e-01, 1.46404594e-01, 1.00188486e-01],
+                ],
+                [
+                    [-1.29182907e-02, -6.15932420e-02, 7.45115355e-02],
+                    [-5.59857301e-02, 2.19830811e-01, -1.63845062e-01],
+                    [-4.97626871e-01, 2.09239945e-01, 2.88386941e-01],
+                ],
+                [
+                    [1.36048580e-02, -3.02196294e-02, 1.66147724e-02],
+                    [1.13924511e-01, 6.27811998e-02, -1.76705718e-01],
+                    [-6.67078257e-01, 3.67658824e-01, 2.99419403e-01],
+                ],
+            ],
+            [
+                [
+                    [-3.56343776e-01, -5.53474613e-02, 4.11691219e-01],
+                    [-9.69219357e-02, 2.94591039e-02, 6.74628317e-02],
+                    [-6.35175705e-02, 2.76544970e-02, 3.58630717e-02],
+                ],
+                [
+                    [-1.54499024e-01, -7.39420280e-02, 2.28441030e-01],
+                    [-1.66789949e-01, -8.78955179e-05, 1.66877866e-01],
+                    [-1.72369644e-01, 1.05565332e-01, 6.68043196e-02],
+                ],
+                [
+                    [2.38748826e-02, -1.18255816e-01, 9.43809375e-02],
+                    [-1.04707085e-01, -1.08934477e-01, 2.13641584e-01],
+                    [-3.69844258e-01, 1.80118099e-01, 1.89726159e-01],
+                ],
+                [
+                    [2.57137045e-02, -7.94617534e-02, 5.37480488e-02],
+                    [1.22328237e-01, -2.38788679e-01, 1.16460443e-01],
+                    [-5.98686993e-01, 3.02203178e-01, 2.96483815e-01],
+                ],
+            ],
+        ]
+
+        activations = np.array(activations).astype(dtype)
+        labels = [[1, 2], [1, 1]]
+
+        cost_threshold = 1e-8 if dtype == np.float32 else 5e-4
+        grad_threshold = 1e-8 if dtype == np.float32 else 1e-4
+        rtol = 1e-3 if dtype == np.float32 else 0.1
+
+        fn_pt = RNNTLossNumba(blank=0, reduction='sum')
+        pt_costs, pt_grads = wrap_and_call(fn_pt, activations, labels, device)
+
+        fn_np = RNNTLoss_Numpy()
+        np_costs, np_grads = wrap_and_call(fn_np, activations, labels, device)
+
+        fn_ag = RNNTLossPytorch(blank=0, reduction='sum')
+        ag_costs, ag_grads = wrap_and_call(fn_ag, activations, labels, device)
+
+        assert np.allclose(pt_costs, sum(expected_costs), atol=cost_threshold), "big_test average costs mismatch."
+        assert np.allclose(
+            pt_grads, expected_grads, atol=grad_threshold, rtol=1e-3
+        ), "big_test grads for average cost mismatch."
+
+        assert np.allclose(pt_costs, np_costs, atol=cost_threshold, rtol=rtol), "big_test average costs mismatch."
+        assert np.allclose(
+            pt_grads, np_grads, atol=grad_threshold, rtol=rtol
+        ), "big_test grads for average cost mismatch."
+
+        assert np.allclose(pt_costs, ag_costs, atol=cost_threshold, rtol=rtol), "big_test average costs mismatch."
+        assert np.allclose(
+            pt_grads, ag_grads, atol=grad_threshold, rtol=rtol
+        ), "big_test grads for average cost mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('device', DEVICES)
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_case_large_random(self, device, dtype):
+        if device == 'cuda':
+            numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        rng = np.random.RandomState(0)
+        acts = rng.randn(4, 8, 11, 5).astype(dtype)
+        labels = [
+            [1, 2, 4, 3, 2, 2, 1, 1, 1, 1],
+            [3, 2, 2, 3, 4, 1, 1, 1, 1, 1],
+            [4, 4, 1, 2, 1, 3, 4, 3, 1, 2],
+            [1, 1, 2, 1, 2, 3, 3, 1, 1, 1],
+        ]
+
+        cost_threshold = 1e-8 if dtype == np.float32 else 5e-4
+        grad_threshold = 1e-8 if dtype == np.float32 else 1e-4
+        rtol = 1e-3 if dtype == np.float32 else 5e-2
+
+        fn_pt = RNNTLossNumba(blank=0, reduction='sum')
+        pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
+
+        fn_np = RNNTLoss_Numpy()
+        np_cost, np_grads = wrap_and_call(fn_np, acts, labels, device)
+
+        fn_ag = RNNTLossPytorch(blank=0, reduction='sum')
+        ag_cost, ag_grads = wrap_and_call(fn_ag, acts, labels, device)
+
+        assert np.allclose(pt_cost, np_cost, atol=cost_threshold, rtol=rtol), "large_random_test costs mismatch."
+        assert np.allclose(ag_cost, np_cost, atol=cost_threshold, rtol=rtol), "large_random_test costs mismatch."
+        assert np.allclose(pt_grads, np_grads, atol=grad_threshold, rtol=rtol), "large_random_test gradient mismatch."
+        assert np.allclose(ag_grads, np_grads, atol=grad_threshold, rtol=rtol), "large_random_test gradient mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('device', DEVICES)
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_case_small_clamp(self, device, dtype):
+        if device == 'cuda':
+            numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        GRAD_CLAMP = 0.1
+        acts = np.array(
+            [
+                [
+                    [[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.2, 0.8, 0.1]],
+                    [[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.2, 0.1, 0.1], [0.7, 0.1, 0.2, 0.1, 0.1]],
+                ]
+            ]
+        ).astype(dtype)
+        labels = [[1, 2]]
+
+        cost_threshold = 1e-8 if dtype == np.float32 else 5e-4
+        grad_threshold = 1e-8 if dtype == np.float32 else 5e-5
+        rtol = 1e-5 if dtype == np.float32 else 1e-3
+
+        fn_pt = RNNTLossNumba(blank=0, reduction='sum', clamp=GRAD_CLAMP)
+        pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
+
+        fn_np = RNNTLoss_Numpy(blank=0, clamp=GRAD_CLAMP)
+        np_cost, np_grads = wrap_and_call(fn_np, acts, labels, device)
+
+        expected_cost = 4.495666
+        expected_grads = np.array(
+            [
+                [
+                    [
+                        [-0.1, -0.1, 0.1, 0.1, 0.1],
+                        [-0.1, 0.1, -0.1, 0.1, 0.1],
+                        [-0.1, 0.06269141, 0.06928472, 0.1, 0.06269141],
+                    ],
+                    [
+                        [0.05456069, -0.1, 0.05456069, 0.05456069, 0.05456069],
+                        [0.1, 0.1, -0.1, 0.1, 0.1],
+                        [-0.1, 0.1, 0.1, 0.1, 0.1],
+                    ],
+                ]
+            ]
+        )
+
+        assert np.allclose(pt_cost, expected_cost, atol=cost_threshold, rtol=rtol), "small_test costs mismatch."
+        assert np.allclose(pt_grads, expected_grads, atol=grad_threshold, rtol=rtol), "small_test gradient mismatch."
+
+        assert np.allclose(pt_cost, np_cost, atol=cost_threshold, rtol=rtol), "small_test costs mismatch."
+        assert np.allclose(pt_grads, np_grads, atol=grad_threshold, rtol=rtol), "small_test gradient mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('device', DEVICES)
+    @pytest.mark.parametrize('dtype', DTYPES)
+    @pytest.mark.parametrize('fastemit_lambda', [1.0, 0.01, 0.00001])
+    def test_case_small_fastemit_clamp(self, device, dtype, fastemit_lambda):
+        if device == 'cuda':
+            numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        GRAD_CLAMP = 0.1
+        acts = np.array(
+            [
+                [
+                    [[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.2, 0.8, 0.1]],
+                    [[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.2, 0.1, 0.1], [0.7, 0.1, 0.2, 0.1, 0.1]],
+                ]
+            ]
+        ).astype(dtype)
+        labels = [[1, 2]]
+
+        cost_threshold = 1e-8 if dtype == np.float32 else 1e-3
+        grad_threshold = 1e-8 if dtype == np.float32 else 5e-4
+        rtol = 1e-5 if dtype == np.float32 else 1e-3
+
+        fn_pt = RNNTLossNumba(blank=0, reduction='sum', fastemit_lambda=fastemit_lambda, clamp=GRAD_CLAMP)
+        pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
+
+        fn_np = RNNTLoss_Numpy(blank=0, fastemit_lambda=fastemit_lambda, clamp=GRAD_CLAMP)
+        np_cost, np_grads = wrap_and_call(fn_np, acts, labels, device)
+
+        expected_cost = 4.495666
+        expected_cost += expected_cost * fastemit_lambda
+
+        assert np.allclose(pt_cost, expected_cost, atol=cost_threshold, rtol=rtol), "small_test costs mismatch."
+        assert np.allclose(pt_cost, np_cost, atol=cost_threshold, rtol=rtol), "small_test costs mismatch."
+        assert np.allclose(pt_grads, np_grads, atol=grad_threshold, rtol=rtol), "small_test gradient mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('device', DEVICES)
+    def test_case_small_random_accumulated(self, device):
+        if device == 'cuda':
+            numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        torch.manual_seed(0)
+        base_layer = torch.randn(3, 5, requires_grad=True)
+
+        mid1 = torch.randn(1, 4, 3, 3, requires_grad=True)
+        labels1 = [[1, 3]]
+
+        mid2 = torch.randn(1, 6, 5, 3, requires_grad=True)
+        labels2 = [[1, 2, 3, 4]]
+
+        def zero_grad():
+            if base_layer.grad is not None:
+                base_layer.grad = None
+            if mid1.grad is not None:
+                mid1.grad = None
+            if mid2.grad is not None:
+                mid2.grad = None
+
+        fn_pt = RNNTLossNumba(blank=0, reduction='sum')
+        fn_np = RNNTLoss_Numpy()
+
+        # run 1
+        acts1 = torch.matmul(mid1, base_layer)  # [1, 4, 3, 5]
+        pt_cost1, _ = wrap_and_call(fn_pt, acts1, labels1, device)
+        pt_grads1 = base_layer.grad.detach().cpu().numpy()
+
+        zero_grad()
+
+        acts1 = torch.matmul(mid1, base_layer)  # [1, 4, 3, 5]
+        np_cost1, _ = wrap_and_call(fn_np, acts1, labels1, device)
+        np_grads1 = base_layer.grad.detach().cpu().numpy()
+
+        zero_grad()
+
+        assert np.allclose(pt_grads1, np_grads1, atol=1e-6)
+
+        # run 2
+        acts2 = torch.matmul(mid2, base_layer)  # [1, 4, 3, 5]
+        pt_cost2, _ = wrap_and_call(fn_pt, acts2, labels2, device)
+        pt_grads2 = base_layer.grad.clone().cpu().numpy()
+
+        zero_grad()
+
+        acts2 = torch.matmul(mid2, base_layer)  # [1, 4, 3, 5]
+        np_cost2, _ = wrap_and_call(fn_np, acts2, labels2, device)
+        np_grads2 = base_layer.grad.clone().cpu().numpy()
+
+        zero_grad()
+
+        assert np.allclose(pt_grads2, np_grads2, atol=1e-6)
+
+        # run 1 + 2
+        acts1 = torch.matmul(mid1, base_layer)  # [1, 4, 3, 5]
+        pt_cost1, _ = wrap_and_call(fn_pt, acts1, labels1, device)
+
+        acts2 = torch.matmul(mid2, base_layer)  # [1, 6, 5, 5]
+        pt_cost2, _ = wrap_and_call(fn_pt, acts2, labels2, device)
+        pt_grads1_p_2 = base_layer.grad.clone().cpu().numpy()
+
+        assert np.allclose(pt_grads1_p_2, np_grads1 + np_grads2, atol=1e-5)
+
+
+class TestMultiblankRNNTLoss:
+    @pytest.mark.unit
+    @pytest.mark.parametrize('device', DEVICES)
+    def test_case_randomized_act_label(self, device):
+        if device == 'cuda':
+            numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+            B, T, U, V = 4, 8, 4, 8  # here V is number of non blank labels
+            big_blank_durations = [2, 4, 8]
+            sigma = 0.1
+
+            acts = torch.rand([B, T, U, V + 1 + len(big_blank_durations)])
+            labels = [[random.randrange(0, V) for i in range(U - 1)] for j in range(B)]
+
+            fn_pt = MultiblankRNNTLossNumba(
+                blank=V + len(big_blank_durations),
+                reduction='sum',
+                big_blank_durations=big_blank_durations,
+                sigma=sigma,
+            )
+            pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
+
+            fn_ag = MultiblankRNNTLossPytorch(
+                blank=V + len(big_blank_durations),
+                reduction='sum',
+                big_blank_durations=big_blank_durations,
+                sigma=sigma,
+            )  # ag for automatic gradient computation
+            ag_cost, ag_grads = wrap_and_call(fn_ag, acts, labels, device)
+
+            assert np.allclose(pt_cost, ag_cost, rtol=1e-6), "multi-blank costs mismatch."
+            assert np.allclose(pt_grads, ag_grads, rtol=1e-2), "multi-blank gradient mismatch."
+
+
+class TestTDTLoss:
+    @pytest.mark.unit
+    @pytest.mark.parametrize('device', CUDA_ONLY_DEVICE)
+    def test_case_randomized_act_label(self, device):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        B, T, U, V = 4, 8, 4, 8  # here V is number of non blank labels
+        durations = [0, 1, 2, 3, 4, 5]
+        sigma = 0.05
+
+        acts = torch.rand([B, T, U, V + 1 + len(durations)])
+        labels = [[random.randrange(0, V) for i in range(U - 1)] for j in range(B)]
+
+        fn_pt = TDTLossNumba(blank=V, reduction='sum', durations=durations, sigma=sigma)
+        pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
+
+        fn_ag = TDTLossPytorch(
+            blank=V, reduction='sum', durations=durations, sigma=sigma
+        )  # ag for automatic gradient computation
+        ag_cost, ag_grads = wrap_and_call(fn_ag, acts, labels, device)
+
+        assert np.allclose(pt_cost, ag_cost, rtol=1e-6), "tdt costs mismatch."
+        assert np.allclose(pt_grads, ag_grads, rtol=1e-2), "td gradient mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('device', CUDA_ONLY_DEVICE)
+    def test_case_randomized_act_label_no_0_duration(self, device):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        B, T, U, V = 4, 8, 4, 8  # here V is number of non blank labels
+        durations = [1, 2, 3, 4, 5]
+        sigma = 0.05
+
+        acts = torch.rand([B, T, U, V + 1 + len(durations)])
+        labels = [[random.randrange(0, V) for i in range(U - 1)] for j in range(B)]
+
+        fn_pt = TDTLossNumba(blank=V, reduction='sum', durations=durations, sigma=sigma)
+        pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
+
+        fn_ag = TDTLossPytorch(
+            blank=V, reduction='sum', durations=durations, sigma=sigma
+        )  # ag for automatic gradient computation
+        ag_cost, ag_grads = wrap_and_call(fn_ag, acts, labels, device)
+
+        assert np.allclose(pt_cost, ag_cost, rtol=1e-6), "tdt costs mismatch."
+        assert np.allclose(pt_grads, ag_grads, rtol=1e-2), "td gradient mismatch."
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('device', CUDA_ONLY_DEVICE)
+    def test_case_fixed_case_act_label(self, device):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        B, T, U, V = 1, 3, 2, 3  # here V is number of non blank labels
+        durations = [0, 1, 2]
+        sigma = 0.05
+
+        acts = torch.zeros([B, T, U, V + 1 + len(durations)])
+        labels = [[(i + j) % (V - 1) for i in range(U - 1)] for j in range(B)]
+
+        fn_pt = TDTLossNumba(blank=V, reduction='sum', durations=durations, sigma=sigma)
+        pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
+
+        expected_cost = 4.155739
+        expected_grads = [
+            [
+                [
+                    [-0.64962804, 0.25, 0.25, 0.14962798, 0.2672583, -0.16792619, -0.09933221],
+                    [0.01651875, 0.01651875, 0.01651875, -0.04955626, 0.022025, -0.01227201, -0.009753],
+                ],
+                [
+                    [-0.04892651, 0.01714851, 0.01714851, 0.01462949, -0.01143234, -0.01143234, 0.02286467],
+                    [0.12531489, 0.12531489, 0.12531489, -0.37594467, 0.16708651, 0.13027048, -0.29735702],
+                ],
+                [
+                    [-0.02572276, 0.00857425, 0.00857425, 0.00857425, -0.02286468, 0.01143234, 0.01143234],
+                    [0.13388914, 0.13388914, 0.13388914, -0.40166742, 0.17851885, -0.35703772, 0.17851885],
+                ],
+            ]
+        ]
+
+        assert np.allclose(pt_cost, expected_cost, rtol=1e-6), "tdt costs mismatch."
+        assert np.allclose(pt_grads, expected_grads, rtol=1e-2), "td gradient mismatch."
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

tests/collections/asr/numba/rnnt_loss/utils/test_gpu_rnnt_kernel.py

@@ -0,0 +1,797 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pytest
+import torch
+from numba import cuda
+
+from nemo.collections.asr.losses.rnnt_pytorch import MultiblankRNNTLossPytorch, TDTLossPytorch
+from nemo.collections.asr.parts.numba.rnnt_loss import rnnt_numpy
+from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_pytorch import certify_inputs
+from nemo.collections.asr.parts.numba.rnnt_loss.utils.cuda_utils import gpu_rnnt_kernel, reduce
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+
+
+DTYPES = [torch.float32]
+if numba_utils.is_numba_cuda_fp16_supported():
+    DTYPES.append(torch.float16)
+
+
+def log_softmax(x, axis=-1):
+    x = torch.from_numpy(x)  # zero-copy
+    x = x.float()
+    x = torch.log_softmax(x, dim=axis)
+    x = x.numpy()
+    return x
+
+
+def log_softmax_grad(x, axis=-1):
+    x = torch.tensor(x, requires_grad=True)  # alloc memory
+    y = torch.log_softmax(x, dim=axis)
+    y.sum().backward()
+    return x.grad.numpy()
+
+
+class TestRNNTCUDAKernels:
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_compute_alphas_kernel(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        random = np.random.RandomState(0)
+        original_shape = [1, 5, 11, 3]
+        B, T, U, V = original_shape
+        threshold = 1e-5 if dtype == torch.float32 else 3e-4
+
+        # Numpy kernel
+        x = random.randn(*original_shape)
+        labels = np.array([[1, 1, 1, 2, 2, 2, 1, 2, 2, 1]])  # [1, 10]
+        label_len = len(labels[0]) + 1
+        blank_idx = 0
+
+        x_np = log_softmax(x, axis=-1)
+        ground_alphas, ground_log_likelihood = rnnt_numpy.forward_pass(
+            x_np[0, :, :label_len, :], labels[0, : label_len - 1], blank_idx
+        )
+
+        # Pytorch kernel
+        device = torch.device('cuda')
+        if hasattr(cuda, 'external_stream'):
+            stream = cuda.external_stream(torch.cuda.current_stream(device).cuda_stream)
+        else:
+            stream = cuda.default_stream()
+
+        x_c = torch.tensor(x, device=device, dtype=dtype)
+        labels_c = torch.tensor(labels, device=device, dtype=torch.int64)
+
+        # Allocate workspace memory
+        denom = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        alphas = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        llForward = torch.zeros(B, device=device, dtype=x_c.dtype)
+        input_lengths = torch.tensor([T], dtype=torch.int64, device=device)
+        label_lengths = torch.tensor([len(labels[0])], dtype=torch.int64, device=device)
+
+        # certify input data
+        certify_inputs(x_c, labels_c, input_lengths, label_lengths)
+
+        # flatten activation tensor (for pointer based indexing)
+        x_c = x_c.view([-1])
+
+        # call kernel
+        # log softmax reduction
+        reduce.reduce_max(x_c, denom, rows=V, cols=B * T * U, minus=False, stream=stream)
+        reduce.reduce_exp(x_c, denom, rows=V, cols=B * T * U, minus=True, stream=stream)
+
+        # alpha kernel
+        gpu_rnnt_kernel.compute_alphas_kernel[B, U, stream, 0](
+            x_c,
+            denom,
+            alphas,
+            llForward,
+            input_lengths,
+            label_lengths,
+            labels_c,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+        )
+
+        # sync kernel
+        stream.synchronize()
+
+        # reshape alphas
+        alphas = alphas.view([B, T, U])
+        diff = ground_alphas - alphas[0].cpu().numpy()
+
+        assert np.abs(diff).mean() <= threshold
+        assert np.square(diff).mean() <= (threshold**2)
+
+        ll_diff = ground_log_likelihood - llForward[0].cpu().numpy()
+
+        assert np.abs(ll_diff).mean() <= threshold
+        assert np.square(ll_diff).mean() <= (threshold**2)
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_compute_betas_kernel(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        random = np.random.RandomState(0)
+        original_shape = [1, 5, 11, 3]
+        B, T, U, V = original_shape
+        threshold = 1e-5 if dtype == torch.float32 else 3e-4
+
+        # Numpy kernel
+        x = random.randn(*original_shape)
+        labels = np.array([[1, 1, 1, 2, 2, 2, 1, 2, 2, 1]])  # [1, 10]
+        label_len = len(labels[0]) + 1
+        blank_idx = 0
+
+        x_np = log_softmax(x, axis=-1)
+        ground_alphas, ground_log_likelihood = rnnt_numpy.backward_pass(
+            x_np[0, :, :label_len, :], labels[0, : label_len - 1], blank_idx
+        )
+
+        # Pytorch kernel
+        device = torch.device('cuda')
+        if hasattr(cuda, 'external_stream'):
+            stream = cuda.external_stream(torch.cuda.current_stream(device).cuda_stream)
+        else:
+            stream = cuda.default_stream()
+
+        x_c = torch.tensor(x, device=device, dtype=dtype)
+        labels_c = torch.tensor(labels, device=device, dtype=torch.int64)
+
+        # Allocate workspace memory
+        denom = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        betas = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        llBackward = torch.zeros(B, device=device, dtype=x_c.dtype)
+        input_lengths = torch.tensor([T], dtype=torch.int64, device=device)
+        label_lengths = torch.tensor([len(labels[0])], dtype=torch.int64, device=device)
+
+        # certify input data
+        certify_inputs(x_c, labels_c, input_lengths, label_lengths)
+
+        # flatten activation tensor (for pointer based indexing)
+        x_c = x_c.view([-1])
+
+        # call kernel
+        # log softmax reduction
+        reduce.reduce_max(x_c, denom, rows=V, cols=B * T * U, minus=False, stream=stream)
+        reduce.reduce_exp(x_c, denom, rows=V, cols=B * T * U, minus=True, stream=stream)
+
+        # beta kernel
+        gpu_rnnt_kernel.compute_betas_kernel[B, U, stream, 0](
+            x_c,
+            denom,
+            betas,
+            llBackward,
+            input_lengths,
+            label_lengths,
+            labels_c,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+        )
+
+        # sync kernel
+        stream.synchronize()
+
+        # reshape alphas
+        betas = betas.view([B, T, U])
+        diff = ground_alphas - betas[0].cpu().numpy()
+
+        assert np.abs(diff).mean() <= threshold
+        assert np.square(diff).mean() <= (threshold**2)
+
+        ll_diff = ground_log_likelihood - llBackward[0].cpu().numpy()
+
+        assert np.abs(ll_diff).mean() <= threshold
+        assert np.square(ll_diff).mean() <= (threshold**2)
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_compute_grads_kernel(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        fastemit_lambda = 0.0
+        clamp = 0.0
+
+        random = np.random.RandomState(0)
+        original_shape = [1, 5, 11, 3]
+        B, T, U, V = original_shape
+        threshold = 1e-5 if dtype == torch.float32 else 3e-5
+
+        # Numpy kernel
+        x = random.randn(*original_shape)
+        labels = torch.from_numpy(np.array([[1, 1, 1, 2, 2, 2, 1, 2, 2, 1]], dtype=np.int64))  # [1, 10]
+        audio_len = torch.from_numpy(np.array([T], dtype=np.int64))
+        label_len = torch.from_numpy(np.array([U - 1], dtype=np.int64))
+        blank_idx = 0
+
+        x_np = torch.from_numpy(x)
+        x_np.requires_grad_(True)
+
+        """
+        Here we will directly utilize the numpy variant of the loss without explicitly calling
+        the numpy functions for alpha, beta and grads. 
+        
+        This is because the grads returned by the rnnt_numpy.transduce_batch() are :
+        d/dx (alpha + beta alignment)(log_softmax(x)). 
+        But according to the chain rule, we'd still need to compute the gradient of log_softmax(x)
+        and update the alignments by hand. Instead, we will rely on pytorch to compute the gradient 
+        of the log_softmax(x) step and propagate it backwards. 
+        """
+        loss_func = rnnt_numpy.RNNTLoss(blank_idx, fastemit_lambda=fastemit_lambda, clamp=clamp)
+        loss_val = loss_func(x_np, labels, audio_len, label_len)
+        loss_val.sum().backward()
+        true_grads = x_np.grad
+
+        # Pytorch kernel
+        device = torch.device('cuda')
+        if hasattr(cuda, 'external_stream'):
+            stream = cuda.external_stream(torch.cuda.current_stream(device).cuda_stream)
+        else:
+            stream = cuda.default_stream()
+
+        x_c = torch.tensor(x, device=device, dtype=dtype)
+        labels_c = labels.clone().to(device=device, dtype=torch.int64)
+
+        # Allocate workspace memory
+        denom = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        alphas = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        betas = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        llForward = torch.zeros(B, device=device, dtype=x_c.dtype)
+        llBackward = torch.zeros(B, device=device, dtype=x_c.dtype)
+        input_lengths = torch.tensor([T], dtype=torch.int64, device=device)
+        label_lengths = torch.tensor([len(labels[0])], dtype=torch.int64, device=device)
+
+        # certify input data
+        certify_inputs(x_c, labels_c, input_lengths, label_lengths)
+
+        # flatten activation tensor (for pointer based indexing)
+        x_c = x_c.view([-1])
+        grads = torch.zeros_like(x_c, requires_grad=False)
+
+        # call kernel
+        # log softmax reduction
+        reduce.reduce_max(x_c, denom, rows=V, cols=B * T * U, minus=False, stream=stream)
+        reduce.reduce_exp(x_c, denom, rows=V, cols=B * T * U, minus=True, stream=stream)
+
+        # alpha kernel
+        gpu_rnnt_kernel.compute_alphas_kernel[B, U, stream, 0](
+            x_c,
+            denom,
+            alphas,
+            llForward,
+            input_lengths,
+            label_lengths,
+            labels_c,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+        )
+
+        # beta kernel
+        gpu_rnnt_kernel.compute_betas_kernel[B, U, stream, 0](
+            x_c,
+            denom,
+            betas,
+            llBackward,
+            input_lengths,
+            label_lengths,
+            labels_c,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+        )
+
+        # gamma kernel
+        grad_blocks_per_grid = B * T * U
+        grad_threads_per_block = gpu_rnnt_kernel.GPU_RNNT_THREAD_SIZE
+        gpu_rnnt_kernel.compute_grad_kernel[grad_blocks_per_grid, grad_threads_per_block, stream, 0](
+            grads,
+            x_c,
+            denom,
+            alphas,
+            betas,
+            llForward,
+            input_lengths,
+            label_lengths,
+            labels_c,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+            fastemit_lambda,
+            clamp,
+        )
+
+        # sync kernel
+        stream.synchronize()
+
+        # reshape grads
+        grads = grads.view([B, T, U, V])
+        diff = true_grads - grads[0].cpu().numpy()
+
+        assert np.abs(diff).mean() <= threshold
+        assert np.square(diff).mean() <= (threshold**2) * 5.0
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_compute_grads_kernel_fastemit(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        fastemit_lambda = 0.001
+        clamp = 0.0
+
+        random = np.random.RandomState(0)
+        original_shape = [1, 5, 11, 3]
+        B, T, U, V = original_shape
+        threshold = 1e-5 if dtype == torch.float32 else 3e-5
+
+        # Numpy kernel
+        x = random.randn(*original_shape)
+        labels = torch.from_numpy(np.array([[1, 1, 1, 2, 2, 2, 1, 2, 2, 1]], dtype=np.int64))  # [1, 10]
+        audio_len = torch.from_numpy(np.array([T], dtype=np.int64))
+        label_len = torch.from_numpy(np.array([U - 1], dtype=np.int64))
+        blank_idx = 0
+
+        x_np = torch.from_numpy(x)
+        x_np.requires_grad_(True)
+
+        """
+        Here we will directly utilize the numpy variant of the loss without explicitly calling
+        the numpy functions for alpha, beta and grads. 
+
+        This is because the grads returned by the rnnt_numpy.transduce_batch() are :
+        d/dx (alpha + beta alignment)(log_softmax(x)). 
+        But according to the chain rule, we'd still need to compute the gradient of log_softmax(x)
+        and update the alignments by hand. Instead, we will rely on pytorch to compute the gradient 
+        of the log_softmax(x) step and propagate it backwards. 
+        """
+        loss_func = rnnt_numpy.RNNTLoss(blank_idx, fastemit_lambda=fastemit_lambda, clamp=clamp)
+        loss_val = loss_func(x_np, labels, audio_len, label_len)
+        loss_val.sum().backward()
+        true_grads = x_np.grad
+
+        # Pytorch kernel
+        device = torch.device('cuda')
+        if hasattr(cuda, 'external_stream'):
+            stream = cuda.external_stream(torch.cuda.current_stream(device).cuda_stream)
+        else:
+            stream = cuda.default_stream()
+
+        x_c = torch.tensor(x, device=device, dtype=dtype)
+        labels_c = labels.clone().to(device=device, dtype=torch.int64)
+
+        # Allocate workspace memory
+        denom = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        alphas = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        betas = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        llForward = torch.zeros(B, device=device, dtype=x_c.dtype)
+        llBackward = torch.zeros(B, device=device, dtype=x_c.dtype)
+        input_lengths = torch.tensor([T], dtype=torch.int64, device=device)
+        label_lengths = torch.tensor([len(labels[0])], dtype=torch.int64, device=device)
+
+        # certify input data
+        certify_inputs(x_c, labels_c, input_lengths, label_lengths)
+
+        # flatten activation tensor (for pointer based indexing)
+        x_c = x_c.view([-1])
+        grads = torch.zeros_like(x_c, requires_grad=False)
+
+        # call kernel
+        # log softmax reduction
+        reduce.reduce_max(x_c, denom, rows=V, cols=B * T * U, minus=False, stream=stream)
+        reduce.reduce_exp(x_c, denom, rows=V, cols=B * T * U, minus=True, stream=stream)
+
+        # alpha kernel
+        gpu_rnnt_kernel.compute_alphas_kernel[B, U, stream, 0](
+            x_c,
+            denom,
+            alphas,
+            llForward,
+            input_lengths,
+            label_lengths,
+            labels_c,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+        )
+
+        # beta kernel
+        gpu_rnnt_kernel.compute_betas_kernel[B, U, stream, 0](
+            x_c,
+            denom,
+            betas,
+            llBackward,
+            input_lengths,
+            label_lengths,
+            labels_c,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+        )
+
+        # gamma kernel
+        grad_blocks_per_grid = B * T * U
+        grad_threads_per_block = gpu_rnnt_kernel.GPU_RNNT_THREAD_SIZE
+        gpu_rnnt_kernel.compute_grad_kernel[grad_blocks_per_grid, grad_threads_per_block, stream, 0](
+            grads,
+            x_c,
+            denom,
+            alphas,
+            betas,
+            llForward,
+            input_lengths,
+            label_lengths,
+            labels_c,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+            fastemit_lambda,
+            clamp,
+        )
+
+        # sync kernel
+        stream.synchronize()
+
+        # reshape grads
+        grads = grads.view([B, T, U, V])
+        diff = true_grads - grads[0].cpu().numpy()
+
+        assert np.abs(diff).mean() <= threshold
+        assert np.square(diff).mean() <= (threshold**2) * 5
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_compute_grads_kernel_clamp(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        fastemit_lambda = 0.0
+        clamp = 0.1
+
+        random = np.random.RandomState(0)
+        original_shape = [1, 5, 11, 3]
+        B, T, U, V = original_shape
+        threshold = 1e-5 if dtype == torch.float32 else 3e-5
+
+        # Numpy kernel
+        x = random.randn(*original_shape)
+        labels = torch.from_numpy(np.array([[1, 1, 1, 2, 2, 2, 1, 2, 2, 1]], dtype=np.int64))  # [1, 10]
+        audio_len = torch.from_numpy(np.array([T], dtype=np.int64))
+        label_len = torch.from_numpy(np.array([U - 1], dtype=np.int64))
+        blank_idx = 0
+
+        x_np = torch.from_numpy(x)
+        x_np.requires_grad_(True)
+
+        """
+        Here we will directly utilize the numpy variant of the loss without explicitly calling
+        the numpy functions for alpha, beta and grads. 
+
+        This is because the grads returned by the rnnt_numpy.transduce_batch() are :
+        d/dx (alpha + beta alignment)(log_softmax(x)). 
+        But according to the chain rule, we'd still need to compute the gradient of log_softmax(x)
+        and update the alignments by hand. Instead, we will rely on pytorch to compute the gradient 
+        of the log_softmax(x) step and propagate it backwards. 
+        """
+        loss_func = rnnt_numpy.RNNTLoss(blank_idx, fastemit_lambda=fastemit_lambda, clamp=clamp)
+        loss_val = loss_func(x_np, labels, audio_len, label_len)
+        loss_val.sum().backward()
+        true_grads = x_np.grad
+
+        # Pytorch kernel
+        device = torch.device('cuda')
+        if hasattr(cuda, 'external_stream'):
+            stream = cuda.external_stream(torch.cuda.current_stream(device).cuda_stream)
+        else:
+            stream = cuda.default_stream()
+
+        x_c = torch.tensor(x, device=device, dtype=dtype)
+        labels_c = labels.clone().to(device=device, dtype=torch.int64)
+
+        # Allocate workspace memory
+        denom = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        alphas = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        betas = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
+        llForward = torch.zeros(B, device=device, dtype=x_c.dtype)
+        llBackward = torch.zeros(B, device=device, dtype=x_c.dtype)
+        input_lengths = torch.tensor([T], dtype=torch.int64, device=device)
+        label_lengths = torch.tensor([len(labels[0])], dtype=torch.int64, device=device)
+
+        # certify input data
+        certify_inputs(x_c, labels_c, input_lengths, label_lengths)
+
+        # flatten activation tensor (for pointer based indexing)
+        x_c = x_c.view([-1])
+        grads = torch.zeros_like(x_c, requires_grad=False)
+
+        # call kernel
+        # log softmax reduction
+        reduce.reduce_max(x_c, denom, rows=V, cols=B * T * U, minus=False, stream=stream)
+        reduce.reduce_exp(x_c, denom, rows=V, cols=B * T * U, minus=True, stream=stream)
+
+        # alpha kernel
+        gpu_rnnt_kernel.compute_alphas_kernel[B, U, stream, 0](
+            x_c,
+            denom,
+            alphas,
+            llForward,
+            input_lengths,
+            label_lengths,
+            labels_c,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+        )
+
+        # beta kernel
+        gpu_rnnt_kernel.compute_betas_kernel[B, U, stream, 0](
+            x_c,
+            denom,
+            betas,
+            llBackward,
+            input_lengths,
+            label_lengths,
+            labels_c,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+        )
+
+        # gamma kernel
+        grad_blocks_per_grid = B * T * U
+        grad_threads_per_block = gpu_rnnt_kernel.GPU_RNNT_THREAD_SIZE
+        gpu_rnnt_kernel.compute_grad_kernel[grad_blocks_per_grid, grad_threads_per_block, stream, 0](
+            grads,
+            x_c,
+            denom,
+            alphas,
+            betas,
+            llForward,
+            input_lengths,
+            label_lengths,
+            labels_c,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+            fastemit_lambda,
+            clamp,
+        )
+
+        # sync kernel
+        stream.synchronize()
+
+        # reshape grads
+        grads = grads.view([B, T, U, V])
+        diff = true_grads - grads[0].cpu().numpy()
+
+        assert np.abs(diff).mean() <= threshold
+        assert np.square(diff).mean() <= (threshold**2) * 5
+
+
+class TestTDTCUDAKernels:
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
+    @pytest.mark.unit
+    def test_compute_alphas_kernel(self):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        random = np.random.RandomState(0)
+        original_shape = [1, 15, 11, 3]
+        durations = [0, 1, 2]
+        B, T, U, V = original_shape
+        Vd = len(durations)
+
+        duration_act_shape = [B, T, U, Vd]
+        sigma = 0.05
+
+        # for passing into the kernel function -- it expected unnormalized logits
+        x = random.randn(*original_shape)
+        # for passing into the pytorch function -- it expected normalized logits
+        normalized_x = log_softmax(x, axis=-1) - 0.05
+
+        xd = random.randn(*duration_act_shape)
+        # duration logits are normalized before passing into the loss computation.
+        xd = log_softmax(xd, axis=-1)
+
+        labels = np.array([[1, 1, 1, 1, 0, 0, 1, 0, 0, 1]])  # [1, 10]
+        blank_idx = V - 1
+
+        pytorch_tdt_loss = TDTLossPytorch(blank_idx, durations, sigma=sigma)
+
+        # Pytorch kernel
+        device = torch.device('cuda')
+        if hasattr(cuda, 'external_stream'):
+            stream = cuda.external_stream(torch.cuda.current_stream(device).cuda_stream)
+        else:
+            stream = cuda.default_stream()
+
+        x = torch.tensor(x, device=device, dtype=torch.float32)
+        normalized_x = torch.tensor(normalized_x, device=device, dtype=torch.float32)
+        xd = torch.tensor(xd, device=device, dtype=torch.float32)
+        labels = torch.tensor(labels, device=device, dtype=torch.long)
+        durations = torch.tensor(durations, device=device, dtype=torch.long)
+
+        # Allocate workspace memory
+        denom = torch.zeros(B * T * U, device=device, dtype=x.dtype)
+        alphas = torch.zeros(B * T * U, device=device, dtype=x.dtype)
+        llForward = torch.zeros(B, device=device, dtype=x.dtype)
+        input_lengths = torch.tensor([T], dtype=torch.long, device=device)
+        label_lengths = torch.tensor([U - 1], dtype=torch.long, device=device)
+
+        ground_log_likelihood, ground_alphas = pytorch_tdt_loss.compute_forward_prob(
+            normalized_x, xd, labels, input_lengths, label_lengths
+        )
+
+        # certify input data
+        certify_inputs(x, labels, input_lengths, label_lengths)
+
+        # flatten activation tensor (for pointer based indexing)
+        x = x.view([-1])
+        xd = xd.view([-1])
+
+        # call kernel
+        # log softmax reduction
+        reduce.reduce_max(x, denom, rows=V, cols=B * T * U, minus=False, stream=stream)
+        reduce.reduce_exp(x, denom, rows=V, cols=B * T * U, minus=True, stream=stream)
+
+        # alpha kernel
+        gpu_rnnt_kernel.compute_tdt_alphas_kernel[B, U, stream, 0](
+            x,
+            xd,
+            denom,
+            sigma,
+            alphas,
+            llForward,
+            input_lengths,
+            label_lengths,
+            labels,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+            durations,
+            Vd,
+        )
+
+        # sync kernel
+        stream.synchronize()
+
+        # reshape alphas
+        alphas = alphas.view([B, T, U])
+        diff = torch.norm(ground_alphas - alphas)
+        ll_diff = torch.norm(ground_log_likelihood - llForward)
+
+        assert diff <= 1e-3
+        assert ll_diff <= 1e-3
+
+
+class TestMultiblankRNNTCUDAKernels:
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
+    @pytest.mark.unit
+    def test_compute_alphas_kernel(self):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        random = np.random.RandomState(0)
+        original_shape = [1, 15, 11, 6]
+        big_blank_durations = [2, 3, 4]
+        B, T, U, V = original_shape
+        num_big_blanks = len(big_blank_durations)
+
+        sigma = 0.05
+
+        # for passing into the kernel function -- it expected unnormalized logits
+        x = random.randn(*original_shape)
+        # for passing into the pytorch function -- it expected normalized logits
+        normalized_x = log_softmax(x, axis=-1) - sigma
+
+        labels = np.array([[1, 1, 1, 1, 0, 0, 1, 0, 0, 1]])  # [1, 10]
+        blank_idx = V - 1
+
+        pytorch_multiblank_loss = MultiblankRNNTLossPytorch(blank_idx, big_blank_durations, sigma=sigma)
+
+        # Pytorch kernel
+        device = torch.device('cuda')
+        if hasattr(cuda, 'external_stream'):
+            stream = cuda.external_stream(torch.cuda.current_stream(device).cuda_stream)
+        else:
+            stream = cuda.default_stream()
+
+        x = torch.tensor(x, device=device, dtype=torch.float32)
+        normalized_x = torch.tensor(normalized_x, device=device, dtype=torch.float32)
+        labels = torch.tensor(labels, device=device, dtype=torch.long)
+        big_blank_durations = torch.tensor(big_blank_durations, device=device, dtype=torch.long)
+
+        # Allocate workspace memory
+        denom = torch.zeros(B * T * U, device=device, dtype=x.dtype)
+        alphas = torch.zeros(B * T * U, device=device, dtype=x.dtype)
+        llForward = torch.zeros(B, device=device, dtype=x.dtype)
+        input_lengths = torch.tensor([T], dtype=torch.long, device=device)
+        label_lengths = torch.tensor([U - 1], dtype=torch.long, device=device)
+
+        ground_log_likelihood, ground_alphas = pytorch_multiblank_loss.compute_forward_prob(
+            normalized_x, labels, input_lengths, label_lengths
+        )
+
+        # certify input data
+        certify_inputs(x, labels, input_lengths, label_lengths)
+
+        # flatten activation tensor (for pointer based indexing)
+        x = x.view([-1])
+
+        # call kernel
+        # log softmax reduction
+        reduce.reduce_max(x, denom, rows=V, cols=B * T * U, minus=False, stream=stream)
+        reduce.reduce_exp(x, denom, rows=V, cols=B * T * U, minus=True, stream=stream)
+
+        # alpha kernel
+        gpu_rnnt_kernel.compute_multiblank_alphas_kernel[B, U, stream, 0](
+            x,
+            denom,
+            sigma,
+            alphas,
+            llForward,
+            input_lengths,
+            label_lengths,
+            labels,
+            B,
+            T,
+            U,
+            V,
+            blank_idx,
+            big_blank_durations,
+            num_big_blanks,
+        )
+
+        # sync kernel
+        stream.synchronize()
+
+        # reshape alphas
+        alphas = alphas.view([B, T, U])
+        diff = torch.norm(ground_alphas - alphas)
+        ll_diff = torch.norm(ground_log_likelihood - llForward)
+
+        assert diff <= 1e-3
+        assert ll_diff <= 1e-3

tests/collections/asr/numba/rnnt_loss/utils/test_reduce.py

@@ -0,0 +1,92 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pytest
+from numba import cuda
+
+from nemo.collections.asr.parts.numba.rnnt_loss.utils.cuda_utils import reduce
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+
+DTYPES = [np.float32]
+if numba_utils.is_numba_cuda_fp16_supported():
+    DTYPES.append(np.float16)
+
+
+class TestRNNTCUDAReductions:
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_reduce_max(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        random = np.random.RandomState(0)
+        original_shape = [1, 5, 4, 3]
+        x = random.randn(*original_shape).reshape([-1]).astype(dtype)
+        dx = random.randn(*x.shape).astype(dtype)
+
+        stream = cuda.stream()
+        x_c = cuda.to_device(x, stream=stream)
+        dx_c = cuda.to_device(dx, stream=stream)
+
+        # call kernel
+        cols = np.prod(original_shape[:3])
+        reduce.reduce_max(x_c, dx_c, rows=original_shape[-1], cols=cols, minus=False, stream=stream)
+
+        # sync kernel
+        stream.synchronize()
+
+        dx_result = dx_c.copy_to_host(stream=stream)
+        del x_c, dx_c
+
+        # collect results in first [B * T * U] values; for all V
+        assert np.abs(dx_result[cols:] - dx[cols:]).sum() <= 1e-7
+        # make sure dx_result updates the [B * T * U] values
+        assert np.abs(dx_result[:cols] - dx[:cols]).sum() > 0
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_reduce_exp(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        random = np.random.RandomState(0)
+        original_shape = [1, 5, 4, 2]
+        x = random.randn(*original_shape).reshape([-1]).astype(dtype)
+        dx = np.zeros_like(x).astype(dtype)
+
+        stream = cuda.stream()
+        x_c = cuda.to_device(x, stream=stream)
+        dx_c = cuda.to_device(dx, stream=stream)
+
+        # call kernel
+        cols = np.prod(original_shape[:3])
+        reduce.reduce_exp(x_c, dx_c, rows=original_shape[-1], cols=cols, minus=False, stream=stream)
+
+        # sync kernel
+        stream.synchronize()
+
+        dx_result = dx_c.copy_to_host(stream=stream)
+        del x_c, dx_c
+
+        # collect results in first [B * T * U] values; for all V
+        assert (dx_result[cols:] - dx[cols:]).sum() <= 1e-7
+
+        # make sure dx_result updates the [B * T * U] values
+        assert np.abs(dx_result[:cols] - dx[:cols]).sum() > 0
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])

tests/collections/asr/numba/rnnt_loss/utils/test_rnnt_helper.py

@@ -0,0 +1,368 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pytest
+from numba import cuda
+
+from nemo.collections.asr.parts.numba.rnnt_loss.utils import global_constants, rnnt_helper
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+
+DTYPES = [np.float32]
+if numba_utils.is_numba_cuda_fp16_supported():
+    DTYPES.append(np.float16)
+
+
+class TestRNNTHelper:
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_log_sum_exp(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        # wrapper kernel for device function that is tested
+        @cuda.jit
+        def _kernel(x, y):
+            x_pos = cuda.grid(1)
+            if x_pos < x.shape[0] and x_pos < y.shape[0]:
+                x[x_pos] = rnnt_helper.log_sum_exp(x[x_pos], y[x_pos])
+
+        x = np.zeros([8]).astype(dtype)  # np.random.rand(8192)
+        y = np.ones([8]).astype(dtype)  # np.random.rand(8192)
+        threshold = 1e-5 if dtype == np.float32 else 2e-3
+
+        stream = cuda.stream()
+        x_c = cuda.to_device(x, stream=stream)
+        y_c = cuda.to_device(y, stream=stream)
+
+        # call kernel
+        threads_per_block = global_constants.threads_per_block()
+        blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
+        _kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c)
+
+        # sync kernel
+        stream.synchronize()
+
+        x_new = x_c.copy_to_host(stream=stream)
+        del x_c, y_c
+
+        assert (x_new.sum() - 10.506093500145782) <= threshold
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_log_sum_exp_neg_inf(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        # wrapper kernel for device function that is tested
+        @cuda.jit
+        def _kernel(x, y):
+            x_pos = cuda.grid(1)
+            if x_pos < x.shape[0] and x_pos < y.shape[0]:
+                x[x_pos] = rnnt_helper.log_sum_exp(x[x_pos], y[x_pos])
+
+        x = np.asarray([global_constants.FP32_NEG_INF] * 8).astype(dtype)
+        y = np.ones([len(x)]).astype(dtype)
+
+        stream = cuda.stream()
+        x_c = cuda.to_device(x, stream=stream)
+        y_c = cuda.to_device(y, stream=stream)
+
+        # call kernel
+        threads_per_block = global_constants.threads_per_block()
+        blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
+        _kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c)
+
+        # sync kernel
+        stream.synchronize()
+
+        x_new = x_c.copy_to_host(stream=stream)
+        del x_c, y_c
+
+        assert np.allclose(x_new, np.ones_like(x_new), atol=1e-5)
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_div_up(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        # wrapper kernel for device function that is tested
+        @cuda.jit
+        def _kernel(x, y):
+            x_pos = cuda.grid(1)
+            if x_pos < x.shape[0] and x_pos < y.shape[0]:
+                x[x_pos] = rnnt_helper.div_up(x[x_pos], y[x_pos])
+
+        x = np.full([8], fill_value=10).astype(dtype)  # np.random.rand(8192)
+        y = np.full([8], fill_value=2).astype(dtype)  # np.random.rand(8192)
+
+        stream = cuda.stream()
+        x_c = cuda.to_device(x, stream=stream)
+        y_c = cuda.to_device(y, stream=stream)
+
+        # call kernel
+        threads_per_block = global_constants.threads_per_block()
+        blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
+        _kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c)
+
+        # sync kernel
+        stream.synchronize()
+
+        x_new = x_c.copy_to_host(stream=stream)
+        del x_c, y_c
+
+        for i in range(len(x_new)):
+            assert x_new[i] == ((10 + 2 - 1) // 2)
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_add(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        # wrapper kernel for device function that is tested
+        @cuda.jit
+        def _kernel(x, y):
+            x_pos = cuda.grid(1)
+            if x_pos < x.shape[0] and x_pos < y.shape[0]:
+                x[x_pos] = rnnt_helper.add(x[x_pos], y[x_pos])
+
+        x = np.full([8], fill_value=10).astype(dtype)  # np.random.rand(8192)
+        y = np.full([8], fill_value=2).astype(dtype)  # np.random.rand(8192)
+
+        stream = cuda.stream()
+        x_c = cuda.to_device(x, stream=stream)
+        y_c = cuda.to_device(y, stream=stream)
+
+        # call kernel
+        threads_per_block = global_constants.threads_per_block()
+        blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
+        _kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c)
+
+        # sync kernel
+        stream.synchronize()
+
+        x_new = x_c.copy_to_host(stream=stream)
+        del x_c, y_c
+
+        for i in range(len(x_new)):
+            assert x_new[i] == 12
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_maximum(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        # wrapper kernel for device function that is tested
+        @cuda.jit
+        def _kernel(x, y):
+            x_pos = cuda.grid(1)
+            if x_pos < x.shape[0] and x_pos < y.shape[0]:
+                x[x_pos] = rnnt_helper.maximum(x[x_pos], y[x_pos])
+
+        x = np.full([8], fill_value=10).astype(dtype)  # np.random.rand(8192)
+        y = np.full([8], fill_value=2).astype(dtype)  # np.random.rand(8192)
+
+        stream = cuda.stream()
+        x_c = cuda.to_device(x, stream=stream)
+        y_c = cuda.to_device(y, stream=stream)
+
+        # call kernel
+        threads_per_block = global_constants.threads_per_block()
+        blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
+        _kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c)
+
+        # sync kernel
+        stream.synchronize()
+
+        x_new = x_c.copy_to_host(stream=stream)
+        del x_c, y_c
+
+        for i in range(len(x_new)):
+            assert x_new[i] == 10
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_identity(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        # wrapper kernel for device function that is tested
+        @cuda.jit
+        def _kernel(x):
+            x_pos = cuda.grid(1)
+            if x_pos < x.shape[0]:
+                x[x_pos] = rnnt_helper.identity(x[x_pos])
+
+        x = np.full([8], fill_value=10).astype(dtype)  # np.random.rand(8192)
+
+        stream = cuda.stream()
+        x_c = cuda.to_device(x, stream=stream)
+
+        # call kernel
+        threads_per_block = global_constants.threads_per_block()
+        blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
+        _kernel[blocks_per_grid, threads_per_block, stream](x_c)
+
+        # sync kernel
+        stream.synchronize()
+
+        x_new = x_c.copy_to_host(stream=stream)
+        del x_c
+
+        for i in range(len(x_new)):
+            assert x_new[i] == x[i]
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', [np.float32, np.float16])
+    def test_negate(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        # wrapper kernel for device function that is tested
+        @cuda.jit
+        def _kernel(x):
+            x_pos = cuda.grid(1)
+            if x_pos < x.shape[0]:
+                x[x_pos] = rnnt_helper.negate(x[x_pos])
+
+        x = np.full([8], fill_value=10).astype(dtype)  # np.random.rand(8192)
+
+        stream = cuda.stream()
+        x_c = cuda.to_device(x, stream=stream)
+
+        # call kernel
+        threads_per_block = global_constants.threads_per_block()
+        blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
+        _kernel[blocks_per_grid, threads_per_block, stream](x_c)
+
+        # sync kernel
+        stream.synchronize()
+
+        x_new = x_c.copy_to_host(stream=stream)
+        del x_c
+
+        for i in range(len(x_new)):
+            assert x_new[i] == -x[i]
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_exponential(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        # wrapper kernel for device function that is tested
+        @cuda.jit
+        def _kernel(x):
+            x_pos = cuda.grid(1)
+            if x_pos < x.shape[0]:
+                x[x_pos] = rnnt_helper.exponential(x[x_pos])
+
+        x = np.random.rand(8).astype(dtype)
+
+        stream = cuda.stream()
+        x_c = cuda.to_device(x, stream=stream)
+
+        # call kernel
+        threads_per_block = global_constants.threads_per_block()
+        blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
+        _kernel[blocks_per_grid, threads_per_block, stream](x_c)
+
+        # sync kernel
+        stream.synchronize()
+
+        x_new = x_c.copy_to_host(stream=stream)
+        del x_c
+
+        y = np.exp(x)
+        for i in range(len(x_new)):
+            assert (x_new[i] - y[i]) < 1e-4
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
+    @pytest.mark.unit
+    @pytest.mark.parametrize('dtype', DTYPES)
+    def test_log_plus(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        # wrapper kernel for device function that is tested
+        @cuda.jit
+        def _kernel(x, y):
+            x_pos = cuda.grid(1)
+            if x_pos < x.shape[0] and x_pos < y.shape[0]:
+                x[x_pos] = rnnt_helper.log_plus(x[x_pos], y[x_pos])
+
+        x = np.full([8], fill_value=10.0).astype(dtype)  # np.random.rand(8192)
+        y = np.full([8], fill_value=2.0).astype(dtype)  # np.random.rand(8192)
+
+        stream = cuda.stream()
+        x_c = cuda.to_device(x, stream=stream)
+        y_c = cuda.to_device(y, stream=stream)
+
+        # call kernel
+        threads_per_block = global_constants.threads_per_block()
+        blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
+        _kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c)
+
+        # sync kernel
+        stream.synchronize()
+
+        x_new = x_c.copy_to_host(stream=stream)
+        del x_c, y_c
+
+        z = np.log1p(np.exp(-np.fabs(x - y))) + np.maximum(x, y)
+
+        for i in range(len(x_new)):
+            assert x_new[i] == z[i]
+
+    @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
+    @pytest.mark.parametrize('batch_size', [8, 128, 256])
+    @pytest.mark.parametrize('fastemit_lambda', [0.0, 0.001])
+    @pytest.mark.parametrize('dtype', DTYPES)
+    @pytest.mark.unit
+    def test_compute_costs_data(self, batch_size, fastemit_lambda, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        np.random.seed(0)
+        x = np.full([batch_size], fill_value=0.0)  # np.random.rand(8192)
+        y = np.random.randn(batch_size).astype(dtype)  # np.random.rand(8192)
+        threshold = 1e-5 if dtype == np.float32 else 1e-5
+
+        stream = cuda.stream()
+        x_c = cuda.to_device(x, stream=stream)
+        y_c = cuda.to_device(y, stream=stream)
+
+        # call kernel
+        threads_per_block = min(x.shape[0], 32)
+        blocks_per_grid = (x.shape[0] + (threads_per_block - 1)) // threads_per_block
+        # Kernel call (source, dest, extra_args_...)
+        rnnt_helper.compute_costs_data[blocks_per_grid, threads_per_block, stream](y_c, x_c, fastemit_lambda)
+
+        # sync kernel
+        stream.synchronize()
+
+        x_new = x_c.copy_to_host(stream=stream)
+        del x_c, y_c
+
+        res = -(y.astype(np.float32).copy())
+        res *= 1.0 + fastemit_lambda
+
+        for i in range(len(x_new)):
+            assert abs(x_new[i] - res[i]) < threshold, f"index failed {i}"
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])

tests/collections/asr/numba/spec_augment/test_spec_aug_numba.py

@@ -0,0 +1,297 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.parts.numba.spec_augment import spec_aug_numba
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+
+
+def get_cfg(seed=0, dtype='float32', **kwargs):
+    # fmt: off
+    default = dict(b=2, f=80, t=20, device='cuda',
+                  freq_masks=2, time_masks=2, freq_width=27, time_width=0.05, mask_value=0.0,
+                  seed=seed, dtype=dtype)
+    default.update(**kwargs)
+    cfg = OmegaConf.create(default)
+    # fmt: on
+    return cfg
+
+
+# fmt: off
+def prepare_data(b, f, t, device='cuda', freq_masks=0, time_masks=0, freq_width=10, time_width=0.1,
+                 seed=0, dtype='float32',
+                 **kwargs):
+    torch.manual_seed(seed)
+
+    if dtype == 'float16':
+        dtype = torch.float16
+    else:
+        dtype = torch.float
+
+    x = torch.randn([b, f, t], dtype=dtype, device=device)
+    x_len = torch.randint(t, size=[b], device=x.device)
+
+    sh = x.shape
+    bs = sh[0]
+
+    if isinstance(time_width, int):
+        adaptive_temporal_width = False
+    else:
+        if time_width > 1.0 or time_width < 0.0:
+            raise ValueError('If `time_width` is a float value, must be in range [0, 1]')
+
+        adaptive_temporal_width = True
+
+    orginal_time_width = time_width
+
+    # Construct the freq and time masks as well as start positions
+    if freq_masks > 0:
+        freq_starts = torch.randint(0, sh[1] - freq_width + 1, size=[bs, freq_masks], device=x.device)
+        freq_lengths = torch.randint(0, freq_width + 1, size=[bs, freq_masks], device=x.device)
+    else:
+        freq_starts = torch.zeros([bs, 1], dtype=torch.int64, device=x.device)
+        freq_lengths = torch.zeros([bs, 1], dtype=torch.int64, device=x.device)
+
+    if time_masks > 0:
+        if adaptive_temporal_width:
+            time_width = (x_len * orginal_time_width).int().clamp(min=1)
+        else:
+            time_width = (
+                    torch.tensor(orginal_time_width, dtype=torch.int32, device=x.device)
+                    .unsqueeze(0)
+                    .repeat(sh[0])
+                )
+
+        time_starts = []
+        time_lengths = []
+        for idx in range(sh[0]):
+            time_starts.append(
+                torch.randint(
+                    0, max(1, x_len[idx] - time_width[idx]), size=[1, time_masks], device=x.device
+                )
+            )
+            time_lengths.append(
+                torch.randint(0, time_width[idx] + 1, size=[1, time_masks], device=x.device)
+            )
+
+        time_starts = torch.cat(time_lengths, 0)
+        time_lengths = torch.cat(time_lengths, 0)
+    else:
+        time_starts = torch.zeros([bs, 1], dtype=torch.int64, device=x.device)
+        time_lengths = torch.zeros([bs, 1], dtype=torch.int64, device=x.device)
+
+    output = dict(
+        x=x,
+        x_len=x_len,
+        freq_starts=freq_starts,
+        freq_lengths=freq_lengths,
+        time_starts=time_starts,
+        time_lengths=time_lengths,
+        sh=sh,
+    )
+    return output
+# fmt: on
+
+
+def launch_kernel(data, cfg):
+    # Launch CUDA kernel
+    # fmt: off
+    data['x'] = spec_aug_numba.launch_spec_augment_kernel(
+        x=data['x'], x_len=data['x_len'],
+        freq_starts=data['freq_starts'], freq_lengths=data['freq_lengths'],
+        time_starts=data['time_starts'], time_lengths=data['time_lengths'],
+        freq_masks=cfg.freq_masks, time_masks=cfg.time_masks, mask_value=cfg.mask_value
+    )
+    # fmt: on
+
+
+def freq_mask_check(x, x_len, f_start, f_len, mask_value, bidx):
+    check_result = True
+    for fidx in range(f_start, f_start + f_len):
+        if not (x[bidx, fidx, :] == mask_value).all():
+            check_result = False
+            break
+    assert check_result
+
+
+def time_mask_check(x, x_len, t_start, t_len, mask_value, bidx):
+    check_result = True
+    for tidx in range(t_start, t_start + t_len):
+        if tidx >= x_len[bidx]:
+            # this sample has smaller length than the time index of mask, ignore
+            continue
+
+        if not (x[bidx, :, tidx] == mask_value).all():
+            check_result = False
+            break
+    assert check_result
+
+
+class TestSpecAugmentNumba:
+    @pytest.mark.unit
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.parametrize('dtype', ['float16', 'float32'])
+    def test_spec_aug_kernel(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        cfg = get_cfg(seed=0, dtype=dtype)
+        cfg.freq_masks = 2
+        cfg.time_masks = 10
+
+        data = prepare_data(**cfg)
+
+        launch_kernel(data, cfg)
+        x, x_len, sh = data['x'], data['x_len'], data['sh']
+
+        # Assert freq masks are correct
+        for bidx in range(sh[0]):
+            for f_start, f_len in zip(data['freq_starts'][bidx], data['freq_lengths'][bidx]):
+                freq_mask_check(x, x_len, f_start, f_len, mask_value=cfg.mask_value, bidx=bidx)
+
+    @pytest.mark.unit
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.parametrize('dtype', ['float16', 'float32'])
+    def test_spec_aug_kernel_large_batch(self, dtype):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        # Change max threads per block temporarily
+        original_buffer = spec_aug_numba.MAX_THREAD_BUFFER
+        spec_aug_numba.MAX_THREAD_BUFFER = 4
+
+        cfg = get_cfg(seed=0, dtype=dtype)
+        cfg.freq_masks = 2
+        cfg.time_masks = 10
+        cfg.b = spec_aug_numba.MAX_THREAD_BUFFER + 1
+
+        data = prepare_data(**cfg)
+
+        launch_kernel(data, cfg)
+        x, x_len, sh = data['x'], data['x_len'], data['sh']
+
+        # Assert freq masks are correct
+        for bidx in range(sh[0]):
+            for f_start, f_len in zip(data['freq_starts'][bidx], data['freq_lengths'][bidx]):
+                freq_mask_check(x, x_len, f_start, f_len, mask_value=cfg.mask_value, bidx=bidx)
+
+        # Assert time masks are correct
+        for bidx in range(sh[0]):
+            for t_start, t_len in zip(data['time_starts'][bidx], data['time_lengths'][bidx]):
+                time_mask_check(x, x_len, t_start, t_len, mask_value=cfg.mask_value, bidx=bidx)
+
+        spec_aug_numba.MAX_THREAD_BUFFER = original_buffer
+
+    @pytest.mark.unit
+    @pytest.mark.run_only_on('GPU')
+    def test_spec_aug_kernel_mask_value(self):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        cfg = get_cfg(seed=0)
+        cfg.freq_masks = 2
+        cfg.time_masks = 10
+        cfg.mask_value = -1.0
+
+        data = prepare_data(**cfg)
+
+        launch_kernel(data, cfg)
+        x, x_len, sh = data['x'], data['x_len'], data['sh']
+
+        # Assert freq masks are correct
+        for bidx in range(sh[0]):
+            for f_start, f_len in zip(data['freq_starts'][bidx], data['freq_lengths'][bidx]):
+                freq_mask_check(x, x_len, f_start, f_len, mask_value=cfg.mask_value, bidx=bidx)
+
+        # Assert time masks are correct
+        for bidx in range(sh[0]):
+            for t_start, t_len in zip(data['time_starts'][bidx], data['time_lengths'][bidx]):
+                time_mask_check(x, x_len, t_start, t_len, mask_value=cfg.mask_value, bidx=bidx)
+
+    @pytest.mark.unit
+    @pytest.mark.run_only_on('GPU')
+    def test_spec_aug_kernel_grad(self):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        cfg = get_cfg(seed=0)
+        cfg.freq_masks = 2
+        cfg.time_masks = 10
+
+        data = prepare_data(**cfg)
+
+        launch_kernel(data, cfg)
+
+        result = data['x']  # inplace modification via kernel
+        y = torch.ones_like(result, requires_grad=True)
+        z = y + result
+        z.mean().backward()
+
+        assert y.grad is not None
+
+    @pytest.mark.unit
+    @pytest.mark.run_only_on('GPU')
+    def test_spec_aug_kernel_no_freq_mask(self):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        cfg = get_cfg(seed=0)
+        cfg.freq_masks = 0
+        cfg.time_masks = 10
+
+        data = prepare_data(**cfg)
+
+        launch_kernel(data, cfg)
+        x, x_len, sh = data['x'], data['x_len'], data['sh']
+
+        # Assert time masks are correct
+        for bidx in range(sh[0]):
+            for t_start, t_len in zip(data['time_starts'][bidx], data['time_lengths'][bidx]):
+                time_mask_check(x, x_len, t_start, t_len, mask_value=cfg.mask_value, bidx=bidx)
+
+    @pytest.mark.unit
+    @pytest.mark.run_only_on('GPU')
+    def test_spec_aug_kernel_no_time_mask(self):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        cfg = get_cfg(seed=0)
+        cfg.freq_masks = 2
+        cfg.time_masks = 0
+
+        data = prepare_data(**cfg)
+
+        launch_kernel(data, cfg)
+        x, x_len, sh = data['x'], data['x_len'], data['sh']
+
+        # Assert freq masks are correct
+        for bidx in range(sh[0]):
+            for f_start, f_len in zip(data['freq_starts'][bidx], data['freq_lengths'][bidx]):
+                freq_mask_check(x, x_len, f_start, f_len, mask_value=cfg.mask_value, bidx=bidx)
+
+    @pytest.mark.unit
+    @pytest.mark.run_only_on('GPU')
+    def test_spec_aug_kernel_no_freq_time_mask(self):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        cfg = get_cfg(seed=0)
+        cfg.freq_masks = 0
+        cfg.time_masks = 0
+
+        data = prepare_data(**cfg)
+
+        x, x_len, sh = data['x'], data['x_len'], data['sh']
+        x_copy = x.clone()
+        launch_kernel(data, cfg)
+
+        # Assert no data edits occured
+        assert (data['x'] - x_copy).abs().mean() <= 1e-9

tests/collections/asr/test_asr_classification_model.py

@@ -0,0 +1,390 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import json
+import os
+
+import tempfile
+
+import lightning.pytorch as pl
+import numpy as np
+import pytest
+import soundfile as sf
+import torch
+from omegaconf import DictConfig, ListConfig
+
+from nemo.collections.asr.data import audio_to_label
+from nemo.collections.asr.models import EncDecClassificationModel, EncDecFrameClassificationModel, configs
+from nemo.utils.config_utils import assert_dataclass_signature_match
+
+
+@pytest.fixture()
+def speech_classification_model():
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': 32,
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                }
+            ],
+        },
+    }
+
+    decoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASRDecoderClassification',
+        'params': {
+            'feat_in': 32,
+            'num_classes': 30,
+        },
+    }
+
+    modelConfig = DictConfig(
+        {
+            'preprocessor': DictConfig(preprocessor),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'labels': ListConfig(["dummy_cls_{}".format(i + 1) for i in range(30)]),
+        }
+    )
+    model = EncDecClassificationModel(cfg=modelConfig)
+    return model
+
+
+@pytest.fixture()
+def frame_classification_model():
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': 32,
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                }
+            ],
+        },
+    }
+
+    decoder = {
+        'cls': 'nemo.collections.common.parts.MultiLayerPerceptron',
+        'params': {
+            'hidden_size': 32,
+            'num_classes': 5,
+        },
+    }
+
+    optim = {
+        'name': 'sgd',
+        'lr': 0.01,
+        'weight_decay': 0.001,
+        'momentum': 0.9,
+    }
+
+    modelConfig = DictConfig(
+        {
+            'preprocessor': DictConfig(preprocessor),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'optim': DictConfig(optim),
+            'labels': ListConfig(["0", "1"]),
+        }
+    )
+    model = EncDecFrameClassificationModel(cfg=modelConfig)
+    return model
+
+
+class TestEncDecClassificationModel:
+    @pytest.mark.unit
+    def test_constructor(self, speech_classification_model):
+        asr_model = speech_classification_model.train()
+
+        conv_cnt = (64 * 32 * 1 + 32) + (64 * 1 * 1 + 32)  # separable kernel + bias + pointwise kernel + bias
+        bn_cnt = (4 * 32) * 2  # 2 * moving averages
+        dec_cnt = 32 * 30 + 30  # fc + bias
+
+        param_count = conv_cnt + bn_cnt + dec_cnt
+        assert asr_model.num_weights == param_count
+
+        # Check to/from config_dict:
+        confdict = asr_model.to_config_dict()
+        instance2 = EncDecClassificationModel.from_config_dict(confdict)
+
+        assert isinstance(instance2, EncDecClassificationModel)
+
+    @pytest.mark.unit
+    def test_forward(self, speech_classification_model):
+        asr_model = speech_classification_model.eval()
+
+        asr_model.preprocessor.featurizer.dither = 0.0
+        asr_model.preprocessor.featurizer.pad_to = 0
+
+        input_signal = torch.randn(size=(4, 512))
+        length = torch.randint(low=321, high=500, size=[4])
+
+        with torch.no_grad():
+            # batch size 1
+            logprobs_instance = []
+            for i in range(input_signal.size(0)):
+                logprobs_ins = asr_model.forward(
+                    input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1]
+                )
+                logprobs_instance.append(logprobs_ins)
+            logprobs_instance = torch.cat(logprobs_instance, 0)
+
+            # batch size 4
+            logprobs_batch = asr_model.forward(input_signal=input_signal, input_signal_length=length)
+
+        assert logprobs_instance.shape == logprobs_batch.shape
+        diff = torch.mean(torch.abs(logprobs_instance - logprobs_batch))
+        assert diff <= 1e-6
+        diff = torch.max(torch.abs(logprobs_instance - logprobs_batch))
+        assert diff <= 1e-6
+
+    @pytest.mark.unit
+    def test_vocab_change(self, speech_classification_model):
+        asr_model = speech_classification_model.train()
+
+        old_labels = copy.deepcopy(asr_model._cfg.labels)
+        nw1 = asr_model.num_weights
+        asr_model.change_labels(new_labels=old_labels)
+        # No change
+        assert nw1 == asr_model.num_weights
+        new_labels = copy.deepcopy(old_labels)
+        new_labels.append('dummy_cls_31')
+        new_labels.append('dummy_cls_32')
+        new_labels.append('dummy_cls_33')
+        asr_model.change_labels(new_labels=new_labels)
+        # fully connected + bias
+        assert asr_model.num_weights == nw1 + 3 * (asr_model.decoder._feat_in + 1)
+
+    @pytest.mark.unit
+    def test_transcription(self, speech_classification_model, test_data_dir):
+        # Ground truth labels = ["yes", "no"]
+        audio_filenames = ['an22-flrp-b.wav', 'an90-fbbh-b.wav']
+        audio_paths = [os.path.join(test_data_dir, "asr", "train", "an4", "wav", fp) for fp in audio_filenames]
+
+        model = speech_classification_model.eval()
+
+        # Test Top 1 classification transcription
+        results = model.transcribe(audio_paths, batch_size=2)
+        assert len(results) == 2
+        assert results[0].shape == torch.Size([1])
+
+        # Test Top 5 classification transcription
+        model._accuracy.top_k = [5]  # set top k to 5 for accuracy calculation
+        results = model.transcribe(audio_paths, batch_size=2)
+        assert len(results) == 2
+        assert results[0].shape == torch.Size([5])
+
+        # Test Top 1 and Top 5 classification transcription
+        model._accuracy.top_k = [1, 5]
+        results = model.transcribe(audio_paths, batch_size=2)
+        assert len(results) == 2
+        assert results[0].shape == torch.Size([2, 1])
+        assert results[1].shape == torch.Size([2, 5])
+        assert model._accuracy.top_k == [1, 5]
+
+        # Test log probs extraction
+        model._accuracy.top_k = [1]
+        results = model.transcribe(audio_paths, batch_size=2, logprobs=True)
+        assert len(results) == 2
+        assert results[0].shape == torch.Size([len(model.cfg.labels)])
+
+        # Test log probs extraction remains same for any top_k
+        model._accuracy.top_k = [5]
+        results = model.transcribe(audio_paths, batch_size=2, logprobs=True)
+        assert len(results) == 2
+        assert results[0].shape == torch.Size([len(model.cfg.labels)])
+
+    @pytest.mark.unit
+    def test_EncDecClassificationDatasetConfig_for_AudioToSpeechLabelDataset(self):
+        # ignore some additional arguments as dataclass is generic
+        IGNORE_ARGS = [
+            'is_tarred',
+            'num_workers',
+            'batch_size',
+            'tarred_audio_filepaths',
+            'shuffle',
+            'pin_memory',
+            'drop_last',
+            'tarred_shard_strategy',
+            'shuffle_n',
+            # `featurizer` is supplied at runtime
+            'featurizer',
+            # additional ignored arguments
+            'vad_stream',
+            'int_values',
+            'sample_rate',
+            'normalize_audio',
+            'augmentor',
+            'bucketing_batch_size',
+            'bucketing_strategy',
+            'bucketing_weights',
+        ]
+
+        REMAP_ARGS = {'trim_silence': 'trim'}
+
+        result = assert_dataclass_signature_match(
+            audio_to_label.AudioToSpeechLabelDataset,
+            configs.EncDecClassificationDatasetConfig,
+            ignore_args=IGNORE_ARGS,
+            remap_args=REMAP_ARGS,
+        )
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+
+class TestEncDecFrameClassificationModel(TestEncDecClassificationModel):
+    @pytest.mark.parametrize(["logits_len", "labels_len"], [(20, 10), (21, 10), (19, 10), (20, 9), (20, 11)])
+    @pytest.mark.unit
+    def test_reshape_labels(self, frame_classification_model, logits_len, labels_len):
+        model = frame_classification_model.eval()
+
+        logits = torch.ones(4, logits_len, 2)
+        labels = torch.ones(4, labels_len)
+        logits_len = torch.tensor([6, 7, 8, 9])
+        labels_len = torch.tensor([5, 6, 7, 8])
+        labels_new, labels_len_new = model.reshape_labels(
+            logits=logits, labels=labels, logits_len=logits_len, labels_len=labels_len
+        )
+        assert labels_new.size(1) == logits.size(1)
+        assert torch.equal(labels_len_new, torch.tensor([6, 7, 8, 9]))
+
+    @pytest.mark.unit
+    def test_EncDecClassificationDatasetConfig_for_AudioToMultiSpeechLabelDataset(self):
+        # ignore some additional arguments as dataclass is generic
+        IGNORE_ARGS = [
+            'is_tarred',
+            'num_workers',
+            'batch_size',
+            'tarred_audio_filepaths',
+            'shuffle',
+            'pin_memory',
+            'drop_last',
+            'tarred_shard_strategy',
+            'shuffle_n',
+            # `featurizer` is supplied at runtime
+            'featurizer',
+            # additional ignored arguments
+            'vad_stream',
+            'int_values',
+            'sample_rate',
+            'normalize_audio',
+            'augmentor',
+            'bucketing_batch_size',
+            'bucketing_strategy',
+            'bucketing_weights',
+            'delimiter',
+            'normalize_audio_db',
+            'normalize_audio_db_target',
+            'window_length_in_sec',
+            'shift_length_in_sec',
+        ]
+
+        REMAP_ARGS = {'trim_silence': 'trim'}
+
+        result = assert_dataclass_signature_match(
+            audio_to_label.AudioToMultiLabelDataset,
+            configs.EncDecClassificationDatasetConfig,
+            ignore_args=IGNORE_ARGS,
+            remap_args=REMAP_ARGS,
+        )
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_frame_classification_model(self, frame_classification_model: EncDecFrameClassificationModel):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            # generate random audio
+            audio = np.random.randn(16000 * 1)
+            # save the audio
+            audio_path = os.path.join(temp_dir, "audio.wav")
+            sf.write(audio_path, audio, 16000)
+
+            dummy_labels = "0 0 0 0 1 1 1 1 0 0 0 0"
+
+            dummy_sample = {
+                "audio_filepath": audio_path,
+                "offset": 0.0,
+                "duration": 1.0,
+                "label": dummy_labels,
+            }
+
+            # create a manifest file
+            manifest_path = os.path.join(temp_dir, "dummy_manifest.json")
+            with open(manifest_path, "w") as f:
+                for i in range(4):
+                    f.write(json.dumps(dummy_sample) + "\n")
+
+            dataloader_cfg = {
+                "batch_size": 2,
+                "manifest_filepath": manifest_path,
+                "sample_rate": 16000,
+                "num_workers": 0,
+                "shuffle": False,
+                "labels": ["0", "1"],
+            }
+
+            trainer_cfg = {
+                "max_epochs": 1,
+                "devices": 1,
+                "accelerator": "auto",
+            }
+
+            optim = {
+                'name': 'sgd',
+                'lr': 0.01,
+                'weight_decay': 0.001,
+                'momentum': 0.9,
+            }
+
+            trainer = pl.Trainer(**trainer_cfg)
+            frame_classification_model.set_trainer(trainer)
+            frame_classification_model.setup_optimization(DictConfig(optim))
+            frame_classification_model.setup_training_data(dataloader_cfg)
+            frame_classification_model.setup_validation_data(dataloader_cfg)
+
+            trainer.fit(frame_classification_model)

tests/collections/asr/test_asr_context_biasing.py

@@ -0,0 +1,161 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+import tempfile
+
+import numpy as np
+import pytest
+import torch
+from lightning.pytorch import Trainer
+
+from nemo.collections.asr.models import EncDecCTCModelBPE
+from nemo.collections.asr.parts import context_biasing
+from nemo.collections.asr.parts.context_biasing.ctc_based_word_spotter import WSHyp
+from nemo.collections.asr.parts.utils import rnnt_utils
+
+
+@pytest.fixture(scope="module")
+def conformer_ctc_bpe_model():
+    model = EncDecCTCModelBPE.from_pretrained(model_name="stt_en_conformer_ctc_small")
+    model.set_trainer(Trainer(devices=1, accelerator="cpu"))
+    model = model.eval()
+    return model
+
+
+class TestContextGraphCTC:
+    @pytest.mark.unit
+    def test_graph_building(self):
+        context_biasing_list = [["gpu", [['▁g', 'p', 'u'], ['▁g', '▁p', '▁u']]]]
+        context_graph = context_biasing.ContextGraphCTC(blank_id=1024)
+        context_graph.add_to_graph(context_biasing_list)
+        assert context_graph.num_nodes == 8
+        assert context_graph.blank_token == 1024
+        assert not context_graph.root.next['▁g'].is_end
+        assert context_graph.root.next['▁g'].next['p'].next['u'].is_end
+        assert context_graph.root.next['▁g'].next['p'].next['u'].word == 'gpu'
+        assert context_graph.root.next['▁g'].next['▁p'].next['▁u'].is_end
+        assert context_graph.root.next['▁g'].next['▁p'].next['▁u'].word == 'gpu'
+
+
+class TestCTCWordSpotter:
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    def test_run_word_spotter(self, test_data_dir, conformer_ctc_bpe_model):
+        asr_model = conformer_ctc_bpe_model
+        audio_file_path = os.path.join(test_data_dir, "asr/test/an4/wav/cen3-mjwl-b.wav")
+        target_text = "nineteen"
+        target_tokenization = asr_model.tokenizer.text_to_ids(target_text)
+        ctc_logprobs = (
+            asr_model.transcribe([audio_file_path], batch_size=1, return_hypotheses=True)[0].alignments.cpu().numpy()
+        )
+        context_biasing_list = [[target_text, [target_tokenization]]]
+        context_graph = context_biasing.ContextGraphCTC(blank_id=asr_model.decoding.blank_id)
+        context_graph.add_to_graph(context_biasing_list)
+
+        # without context biasing
+        ws_results = context_biasing.run_word_spotter(
+            ctc_logprobs,
+            context_graph,
+            asr_model,
+            blank_idx=asr_model.decoding.blank_id,
+            beam_threshold=5.0,
+            cb_weight=0.0,
+            ctc_ali_token_weight=0.6,
+        )
+        assert len(ws_results) == 0
+
+        # with context biasing
+        ws_results = context_biasing.run_word_spotter(
+            ctc_logprobs,
+            context_graph,
+            asr_model,
+            blank_idx=asr_model.decoding.blank_id,
+            beam_threshold=5.0,
+            cb_weight=3.0,
+            ctc_ali_token_weight=0.6,
+        )
+        assert len(ws_results) == 1
+        assert ws_results[0].word == target_text
+        assert ws_results[0].start_frame == 9
+        assert ws_results[0].end_frame == 19
+        torch.testing.assert_close(ws_results[0].score, 8.9967, atol=1e-3, rtol=1e-4)
+
+
+class TestContextBiasingUtils:
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    def test_merge_alignment_with_ws_hyps(self, conformer_ctc_bpe_model):
+        asr_model = conformer_ctc_bpe_model
+        blank_idx = asr_model.decoding.blank_id
+        ws_results = [WSHyp(word="gpu", score=6.0, start_frame=0, end_frame=2)]
+
+        # ctc argmax predictions
+        preds = np.array([120, 29, blank_idx, blank_idx])
+        pred_text, raw_text = context_biasing.merge_alignment_with_ws_hyps(
+            preds,
+            asr_model,
+            ws_results,
+            decoder_type="ctc",
+            blank_idx=blank_idx,
+        )
+        assert raw_text == "gp"
+        assert pred_text == "gpu"
+
+        # rnnt token predictions
+        preds = rnnt_utils.Hypothesis(
+            y_sequence=torch.tensor([120, 29]),
+            score=0.0,
+            timestamp=torch.tensor([0, 1, 2, 3]),
+        )
+        pred_text, raw_text = context_biasing.merge_alignment_with_ws_hyps(
+            preds,
+            asr_model,
+            ws_results,
+            decoder_type="rnnt",
+            blank_idx=blank_idx,
+        )
+        assert raw_text == "gp"
+        assert pred_text == "gpu"
+
+        # rnnt empty token predictions
+        preds = rnnt_utils.Hypothesis(
+            y_sequence=[],
+            score=0.0,
+            timestamp=[],
+        )
+        pred_text, raw_text = context_biasing.merge_alignment_with_ws_hyps(
+            preds,
+            asr_model,
+            ws_results,
+            decoder_type="rnnt",
+            blank_idx=blank_idx,
+        )
+        assert raw_text == ""
+        assert pred_text == "gpu"
+
+    @pytest.mark.unit
+    def test_compute_fscore(self):
+        recog_manifest = """{"audio_filepath": "test.wav", "duration": 1.0, "text": "a new gpu for nvidia", "pred_text": "a new gpu for invidia"}\n"""
+        context_words = ["gpu", "cpu", "nvidia"]
+        with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8') as f:
+            f.write(recog_manifest)
+            f.seek(0)
+            fscore_stats = context_biasing.compute_fscore(f.name, context_words)
+            assert (round(fscore_stats[0], 4), round(fscore_stats[1], 4), round(fscore_stats[2], 4)) == (
+                1.0,
+                0.5,
+                0.6667,
+            )

tests/collections/asr/test_asr_ctc_encoder_model_bpe.py

@@ -0,0 +1,429 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import os
+import shutil
+import tempfile
+
+import pytest
+import torch
+from lhotse import CutSet, MonoCut
+from lhotse.testing.dummies import DummyManifest
+from omegaconf import DictConfig
+
+from nemo.collections.asr.data import audio_to_text
+from nemo.collections.asr.data.audio_to_text_lhotse import LhotseSpeechToTextBpeDataset
+from nemo.collections.asr.models import configs
+from nemo.collections.asr.models.ctc_bpe_models import EncDecCTCModelBPE
+from nemo.collections.asr.parts.submodules import ctc_beam_decoding as beam_decode
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCBPEDecoding, CTCBPEDecodingConfig
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.collections.common import tokenizers
+from nemo.utils.config_utils import assert_dataclass_signature_match
+
+
+@pytest.fixture()
+def asr_model(test_data_dir):
+    preprocessor = {'_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor'}
+    encoder = {
+        '_target_': 'nemo.collections.asr.modules.ConvASREncoder',
+        'feat_in': 64,
+        'activation': 'relu',
+        'conv_mask': True,
+        'jasper': [
+            {
+                'filters': 1024,
+                'repeat': 1,
+                'kernel': [1],
+                'stride': [1],
+                'dilation': [1],
+                'dropout': 0.0,
+                'residual': False,
+                'separable': True,
+                'se': True,
+                'se_context_size': -1,
+            }
+        ],
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+        'feat_in': 1024,
+        'num_classes': -1,
+        'vocabulary': None,
+    }
+
+    tokenizer = {'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'}
+
+    modelConfig = DictConfig(
+        {
+            'preprocessor': DictConfig(preprocessor),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'tokenizer': DictConfig(tokenizer),
+        }
+    )
+
+    model_instance = EncDecCTCModelBPE(cfg=modelConfig)
+    return model_instance
+
+
+class TestEncDecCTCModel:
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_constructor(self, asr_model):
+        asr_model.train()
+        # TODO: make proper config and assert correct number of weights
+        # Check to/from config_dict:
+        confdict = asr_model.to_config_dict()
+        instance2 = EncDecCTCModelBPE.from_config_dict(confdict)
+        assert isinstance(instance2, EncDecCTCModelBPE)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_forward(self, asr_model):
+        asr_model = asr_model.eval()
+
+        asr_model.preprocessor.featurizer.dither = 0.0
+        asr_model.preprocessor.featurizer.pad_to = 0
+
+        input_signal = torch.randn(size=(4, 512))
+        length = torch.randint(low=321, high=500, size=[4])
+
+        with torch.no_grad():
+            # batch size 1
+            logprobs_instance = []
+            for i in range(input_signal.size(0)):
+                logprobs_ins, _, _ = asr_model.forward(
+                    input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1]
+                )
+                logprobs_instance.append(logprobs_ins)
+                print(len(logprobs_ins))
+            logprobs_instance = torch.cat(logprobs_instance, 0)
+
+            # batch size 4
+            logprobs_batch, _, _ = asr_model.forward(input_signal=input_signal, input_signal_length=length)
+
+        assert logprobs_instance.shape == logprobs_batch.shape
+        diff = torch.mean(torch.abs(logprobs_instance - logprobs_batch))
+        assert diff <= 1e-6
+        diff = torch.max(torch.abs(logprobs_instance - logprobs_batch))
+        assert diff <= 1e-6
+
+    @pytest.mark.unit
+    def test_predict_step(self, asr_model):
+        asr_model = asr_model.eval()
+        cuts = DummyManifest(CutSet, begin_id=0, end_id=1, with_data=True)
+        dataset = LhotseSpeechToTextBpeDataset(tokenizer=asr_model.tokenizer, return_cuts=True)
+        batch = dataset[cuts]
+        outputs = asr_model.predict_step(batch, 0)
+        assert len(outputs) == 1
+        assert len(outputs[0]) == 2
+        assert isinstance(outputs[0][0], MonoCut)
+        assert isinstance(outputs[0][1], Hypothesis)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_save_restore_artifact(self, asr_model):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            save_path = os.path.join(tmpdir, 'ctc_bpe.nemo')
+            asr_model.train()
+            asr_model.save_to(save_path)
+
+            new_model = EncDecCTCModelBPE.restore_from(save_path)
+            assert isinstance(new_model, type(asr_model))
+            assert new_model.vocab_path.endswith('_vocab.txt')
+
+            assert len(new_model.tokenizer.tokenizer.get_vocab()) == 128
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_save_restore_artifact_spe(self, asr_model, test_data_dir):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128")
+            asr_model.change_vocabulary(new_tokenizer_dir=tokenizer_dir, new_tokenizer_type='bpe')
+
+            save_path = os.path.join(tmpdir, 'ctc_bpe.nemo')
+            asr_model.train()
+            asr_model.save_to(save_path)
+
+            new_model = EncDecCTCModelBPE.restore_from(save_path)
+            assert isinstance(new_model, type(asr_model))
+            assert isinstance(new_model.tokenizer, tokenizers.SentencePieceTokenizer)
+            assert new_model.model_path.endswith('_tokenizer.model')
+            assert new_model.vocab_path.endswith('_vocab.txt')
+            assert new_model.spe_vocab_path.endswith('_tokenizer.vocab')
+
+            assert new_model.tokenizer.tokenizer.vocab_size == 128
+            assert len(new_model.tokenizer.tokenizer.get_vocab()) == 128
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_save_restore_artifact_agg(self, asr_model, test_data_dir):
+        tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128")
+        tok_en = {"dir": tokenizer_dir, "type": "wpe"}
+        # the below is really an english tokenizer but we pretend it is spanish
+        tok_es = {"dir": tokenizer_dir, "type": "wpe"}
+        tcfg = DictConfig({"type": "agg", "langs": {"en": tok_en, "es": tok_es}})
+        with tempfile.TemporaryDirectory() as tmpdir:
+            asr_model.change_vocabulary(new_tokenizer_dir=tcfg, new_tokenizer_type="agg")
+
+            save_path = os.path.join(tmpdir, "ctc_agg.nemo")
+            asr_model.train()
+            asr_model.save_to(save_path)
+
+            new_model = EncDecCTCModelBPE.restore_from(save_path)
+            assert isinstance(new_model, type(asr_model))
+            assert isinstance(new_model.tokenizer, tokenizers.AggregateTokenizer)
+
+            # should be double
+            assert new_model.tokenizer.tokenizer.vocab_size == 264
+            assert len(new_model.tokenizer.tokenizer.get_vocab()) == 264
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_vocab_change(self, test_data_dir, asr_model):
+        old_vocab = copy.deepcopy(asr_model.decoder.vocabulary)
+
+        with tempfile.TemporaryDirectory() as save_dir:
+            save_path = os.path.join(save_dir, 'temp.nemo')
+
+            with tempfile.TemporaryDirectory() as tmpdir:
+                old_tmpdir_path = tmpdir
+
+                old_tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
+                new_tokenizer_dir = os.path.join(tmpdir, 'tokenizer')
+
+                os.makedirs(new_tokenizer_dir, exist_ok=True)
+                shutil.copy2(old_tokenizer_dir, new_tokenizer_dir)
+
+                nw1 = asr_model.num_weights
+                asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe')
+                # No change
+                assert nw1 == asr_model.num_weights
+
+                with open(os.path.join(new_tokenizer_dir, 'vocab.txt'), 'a+') as f:
+                    f.write("!\n")
+                    f.write('$\n')
+                    f.write('@\n')
+
+                asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe')
+                # fully connected + bias
+                assert asr_model.num_weights == nw1 + 3 * (asr_model.decoder._feat_in + 1)
+
+                new_vocab = copy.deepcopy(asr_model.decoder.vocabulary)
+                assert len(old_vocab) != len(new_vocab)
+
+                # save the model (after change of vocabulary)
+                asr_model.save_to(save_path)
+                assert os.path.exists(save_path)
+                # delete copied version of the vocabulary from nested tmpdir (by scope)
+
+            # assert copied vocab no longer exists
+            assert not os.path.exists(os.path.join(old_tmpdir_path, 'tokenizer', 'vocab.txt'))
+
+            # make a copy of the tokenizer before renaming
+            try:
+                os.rename(old_tokenizer_dir, old_tokenizer_dir + '.bkp')
+                assert not os.path.exists(old_tokenizer_dir)
+
+                # restore model from .nemo
+                asr_model2 = EncDecCTCModelBPE.restore_from(save_path)
+                assert isinstance(asr_model2, EncDecCTCModelBPE)
+
+                # Check if vocabulary size is same
+                assert asr_model.tokenizer.tokenizer.vocab_size == asr_model2.tokenizer.tokenizer.vocab_size
+
+                # Make a copy of the tokenizer
+                new_tokenizer_dir = os.path.join(save_dir, 'tokenizer')
+
+                os.makedirs(new_tokenizer_dir, exist_ok=True)
+                new_tokenizer_path = os.path.join(new_tokenizer_dir, 'vocab.txt')
+                with open(new_tokenizer_path, 'w') as f:
+                    for v in asr_model2.tokenizer.tokenizer.get_vocab():
+                        f.write(f"{v}\n")
+
+                    # Add some new tokens too
+                    f.write("^\n")
+                    f.write("^^\n")
+                    f.write("^^^\n")
+
+                assert os.path.exists(new_tokenizer_path)
+
+                # change vocabulary
+                asr_model2.change_vocabulary(new_tokenizer_dir, new_tokenizer_type='wpe')
+                assert asr_model.tokenizer.vocab_size != asr_model2.tokenizer.vocab_size
+
+                new_save_path = os.path.join(save_dir, 'temp2.nemo')
+                asr_model2.save_to(new_save_path)
+
+                asr_model3 = EncDecCTCModelBPE.restore_from(new_save_path)
+                assert isinstance(asr_model3, EncDecCTCModelBPE)
+
+                # Check if vocabulary size is same
+                assert asr_model2.tokenizer.tokenizer.vocab_size == asr_model3.tokenizer.tokenizer.vocab_size
+                assert asr_model2.vocab_path != asr_model3.vocab_path
+
+                # Model PT level checks
+                assert len(asr_model2.artifacts) == 1
+
+            finally:
+                os.rename(old_tokenizer_dir + '.bkp', old_tokenizer_dir)
+
+    @pytest.mark.unit
+    def test_decoding_change(self, asr_model):
+        assert asr_model.decoding is not None
+        assert isinstance(asr_model.decoding, CTCBPEDecoding)
+        assert asr_model.decoding.cfg.strategy == "greedy_batch"
+        assert asr_model.decoding.preserve_alignments is False
+        assert asr_model.decoding.compute_timestamps is False
+
+        cfg = CTCBPEDecodingConfig(preserve_alignments=True, compute_timestamps=True)
+        asr_model.change_decoding_strategy(cfg)
+
+        assert asr_model.decoding.preserve_alignments is True
+        assert asr_model.decoding.compute_timestamps is True
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'beam'
+        new_strategy.beam = DictConfig({'beam_size': 1})
+        asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(asr_model.decoding.decoding, beam_decode.BeamCTCInfer)
+        assert asr_model.decoding.decoding.search_type == "default"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'pyctcdecode'
+        new_strategy.beam = DictConfig({'beam_size': 1})
+        asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(asr_model.decoding.decoding, beam_decode.BeamCTCInfer)
+        assert asr_model.decoding.decoding.search_type == "pyctcdecode"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'flashlight'
+        new_strategy.beam = DictConfig({'beam_size': 1})
+        asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(asr_model.decoding.decoding, beam_decode.BeamCTCInfer)
+        assert asr_model.decoding.decoding.search_type == "flashlight"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'wfst'
+        new_strategy.beam = DictConfig({'beam_size': 1})
+        asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(asr_model.decoding.decoding, beam_decode.WfstCTCInfer)
+        assert asr_model.decoding.decoding.search_type == "riva"
+
+    @pytest.mark.unit
+    def test_ASRDatasetConfig_for_AudioToBPEDataset(self):
+        # ignore some additional arguments as dataclass is generic
+        IGNORE_ARGS = [
+            'is_tarred',
+            'num_workers',
+            'batch_size',
+            'tarred_audio_filepaths',
+            'shuffle',
+            'pin_memory',
+            'drop_last',
+            'tarred_shard_strategy',
+            'shard_manifests',
+            'shuffle_n',
+            'parser',
+            'normalize',
+            'unk_index',
+            'pad_id',
+            'bos_id',
+            'eos_id',
+            'blank_index',
+            'bucketing_batch_size',
+            'bucketing_strategy',
+            'bucketing_weights',
+            'channel_selector',
+            'use_lhotse',
+            'tarred_random_access',
+            'use_bucketing',
+            'batch_duration',
+            'quadratic_duration',
+            'bucket_batch_size',
+            'bucket_duration_bins',
+            'num_buckets',
+            'pin_memory',
+        ]
+
+        REMAP_ARGS = {'trim_silence': 'trim', 'labels': 'tokenizer'}
+
+        result = assert_dataclass_signature_match(
+            audio_to_text.AudioToBPEDataset,
+            configs.ASRDatasetConfig,
+            ignore_args=IGNORE_ARGS,
+            remap_args=REMAP_ARGS,
+        )
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_ASRDatasetConfig_for_TarredAudioToBPEDataset(self):
+        # ignore some additional arguments as dataclass is generic
+        IGNORE_ARGS = [
+            'is_tarred',
+            'num_workers',
+            'batch_size',
+            'shuffle',
+            'pin_memory',
+            'drop_last',
+            'parser',
+            'normalize',
+            'unk_index',
+            'pad_id',
+            'bos_id',
+            'eos_id',
+            'blank_index',
+            'global_rank',
+            'world_size',
+            'bucketing_batch_size',
+            'bucketing_strategy',
+            'bucketing_weights',
+            'max_utts',
+            'use_lhotse',
+            'tarred_random_access',
+            'use_bucketing',
+            'batch_duration',
+            'quadratic_duration',
+            'bucket_batch_size',
+            'bucket_duration_bins',
+            'num_buckets',
+            'pin_memory',
+        ]
+
+        REMAP_ARGS = {
+            'trim_silence': 'trim',
+            'tarred_audio_filepaths': 'audio_tar_filepaths',
+            'tarred_shard_strategy': 'shard_strategy',
+            'shuffle_n': 'shuffle',
+            'labels': 'tokenizer',
+        }
+
+        result = assert_dataclass_signature_match(
+            audio_to_text.TarredAudioToBPEDataset,
+            configs.ASRDatasetConfig,
+            ignore_args=IGNORE_ARGS,
+            remap_args=REMAP_ARGS,
+        )
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None

tests/collections/asr/test_asr_ctcencdec_model.py

@@ -0,0 +1,365 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import copy
+
+import pytest
+import torch
+from lhotse import CutSet, MonoCut
+from lhotse.testing.dummies import DummyManifest
+from omegaconf import DictConfig, OmegaConf, open_dict
+
+import nemo.collections.asr as nemo_asr
+from nemo.collections.asr.data import audio_to_text
+from nemo.collections.asr.data.audio_to_text_lhotse import LhotseSpeechToTextBpeDataset
+from nemo.collections.asr.models import EncDecCTCModel, configs
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecoding, CTCDecodingConfig
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.collections.common.parts.preprocessing.parsers import make_parser
+from nemo.utils.config_utils import assert_dataclass_signature_match, update_model_config
+
+
+@pytest.fixture()
+def asr_model():
+    preprocessor = {'_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor'}
+    encoder = {
+        '_target_': 'nemo.collections.asr.modules.ConvASREncoder',
+        'feat_in': 64,
+        'activation': 'relu',
+        'conv_mask': True,
+        'jasper': [
+            {
+                'filters': 1024,
+                'repeat': 1,
+                'kernel': [1],
+                'stride': [1],
+                'dilation': [1],
+                'dropout': 0.0,
+                'residual': False,
+                'separable': True,
+                'se': True,
+                'se_context_size': -1,
+            }
+        ],
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+        'feat_in': 1024,
+        'num_classes': 28,
+        'vocabulary': [
+            ' ',
+            'a',
+            'b',
+            'c',
+            'd',
+            'e',
+            'f',
+            'g',
+            'h',
+            'i',
+            'j',
+            'k',
+            'l',
+            'm',
+            'n',
+            'o',
+            'p',
+            'q',
+            'r',
+            's',
+            't',
+            'u',
+            'v',
+            'w',
+            'x',
+            'y',
+            'z',
+            "'",
+        ],
+    }
+    modelConfig = DictConfig(
+        {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
+    )
+
+    model_instance = EncDecCTCModel(cfg=modelConfig)
+    return model_instance
+
+
+class TestEncDecCTCModel:
+    @pytest.mark.unit
+    def test_constructor(self, asr_model):
+        asr_model.train()
+        # TODO: make proper config and assert correct number of weights
+        # Check to/from config_dict:
+        confdict = asr_model.to_config_dict()
+        instance2 = EncDecCTCModel.from_config_dict(confdict)
+        assert isinstance(instance2, EncDecCTCModel)
+
+    @pytest.mark.unit
+    def test_forward(self, asr_model):
+        asr_model = asr_model.eval()
+
+        asr_model.preprocessor.featurizer.dither = 0.0
+        asr_model.preprocessor.featurizer.pad_to = 0
+
+        input_signal = torch.randn(size=(4, 512))
+        length = torch.randint(low=321, high=500, size=[4])
+
+        with torch.no_grad():
+            # batch size 1
+            logprobs_instance = []
+            for i in range(input_signal.size(0)):
+                logprobs_ins, _, _ = asr_model.forward(
+                    input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1]
+                )
+                logprobs_instance.append(logprobs_ins)
+                print(len(logprobs_ins))
+            logprobs_instance = torch.cat(logprobs_instance, 0)
+
+            # batch size 4
+            logprobs_batch, _, _ = asr_model.forward(input_signal=input_signal, input_signal_length=length)
+
+        assert logprobs_instance.shape == logprobs_batch.shape
+        diff = torch.mean(torch.abs(logprobs_instance - logprobs_batch))
+        assert diff <= 1e-6
+        diff = torch.max(torch.abs(logprobs_instance - logprobs_batch))
+        assert diff <= 1e-6
+
+    @pytest.mark.unit
+    def test_predict_step(self, asr_model):
+        token_list = [" ", "a", "b", "c"]
+        asr_model = asr_model.eval()
+        cuts = DummyManifest(CutSet, begin_id=0, end_id=1, with_data=True)
+        dataset = LhotseSpeechToTextBpeDataset(tokenizer=make_parser(labels=token_list), return_cuts=True)
+        batch = dataset[cuts]
+        outputs = asr_model.predict_step(batch, 0)
+        assert len(outputs) == 1
+        assert len(outputs[0]) == 2
+        assert isinstance(outputs[0][0], MonoCut)
+        assert isinstance(outputs[0][1], Hypothesis)
+
+    @pytest.mark.unit
+    def test_vocab_change(self, asr_model):
+        old_vocab = copy.deepcopy(asr_model.decoder.vocabulary)
+        nw1 = asr_model.num_weights
+        asr_model.change_vocabulary(new_vocabulary=old_vocab)
+        # No change
+        assert nw1 == asr_model.num_weights
+        new_vocab = copy.deepcopy(old_vocab)
+        new_vocab.append('!')
+        new_vocab.append('$')
+        new_vocab.append('@')
+        asr_model.change_vocabulary(new_vocabulary=new_vocab)
+        # fully connected + bias
+        assert asr_model.num_weights == nw1 + 3 * (asr_model.decoder._feat_in + 1)
+
+    @pytest.mark.unit
+    def test_decoding_change(self, asr_model):
+        assert asr_model.decoding is not None
+        assert isinstance(asr_model.decoding, CTCDecoding)
+        assert asr_model.decoding.cfg.strategy == "greedy_batch"
+        assert asr_model.decoding.preserve_alignments is False
+        assert asr_model.decoding.compute_timestamps is False
+
+        cfg = CTCDecodingConfig(preserve_alignments=True, compute_timestamps=True)
+        asr_model.change_decoding_strategy(cfg)
+
+        assert asr_model.decoding.preserve_alignments is True
+        assert asr_model.decoding.compute_timestamps is True
+
+    @pytest.mark.unit
+    def test_change_conv_asr_se_context_window(self, asr_model):
+        old_cfg = copy.deepcopy(asr_model.cfg)
+        asr_model.change_conv_asr_se_context_window(context_window=32)  # 32 * 0.01s context
+        new_config = asr_model.cfg
+
+        assert old_cfg.encoder.jasper[0].se_context_size == -1
+        assert new_config.encoder.jasper[0].se_context_size == 32
+
+        for name, m in asr_model.encoder.named_modules():
+            if type(m).__class__.__name__ == 'SqueezeExcite':
+                assert m.context_window == 32
+
+    @pytest.mark.unit
+    def test_change_conv_asr_se_context_window_no_config_update(self, asr_model):
+        old_cfg = copy.deepcopy(asr_model.cfg)
+        asr_model.change_conv_asr_se_context_window(context_window=32, update_config=False)  # 32 * 0.01s context
+        new_config = asr_model.cfg
+
+        assert old_cfg.encoder.jasper[0].se_context_size == -1
+        assert new_config.encoder.jasper[0].se_context_size == -1  # no change
+
+        for name, m in asr_model.encoder.named_modules():
+            if type(m).__class__.__name__ == 'SqueezeExcite':
+                assert m.context_window == 32
+
+    @pytest.mark.unit
+    def test_dataclass_instantiation(self, asr_model):
+        model_cfg = configs.EncDecCTCModelConfig()
+
+        # Update mandatory values
+        vocabulary = asr_model.decoder.vocabulary
+        model_cfg.model.labels = vocabulary
+
+        # Update encoder
+        model_cfg.model.encoder.activation = 'relu'
+        model_cfg.model.encoder.feat_in = 64
+        model_cfg.model.encoder.jasper = [
+            nemo_asr.modules.conv_asr.JasperEncoderConfig(
+                filters=1024,
+                repeat=1,
+                kernel=[1],
+                stride=[1],
+                dilation=[1],
+                dropout=0.0,
+                residual=False,
+                se=True,
+                se_context_size=-1,
+            )
+        ]
+
+        # Update decoder
+        model_cfg.model.decoder.feat_in = 1024
+        model_cfg.model.decoder.num_classes = 28
+        model_cfg.model.decoder.vocabulary = vocabulary
+
+        # Construct the model
+        asr_cfg = OmegaConf.create({'model': asr_model.cfg})
+        model_cfg_v1 = update_model_config(model_cfg, asr_cfg)
+        new_model = EncDecCTCModel(cfg=model_cfg_v1.model)
+
+        assert new_model.num_weights == asr_model.num_weights
+        # trainer and exp manager should be there
+        # assert 'trainer' in model_cfg_v1
+        # assert 'exp_manager' in model_cfg_v1
+        # datasets and optim/sched should not be there after ModelPT.update_model_dataclass()
+        assert 'train_ds' not in model_cfg_v1.model
+        assert 'validation_ds' not in model_cfg_v1.model
+        assert 'test_ds' not in model_cfg_v1.model
+        assert 'optim' not in model_cfg_v1.model
+
+        # Construct the model, without dropping additional keys
+        asr_cfg = OmegaConf.create({'model': asr_model.cfg})
+        model_cfg_v2 = update_model_config(model_cfg, asr_cfg, drop_missing_subconfigs=False)
+
+        # Assert all components are in config
+        # assert 'trainer' in model_cfg_v2
+        # assert 'exp_manager' in model_cfg_v2
+        assert 'train_ds' in model_cfg_v2.model
+        assert 'validation_ds' in model_cfg_v2.model
+        assert 'test_ds' in model_cfg_v2.model
+        assert 'optim' in model_cfg_v2.model
+
+        # Remove extra components (optim and sched can be kept without issue)
+        with open_dict(model_cfg_v2.model):
+            model_cfg_v2.model.pop('train_ds')
+            model_cfg_v2.model.pop('validation_ds')
+            model_cfg_v2.model.pop('test_ds')
+
+        new_model = EncDecCTCModel(cfg=model_cfg_v2.model)
+
+        assert new_model.num_weights == asr_model.num_weights
+        # trainer and exp manager should be there
+
+    @pytest.mark.unit
+    def test_ASRDatasetConfig_for_AudioToCharDataset(self):
+        # ignore some additional arguments as dataclass is generic
+        IGNORE_ARGS = [
+            'is_tarred',
+            'num_workers',
+            'batch_size',
+            'tarred_audio_filepaths',
+            'shuffle',
+            'pin_memory',
+            'drop_last',
+            'tarred_shard_strategy',
+            'shard_manifests',
+            'shuffle_n',
+            'use_start_end_token',
+            'use_start_end_token',
+            'bucketing_batch_size',
+            'bucketing_strategy',
+            'bucketing_weights',
+            'channel_selector',
+            'use_lhotse',
+            'tarred_random_access',
+            'use_bucketing',
+            'batch_duration',
+            'quadratic_duration',
+            'bucket_batch_size',
+            'bucket_duration_bins',
+            'num_buckets',
+            'pin_memory',
+        ]
+
+        REMAP_ARGS = {'trim_silence': 'trim'}
+
+        result = assert_dataclass_signature_match(
+            audio_to_text.AudioToCharDataset,
+            configs.ASRDatasetConfig,
+            ignore_args=IGNORE_ARGS,
+            remap_args=REMAP_ARGS,
+        )
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_ASRDatasetConfig_for_TarredAudioToCharDataset(self):
+        # ignore some additional arguments as dataclass is generic
+        IGNORE_ARGS = [
+            'is_tarred',
+            'num_workers',
+            'batch_size',
+            'shuffle',
+            'pin_memory',
+            'drop_last',
+            'global_rank',
+            'world_size',
+            'use_start_end_token',
+            'bucketing_batch_size',
+            'bucketing_strategy',
+            'bucketing_weights',
+            'max_utts',
+            'use_lhotse',
+            'tarred_random_access',
+            'use_bucketing',
+            'batch_duration',
+            'quadratic_duration',
+            'bucket_batch_size',
+            'bucket_duration_bins',
+            'num_buckets',
+            'pin_memory',
+        ]
+
+        REMAP_ARGS = {
+            'trim_silence': 'trim',
+            'tarred_audio_filepaths': 'audio_tar_filepaths',
+            'tarred_shard_strategy': 'shard_strategy',
+            'shuffle_n': 'shuffle',
+        }
+
+        result = assert_dataclass_signature_match(
+            audio_to_text.TarredAudioToCharDataset,
+            configs.ASRDatasetConfig,
+            ignore_args=IGNORE_ARGS,
+            remap_args=REMAP_ARGS,
+        )
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None

tests/collections/asr/test_asr_datasets.py

@@ -0,0 +1,866 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import copy
+import filecmp
+import json
+import os
+import shutil
+import tempfile
+from unittest import mock
+
+import numpy as np
+import pytest
+import soundfile as sf
+import torch.cuda
+from omegaconf import DictConfig, OmegaConf
+from torch.utils.data import DataLoader
+
+from nemo.collections.asr.data import audio_to_text_dataset
+from nemo.collections.asr.data.audio_to_text import (
+    DataStoreObject,
+    TarredAudioToBPEDataset,
+    TarredAudioToCharDataset,
+    cache_datastore_manifests,
+)
+from nemo.collections.asr.data.audio_to_text_dali import (
+    __DALI_MINIMUM_VERSION__,
+    AudioToBPEDALIDataset,
+    AudioToCharDALIDataset,
+    is_dali_supported,
+)
+from nemo.collections.asr.data.audio_to_text_dataset import inject_dataloader_value_from_model_config
+from nemo.collections.asr.data.feature_to_text import FeatureToBPEDataset, FeatureToCharDataset
+from nemo.collections.asr.models.ctc_models import EncDecCTCModel
+from nemo.collections.asr.parts.utils.manifest_utils import write_manifest
+from nemo.collections.common import tokenizers
+from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config
+from nemo.utils import logging
+
+try:
+    HAVE_DALI = is_dali_supported(__DALI_MINIMUM_VERSION__)
+except (ImportError, ModuleNotFoundError):
+    HAVE_DALI = False
+
+
+def decode_chars(tokens, token_length, mapping):
+    text = []
+    tokens = tokens.cpu().numpy()
+    for idx in tokens:
+        text_token = mapping[idx]
+        text.append(text_token)
+
+    text = text[:token_length]
+    text = ''.join(text)
+    return text
+
+
+def decode_subwords(tokens, token_length, tokenizer: tokenizers.TokenizerSpec):
+    tokens = tokens.cpu().numpy()
+    tokens = tokens[:token_length]
+    text = tokenizer.ids_to_text(tokens)
+    return text
+
+
+class TestASRDatasets:
+    labels = [
+        " ",
+        "a",
+        "b",
+        "c",
+        "d",
+        "e",
+        "f",
+        "g",
+        "h",
+        "i",
+        "j",
+        "k",
+        "l",
+        "m",
+        "n",
+        "o",
+        "p",
+        "q",
+        "r",
+        "s",
+        "t",
+        "u",
+        "v",
+        "w",
+        "x",
+        "y",
+        "z",
+        "'",
+    ]
+
+    @pytest.mark.unit
+    def test_tarred_dataset(self, test_data_dir):
+        manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/tarred_audio_manifest.json'))
+
+        # Test braceexpand loading
+        tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar'))
+        ds_braceexpand = TarredAudioToCharDataset(
+            audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, sample_rate=16000
+        )
+        assert len(ds_braceexpand) == 32
+        count = 0
+        for _ in ds_braceexpand:
+            count += 1
+        assert count == 32
+
+        # Test loading via list
+        tarpath = [os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar')) for i in range(2)]
+        ds_list_load = TarredAudioToCharDataset(
+            audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, sample_rate=16000
+        )
+        count = 0
+        for _ in ds_list_load:
+            count += 1
+        assert count == 32
+
+    @pytest.mark.unit
+    def test_tarred_dataset_filter(self, test_data_dir):
+        """
+        Checks for
+            1. file count when manifest len is less than tarred dataset
+            2. Ignoring files in manifest that are not in tarred balls
+
+        """
+        manifest_path = os.path.abspath(
+            os.path.join(test_data_dir, 'asr/tarred_an4/tarred_duplicate_audio_manifest.json')
+        )
+
+        # Test braceexpand loading
+        tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar'))
+        ds_braceexpand = TarredAudioToCharDataset(
+            audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, sample_rate=16000
+        )
+        assert len(ds_braceexpand) == 6
+        count = 0
+        for _ in ds_braceexpand:
+            count += 1
+        assert count == 5  # file ending with sub is not part of tar ball
+
+    @pytest.mark.unit
+    def test_mismatch_in_model_dataloader_config(self, caplog):
+        logging._logger.propagate = True
+        caplog.set_level(logging.WARNING)
+
+        model_cfg = OmegaConf.create(dict(labels=OmegaConf.create(["a", "b", "c"])))
+        dataloader_cfg = OmegaConf.create(dict(labels=copy.deepcopy(self.labels)))
+
+        inject_dataloader_value_from_model_config(model_cfg, dataloader_cfg, key='labels')
+
+        assert (
+            """`labels` is explicitly provided to the data loader, and is different from the `labels` provided at the model level config."""
+            in caplog.text
+        )
+
+        logging._logger.propagate = False
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_tarred_bpe_dataset(self, test_data_dir):
+        manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/tarred_audio_manifest.json'))
+
+        tokenizer_path = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
+        tokenizer = tokenizers.AutoTokenizer(pretrained_model_name='bert-base-cased', vocab_file=tokenizer_path)
+
+        # Test braceexpand loading
+        tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar'))
+        ds_braceexpand = TarredAudioToBPEDataset(
+            audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, tokenizer=tokenizer, sample_rate=16000
+        )
+        assert len(ds_braceexpand) == 32
+        count = 0
+        for _ in ds_braceexpand:
+            count += 1
+        assert count == 32
+
+        # Test loading via list
+        tarpath = [os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar')) for i in range(2)]
+        ds_list_load = TarredAudioToBPEDataset(
+            audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, tokenizer=tokenizer, sample_rate=16000
+        )
+        count = 0
+        for _ in ds_list_load:
+            count += 1
+        assert count == 32
+
+    @pytest.mark.skipif(not HAVE_DALI, reason="NVIDIA DALI is not installed or incompatible version")
+    @pytest.mark.unit
+    def test_dali_char_dataset(self, test_data_dir):
+        manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/an4_val.json'))
+
+        num_samples = 10
+        batch_size = 2
+        device = 'gpu' if torch.cuda.is_available() else 'cpu'
+        texts = []
+
+        with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8') as f:
+            with open(manifest_path, 'r', encoding='utf-8') as m:
+                for ix, line in enumerate(m):
+                    if ix >= num_samples:
+                        break
+
+                    line = line.replace("tests/data/", "tests/.data/").replace("\n", "")
+                    f.write(f"{line}\n")
+
+                    data = json.loads(line)
+                    texts.append(data['text'])
+
+            f.seek(0)
+
+            dataset = AudioToCharDALIDataset(
+                manifest_filepath=f.name,
+                device=device,
+                batch_size=batch_size,
+                labels=self.labels,
+                max_duration=16.0,
+                parser='en',
+                shuffle=False,
+            )
+
+            assert len(dataset) == (num_samples // batch_size)  # num batches
+            count = 0
+            original_transcripts = []
+            for batch in dataset:
+                transcripts = batch[2]  # transcript index in DALIOutputs
+                transcripts_lengths = batch[3]  # transcript length index in DALIOutputs
+                transcripts = [
+                    decode_chars(transcript, transcripts_length, mapping=self.labels)
+                    for transcript, transcripts_length in zip(transcripts, transcripts_lengths)
+                ]
+                original_transcripts.extend(transcripts)
+                count += len(transcripts)
+            assert count == num_samples
+
+            # Assert transcripts are correct
+            for text, og_transcript in zip(texts, original_transcripts):
+                assert text == og_transcript
+
+            # Repeat, now with shuffle enabled
+            f.seek(0)
+
+            dataset = AudioToCharDALIDataset(
+                manifest_filepath=f.name,
+                device=device,
+                batch_size=batch_size,
+                labels=self.labels,
+                max_duration=16.0,
+                parser='en',
+                shuffle=True,
+            )
+
+            assert len(dataset) == (num_samples // batch_size)  # num batches
+            count = 0
+            shuffled_transcripts = []
+            for batch in dataset:
+                transcripts = batch[2]  # transcript index in DALIOutputs
+                transcripts_lengths = batch[3]  # transcript length index in DALIOutputs
+                transcripts = [
+                    decode_chars(transcript, transcripts_length, mapping=self.labels)
+                    for transcript, transcripts_length in zip(transcripts, transcripts_lengths)
+                ]
+                shuffled_transcripts.extend(transcripts)
+                count += len(transcripts)
+            assert count == num_samples
+
+            samples_changed = 0
+            for orig, shuffled in zip(original_transcripts, shuffled_transcripts):
+                if orig != shuffled:
+                    samples_changed += 1
+            assert samples_changed > 1  # assume after shuffling at least 1 sample was displaced
+
+            for og_transcript, shuffled_transcript in zip(sorted(original_transcripts), sorted(shuffled_transcripts)):
+                assert og_transcript == shuffled_transcript
+
+    @pytest.mark.skipif(not HAVE_DALI, reason="NVIDIA DALI is not installed or incompatible version")
+    @pytest.mark.unit
+    def test_dali_bpe_dataset(self, test_data_dir):
+        manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/an4_val.json'))
+
+        num_samples = 10
+        batch_size = 2
+        device = 'gpu' if torch.cuda.is_available() else 'cpu'
+        texts = []
+
+        tokenizer_path = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
+        tokenizer = tokenizers.AutoTokenizer(pretrained_model_name='bert-base-cased', vocab_file=tokenizer_path)
+
+        with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8') as f:
+            with open(manifest_path, 'r', encoding='utf-8') as m:
+                for ix, line in enumerate(m):
+                    if ix >= num_samples:
+                        break
+
+                    line = line.replace("tests/data/", "tests/.data/").replace("\n", "")
+                    f.write(f"{line}\n")
+
+                    data = json.loads(line)
+                    texts.append(data['text'])
+
+            f.seek(0)
+
+            dataset = AudioToBPEDALIDataset(
+                manifest_filepath=f.name,
+                tokenizer=tokenizer,
+                device=device,
+                batch_size=batch_size,
+                max_duration=16.0,
+                shuffle=False,
+            )
+
+            assert len(dataset) == (num_samples // batch_size)  # num batches
+            count = 0
+            original_transcripts = []
+            for batch in dataset:
+                transcripts = batch[2]  # transcript index in DALIOutputs
+                transcripts_lengths = batch[3]  # transcript length index in DALIOutputs
+                transcripts = [
+                    decode_subwords(transcript, transcripts_length, tokenizer=tokenizer)
+                    for transcript, transcripts_length in zip(transcripts, transcripts_lengths)
+                ]
+                original_transcripts.extend(transcripts)
+                count += len(transcripts)
+            assert count == num_samples
+
+            # Assert transcripts are correct
+            for text, og_transcript in zip(texts, original_transcripts):
+                assert text == og_transcript
+
+            # Repeat, now with shuffle enabled
+            f.seek(0)
+
+            dataset = AudioToBPEDALIDataset(
+                manifest_filepath=f.name,
+                tokenizer=tokenizer,
+                device=device,
+                batch_size=batch_size,
+                max_duration=16.0,
+                shuffle=True,
+            )
+
+            assert len(dataset) == (num_samples // batch_size)  # num batches
+            count = 0
+            shuffled_transcripts = []
+            for batch in dataset:
+                transcripts = batch[2]  # transcript index in DALIOutputs
+                transcripts_lengths = batch[3]  # transcript length index in DALIOutputs
+                transcripts = [
+                    decode_subwords(transcript, transcripts_length, tokenizer=tokenizer)
+                    for transcript, transcripts_length in zip(transcripts, transcripts_lengths)
+                ]
+                shuffled_transcripts.extend(transcripts)
+                count += len(transcripts)
+            assert count == num_samples
+
+            samples_changed = 0
+            for orig, shuffled in zip(original_transcripts, shuffled_transcripts):
+                if orig != shuffled:
+                    samples_changed += 1
+            assert samples_changed > 1  # assume after shuffling at least 1 sample was displaced
+
+            for og_transcript, shuffled_transcript in zip(sorted(original_transcripts), sorted(shuffled_transcripts)):
+                assert og_transcript == shuffled_transcript
+
+    @pytest.mark.xfail(
+        reason="DALI ASR Dataset's preprocessor is not patched with padding inconsistency fix (PR #13827)"
+    )
+    @pytest.mark.skipif(not HAVE_DALI, reason="NVIDIA DALI is not installed or incompatible version")
+    @pytest.mark.unit
+    def test_dali_char_vs_ref_dataset(self, test_data_dir):
+        manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/an4_val.json'))
+
+        num_samples = 10
+        batch_size = 1
+        texts = []
+
+        with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8') as f:
+            with open(manifest_path, 'r') as m:
+                for ix, line in enumerate(m):
+                    if ix >= num_samples:
+                        break
+
+                    line = line.replace("tests/data/", "tests/.data/").replace("\n", "")
+                    f.write(f"{line}\n")
+
+                    data = json.loads(line)
+                    texts.append(data['text'])
+
+            f.seek(0)
+
+            preprocessor = {
+                '_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor',
+                'dither': 0.0,
+            }
+            preprocessor_cfg = DictConfig(preprocessor)
+
+            dataset_cfg = {
+                'manifest_filepath': f.name,
+                'sample_rate': 16000,
+                'labels': self.labels,
+                'batch_size': batch_size,
+                'trim_silence': False,
+                'max_duration': 16.7,
+                'shuffle': False,
+                'is_tarred': False,
+            }
+            dali_dataset = audio_to_text_dataset.get_dali_char_dataset(
+                config=dataset_cfg,
+                shuffle=False,
+                device_id=0,
+                global_rank=0,
+                world_size=1,
+                preprocessor_cfg=preprocessor_cfg,
+            )
+            ref_dataset = audio_to_text_dataset.get_char_dataset(
+                config=dataset_cfg,
+            )
+            ref_dataloader = DataLoader(
+                dataset=ref_dataset,
+                batch_size=batch_size,
+                collate_fn=ref_dataset.collate_fn,
+                drop_last=False,
+                shuffle=False,
+                num_workers=0,
+                pin_memory=False,
+            )
+            ref_preprocessor = EncDecCTCModel.from_config_dict(preprocessor_cfg)
+
+            for ref_data, dali_data in zip(ref_dataloader, dali_dataset):
+                ref_audio, ref_audio_len, _, _ = ref_data
+                ref_features, ref_features_len = ref_preprocessor(input_signal=ref_audio, length=ref_audio_len)
+
+                dali_features, dali_features_len, _, _ = dali_data
+
+                a = ref_features.cpu().numpy()[:, :, :ref_features_len]
+                b = dali_features.cpu().numpy()[:, :, :dali_features_len]
+
+                err = np.abs(a - b)
+                assert np.mean(err) < 0.0001
+                assert np.max(err) < 0.01
+
+    @pytest.mark.skipif(not HAVE_DALI, reason="NVIDIA DALI is not installed or incompatible version")
+    @pytest.mark.unit
+    def test_tarred_dali_char_dataset(self, test_data_dir):
+        manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/tarred_audio_manifest.json'))
+        audio_tar_filepaths = [
+            os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{idx}.tar')) for idx in range(2)
+        ]
+        audio_tar_index_filepaths = [
+            os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/dali_index/audio_{idx}.index'))
+            for idx in range(2)
+        ]
+
+        batch_size = 8
+        device = 'gpu' if torch.cuda.is_available() else 'cpu'
+        texts = []
+
+        with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8') as f:
+            num_samples = 0
+            with open(manifest_path, 'r') as m:
+                num_samples = len(m.readlines())
+
+            dataset = AudioToCharDALIDataset(
+                manifest_filepath=manifest_path,
+                audio_tar_filepaths=audio_tar_filepaths,
+                audio_tar_index_filepaths=audio_tar_index_filepaths,
+                device=device,
+                batch_size=batch_size,
+                labels=self.labels,
+                max_duration=16.0,
+                parser='en',
+                shuffle=False,
+            )
+
+            assert len(dataset) == (num_samples // batch_size)  # num batches
+            count = 0
+            original_transcripts = []
+            for batch in dataset:
+                transcripts = batch[2]  # transcript index in DALIOutputs
+                transcripts_lengths = batch[3]  # transcript length index in DALIOutputs
+                transcripts = [
+                    decode_chars(transcript, transcripts_length, mapping=self.labels)
+                    for transcript, transcripts_length in zip(transcripts, transcripts_lengths)
+                ]
+                original_transcripts.extend(transcripts)
+                count += len(transcripts)
+            assert count == num_samples
+
+            # Assert transcripts are correct
+            for text, og_transcript in zip(texts, original_transcripts):
+                assert text == og_transcript
+
+            dataset = AudioToCharDALIDataset(
+                manifest_filepath=manifest_path,  # f.name,
+                audio_tar_filepaths=audio_tar_filepaths,
+                audio_tar_index_filepaths=audio_tar_index_filepaths,
+                device=device,
+                batch_size=batch_size,
+                labels=self.labels,
+                max_duration=16.0,
+                parser='en',
+                shuffle=True,
+            )
+
+            assert len(dataset) == (num_samples // batch_size)  # num batches
+            count = 0
+            shuffled_transcripts = []
+            for batch in dataset:
+                transcripts = batch[2]  # transcript index in DALIOutputs
+                transcripts_lengths = batch[3]  # transcript length index in DALIOutputs
+                transcripts = [
+                    decode_chars(transcript, transcripts_length, mapping=self.labels)
+                    for transcript, transcripts_length in zip(transcripts, transcripts_lengths)
+                ]
+                shuffled_transcripts.extend(transcripts)
+                count += len(transcripts)
+            assert count == num_samples
+
+            samples_changed = 0
+            for orig, shuffled in zip(original_transcripts, shuffled_transcripts):
+                if orig != shuffled:
+                    samples_changed += 1
+            assert samples_changed > 1  # assume after shuffling at least 1 sample was displaced
+
+            for og_transcript, shuffled_transcript in zip(sorted(original_transcripts), sorted(shuffled_transcripts)):
+                assert og_transcript == shuffled_transcript
+
+    @pytest.mark.skipif(not HAVE_DALI, reason="NVIDIA DALI is not installed or incompatible version")
+    @pytest.mark.unit
+    def test_dali_tarred_char_vs_ref_dataset(self, test_data_dir):
+        manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/tarred_audio_manifest.json'))
+        audio_tar_filepaths = [
+            os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{idx}.tar')) for idx in range(2)
+        ]
+        audio_tar_index_filepaths = [
+            os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/dali_index/audio_{idx}.index'))
+            for idx in range(2)
+        ]
+
+        batch_size = 8
+        texts = []
+
+        with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8') as f:
+            num_samples = 0
+            with open(manifest_path, 'r') as m:
+                for ix, line in enumerate(m):
+                    data = json.loads(line)
+                    texts.append(data['text'])
+                    num_samples = ix
+
+            preprocessor = {
+                '_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor',
+                'dither': 0.0,
+            }
+            preprocessor_cfg = DictConfig(preprocessor)
+
+            dataset_cfg = {
+                'manifest_filepath': f.name,
+                'tarred_audio_filepaths': audio_tar_filepaths,
+                'tarred_audio_index_filepaths': audio_tar_index_filepaths,
+                'sample_rate': 16000,
+                'labels': self.labels,
+                'batch_size': batch_size,
+                'trim_silence': False,
+                'max_duration': 16.7,
+                'shuffle': False,
+                'is_tarred': False,
+            }
+            dali_dataset = audio_to_text_dataset.get_dali_char_dataset(
+                config=dataset_cfg,
+                shuffle=False,
+                device_id=0,
+                global_rank=0,
+                world_size=1,
+                preprocessor_cfg=preprocessor_cfg,
+            )
+            ref_dataset = audio_to_text_dataset.get_tarred_dataset(
+                config=dataset_cfg, shuffle_n=0, global_rank=0, world_size=1
+            )
+            ref_dataloader = DataLoader(
+                dataset=ref_dataset,
+                batch_size=batch_size,
+                collate_fn=ref_dataset.collate_fn,
+                drop_last=False,
+                shuffle=False,
+                num_workers=0,
+                pin_memory=False,
+            )
+            ref_preprocessor = EncDecCTCModel.from_config_dict(preprocessor_cfg)
+
+            for ref_data, dali_data in zip(ref_dataloader, dali_dataset):
+                ref_audio, ref_audio_len, _, _ = ref_data
+                ref_features, ref_features_len = ref_preprocessor(input_signal=ref_audio, length=ref_audio_len)
+
+                dali_features, dali_features_len, _, _ = dali_data
+
+                a = ref_features.cpu().numpy()[:, :, :ref_features_len]
+                b = dali_features.cpu().numpy()[:, :, :dali_features_len]
+
+                err = np.abs(a - b)
+                assert np.mean(err) < 0.0001
+                assert np.max(err) < 0.01
+
+    @pytest.mark.unit
+    def test_feature_to_text_char_dataset(self):
+        num_samples = 5
+        golden_feat_shape = (80, 5)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manifest_path = os.path.join(tmpdir, 'manifest_input.json')
+            with open(manifest_path, 'w', encoding='utf-8') as fp:
+                for i in range(num_samples):
+                    feat_file = os.path.join(tmpdir, f"feat_{i}.pt")
+                    torch.save(torch.randn(80, 5), feat_file)
+                    entry = {'audio_filepath': "", 'feature_file': feat_file, 'duration': 100000, "text": "a b c"}
+                    fp.write(json.dumps(entry) + '\n')
+
+            dataset = FeatureToCharDataset(manifest_path, labels=self.labels)
+            cnt = 0
+            for item in dataset:
+                cnt += 1
+                feat = item[0]
+                token_len = item[3]
+                assert feat.shape == golden_feat_shape
+                assert torch.equal(token_len, torch.tensor(5))
+            assert cnt == num_samples
+
+    @pytest.mark.unit
+    def test_feature_to_text_bpe_dataset(self, test_data_dir):
+        num_samples = 5
+        golden_feat_shape = (80, 5)
+        tokenizer_path = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
+        tokenizer = tokenizers.AutoTokenizer(pretrained_model_name='bert-base-cased', vocab_file=tokenizer_path)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manifest_path = os.path.join(tmpdir, 'manifest_input.json')
+            with open(manifest_path, 'w', encoding='utf-8') as fp:
+                for i in range(num_samples):
+                    feat_file = os.path.join(tmpdir, f"feat_{i}.pt")
+                    torch.save(torch.randn(80, 5), feat_file)
+                    entry = {'audio_filepath': "", 'feature_file': feat_file, 'duration': 100000, "text": "a b c"}
+                    fp.write(json.dumps(entry) + '\n')
+
+            dataset = FeatureToBPEDataset(manifest_path, tokenizer=tokenizer)
+            cnt = 0
+            for item in dataset:
+                cnt += 1
+                feat = item[0]
+                token_len = item[3]
+                assert feat.shape == golden_feat_shape
+                assert torch.equal(token_len, torch.tensor(5))
+            assert cnt == num_samples
+
+    @pytest.mark.unit
+    def test_feature_with_rttm_to_text_char_dataset(self):
+        num_samples = 2
+        golden_feat_shape = (80, 10)
+        sample = torch.ones(80, 10)
+        masked_sample = sample * FeatureToCharDataset.ZERO_LEVEL_SPEC_DB_VAL
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manifest_path = os.path.join(tmpdir, 'manifest_input.json')
+            with open(manifest_path, 'w', encoding='utf-8') as fp:
+                feat_file = os.path.join(tmpdir, f"feat_0.pt")
+                torch.save(sample, feat_file)
+
+                rttm_file = os.path.join(tmpdir, f"rttm_0.rttm")
+                with open(rttm_file, "w") as fout:
+                    fout.write(f"SPEAKER <NA> 1 0 1 <NA> <NA> speech <NA> <NA>\n")
+
+                entry = {
+                    'audio_filepath': "",
+                    'feature_file': feat_file,
+                    'rttm_file': rttm_file,
+                    'duration': 100000,
+                    "text": "a b c",
+                }
+                fp.write(json.dumps(entry) + '\n')
+
+                # second sample where all frames are not masked
+                feat_file = os.path.join(tmpdir, f"feat_1.pt")
+                torch.save(sample, feat_file)
+
+                rttm_file = os.path.join(tmpdir, f"rttm_1.rttm")
+                with open(rttm_file, "w") as fout:
+                    fout.write(f"SPEAKER <NA> 1 0 0 <NA> <NA> speech <NA> <NA>\n")
+
+                entry = {
+                    'audio_filepath': "",
+                    'feature_file': feat_file,
+                    'rttm_file': rttm_file,
+                    'duration': 100000,
+                    "text": "a b c",
+                }
+                fp.write(json.dumps(entry) + '\n')
+
+            dataset = FeatureToCharDataset(manifest_path, labels=self.labels, normalize=None, use_rttm=True)
+            cnt = 0
+            for item in dataset:
+                cnt += 1
+                feat = item[0]
+                token_len = item[3]
+                assert feat.shape == golden_feat_shape
+                assert torch.equal(token_len, torch.tensor(5))
+
+                if cnt == 1:
+                    assert torch.equal(feat, sample)
+                else:
+                    assert torch.equal(feat, masked_sample)
+
+            assert cnt == num_samples
+
+    @pytest.mark.unit
+    def test_feature_with_rttm_to_text_bpe_dataset(self, test_data_dir):
+        tokenizer_path = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
+        tokenizer = tokenizers.AutoTokenizer(pretrained_model_name='bert-base-cased', vocab_file=tokenizer_path)
+        num_samples = 2
+        golden_feat_shape = (80, 10)
+        sample = torch.ones(80, 10)
+        masked_sample = sample * FeatureToCharDataset.ZERO_LEVEL_SPEC_DB_VAL
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manifest_path = os.path.join(tmpdir, 'manifest_input.json')
+            with open(manifest_path, 'w', encoding='utf-8') as fp:
+                feat_file = os.path.join(tmpdir, f"feat_0.pt")
+                torch.save(sample, feat_file)
+
+                rttm_file = os.path.join(tmpdir, f"rttm_0.rttm")
+                with open(rttm_file, "w") as fout:
+                    fout.write(f"SPEAKER <NA> 1 0 1 <NA> <NA> speech <NA> <NA>\n")
+
+                entry = {
+                    'audio_filepath': "",
+                    'feature_file': feat_file,
+                    'rttm_file': rttm_file,
+                    'duration': 100000,
+                    "text": "a b c",
+                }
+                fp.write(json.dumps(entry) + '\n')
+
+                # second sample where all frames are not masked
+                feat_file = os.path.join(tmpdir, f"feat_1.pt")
+                torch.save(sample, feat_file)
+
+                rttm_file = os.path.join(tmpdir, f"rttm_1.rttm")
+                with open(rttm_file, "w") as fout:
+                    fout.write(f"SPEAKER <NA> 1 0 0 <NA> <NA> speech <NA> <NA>\n")
+
+                entry = {
+                    'audio_filepath': "",
+                    'feature_file': feat_file,
+                    'rttm_file': rttm_file,
+                    'duration': 100000,
+                    "text": "a b c",
+                }
+                fp.write(json.dumps(entry) + '\n')
+
+            dataset = FeatureToBPEDataset(manifest_path, tokenizer=tokenizer, normalize=None, use_rttm=True)
+            cnt = 0
+            for item in dataset:
+                cnt += 1
+                feat = item[0]
+                token_len = item[3]
+                assert feat.shape == golden_feat_shape
+                assert torch.equal(token_len, torch.tensor(5))
+
+                if cnt == 1:
+                    assert torch.equal(feat, sample)
+                else:
+                    assert torch.equal(feat, masked_sample)
+
+            assert cnt == num_samples
+
+
+class TestUtilityFunctions:
+    @pytest.mark.unit
+    @pytest.mark.parametrize('cache_audio', [False, True])
+    def test_cache_datastore_manifests(self, cache_audio: bool):
+        """Test caching of manifest and audio files."""
+        # Data setup
+        random_seed = 42
+        sample_rate = 16000
+        num_examples = 10
+        num_manifests = 2
+        data_duration = 1.0
+
+        # Generate random signals
+        _rng = np.random.default_rng(seed=random_seed)
+
+        # Input and target signals have the same duration
+        data_duration_samples = int(data_duration * sample_rate)
+
+        with tempfile.TemporaryDirectory() as test_dir:
+            test_store_dir = os.path.join(test_dir, 'store')
+            os.mkdir(test_store_dir)
+
+            # Prepare metadata and audio files
+            manifest_filepaths = []
+            audio_files = []
+            for m in range(num_manifests):
+                manifest_dir = os.path.join(test_store_dir, f'manifest_{m}')
+                os.mkdir(manifest_dir)
+                manifest_filepath = os.path.join(manifest_dir, 'manifest.json')
+
+                metadata = []
+                data = _rng.uniform(low=-0.5, high=0.5, size=(data_duration_samples, num_examples))
+                for n in range(num_examples):
+                    audio_filepath = f'manifest_{m}_audio_{n:02d}.wav'
+                    audio_file = os.path.join(manifest_dir, audio_filepath)
+                    # Write audio file
+                    sf.write(audio_file, data[:, n], sample_rate, 'float')
+                    # Update metadata
+                    metadata.append(
+                        {
+                            'audio_filepath': audio_filepath,
+                            'duration': data_duration,
+                            'text': f'text for example {n:02d}',
+                        }
+                    )
+                    # Update audio files
+                    audio_files.append(audio_file)
+
+                # Save manifest
+                write_manifest(manifest_filepath, metadata)
+                manifest_filepaths.append(manifest_filepath)
+
+            # Cache location
+            test_cache_dir = os.path.join(test_dir, 'cache')
+
+            # Instead of using AIS, copy object from store dir to cache dir
+            def fake_get(self):
+                # Object path relative to store path
+                object_path = os.path.relpath(self.store_path, start=test_store_dir)
+                # Copy to fake local path
+                self._local_path = os.path.join(test_cache_dir, object_path)
+                os.makedirs(os.path.dirname(self.local_path), exist_ok=True)
+                shutil.copy(self.store_path, self.local_path)
+                # Return path as in the original get
+                return self.local_path
+
+            with (
+                mock.patch('nemo.collections.asr.data.audio_to_text.is_datastore_path', lambda x: True),
+                mock.patch.object(DataStoreObject, 'get', fake_get),
+            ):
+                # Use a single worker for this test to avoid failure with mock & multiprocessing (#5607)
+                cache_datastore_manifests(manifest_filepaths, cache_audio=cache_audio, num_workers=1)
+
+            # Manifests need to be compared
+            store_files_to_compare = manifest_filepaths
+            if cache_audio:
+                # Audio needs to be compared
+                store_files_to_compare += audio_files
+
+            # Compare files
+            for f_store in store_files_to_compare:
+                f_cache = os.path.join(test_cache_dir, os.path.relpath(f_store, test_store_dir))
+                assert filecmp.cmp(f_store, f_cache, shallow=False), f'Files {f_store} and {f_cache} do not match.'

tests/collections/asr/test_asr_eou.py

@@ -0,0 +1,91 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+import numpy as np
+import pytest
+
+from nemo.collections.asr.parts.utils.eou_utils import EOUResult, cal_eou_metrics_from_frame_labels
+
+
+def make_eou_frame_labels(duration: float, eou_time: float, frame_len_in_secs: float = 0.08) -> List[float]:
+    """
+    Make EOU frame labels.
+    Args:
+        duration (float): Duration of the audio in seconds.
+        eou_time (float): Time of the EOU in seconds.
+        frame_len_in_secs (float): Length of each frame in seconds.
+    Returns:
+        List[float]: List of EOU frame labels.
+    """
+    if eou_time < 0 or eou_time > duration:
+        raise ValueError(f"EOU time ({eou_time}) is out of range for duration ({duration}).")
+
+    labels = [0] * int(np.ceil(duration / frame_len_in_secs) + 1)
+    labels[int(np.ceil(eou_time / frame_len_in_secs))] = 1
+    return labels
+
+
+class TestEOUMetrics:
+    @pytest.mark.unit
+    def test_cal_eou_metrics_from_frame_labels(self):
+        duration = 1.6
+        eou_time = 0.64
+        frame_len_in_secs = 0.08
+        ref_labels = make_eou_frame_labels(duration, eou_time, frame_len_in_secs)
+
+        # Test case 1: Early cutoff
+        pred_eou_time = 0.32
+        preds = make_eou_frame_labels(duration, pred_eou_time, frame_len_in_secs)
+        eou_metrics: EOUResult = cal_eou_metrics_from_frame_labels(
+            prediction=preds, reference=ref_labels, frame_len_in_secs=frame_len_in_secs
+        )
+        assert eou_metrics.true_positives == 0
+        assert eou_metrics.false_positives == 1
+        assert eou_metrics.false_negatives == 0
+        assert eou_metrics.num_utterances == 1
+        assert eou_metrics.num_predictions == 1
+        assert eou_metrics.missing == 0
+        assert eou_metrics.latency == []
+        assert np.isclose(eou_metrics.early_cutoff, [0.32])
+
+        # Test case 2: Latency
+        pred_eou_time = 0.96
+        preds = make_eou_frame_labels(duration, pred_eou_time, frame_len_in_secs)
+        eou_metrics: EOUResult = cal_eou_metrics_from_frame_labels(
+            prediction=preds, reference=ref_labels, frame_len_in_secs=frame_len_in_secs
+        )
+        assert eou_metrics.true_positives == 0
+        assert eou_metrics.false_positives == 0
+        assert eou_metrics.false_negatives == 1
+        assert eou_metrics.num_utterances == 1
+        assert eou_metrics.num_predictions == 1
+        assert eou_metrics.missing == 0
+        assert np.isclose(eou_metrics.latency, [0.32])
+        assert eou_metrics.early_cutoff == []
+
+        # Test case 3: miss detection
+        preds = [0] * len(ref_labels)
+        eou_metrics: EOUResult = cal_eou_metrics_from_frame_labels(
+            prediction=preds, reference=ref_labels, frame_len_in_secs=frame_len_in_secs
+        )
+        assert eou_metrics.true_positives == 0
+        assert eou_metrics.false_positives == 0
+        assert eou_metrics.false_negatives == 1
+        assert eou_metrics.num_utterances == 1
+        assert eou_metrics.num_predictions == 0
+        assert eou_metrics.missing == 1
+        assert eou_metrics.latency == []
+        assert eou_metrics.early_cutoff == []

tests/collections/asr/test_asr_exportables.py

@@ -0,0 +1,628 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import tempfile
+
+import onnx
+import pytest
+import torch.cuda
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+from nemo.collections.asr.models import (
+    EncDecClassificationModel,
+    EncDecCTCModel,
+    EncDecRNNTModel,
+    EncDecSpeakerLabelModel,
+)
+from nemo.collections.asr.parts.utils import asr_module_utils
+from nemo.collections.common.parts.adapter_modules import LinearAdapterConfig
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+
+NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+
+
+class TestExportable:
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    def test_EncDecCTCModel_export_to_onnx(self):
+        model_config = DictConfig(
+            {
+                'preprocessor': DictConfig(self.preprocessor),
+                'encoder': DictConfig(self.encoder_dict),
+                'decoder': DictConfig(self.decoder_dict),
+            }
+        )
+        model = EncDecCTCModel(cfg=model_config).cuda()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            filename = os.path.join(tmpdir, 'qn.onnx')
+            model.export(
+                output=filename,
+                check_trace=True,
+            )
+            onnx_model = onnx.load(filename)
+            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
+            assert onnx_model.graph.input[0].name == 'audio_signal'
+            assert onnx_model.graph.output[0].name == 'logprobs'
+
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    def test_EncDecClassificationModel_export_to_onnx(self, speech_classification_model):
+        model = speech_classification_model.cuda()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            filename = os.path.join(tmpdir, 'edc.onnx')
+            model.export(
+                output=filename,
+                check_trace=True,
+            )
+            onnx_model = onnx.load(filename)
+            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
+            assert onnx_model.graph.input[0].name == 'audio_signal'
+            assert onnx_model.graph.output[0].name == 'logits'
+
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    def test_EncDecSpeakerLabelModel_export_to_onnx(self, speaker_label_model):
+        model = speaker_label_model.cuda()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            filename = os.path.join(tmpdir, 'sl.onnx')
+            model.export(output=filename)
+            onnx_model = onnx.load(filename)
+            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
+            assert onnx_model.graph.input[0].name == 'audio_signal'
+            assert onnx_model.graph.output[0].name == 'logits'
+
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    def test_EncDecCitrinetModel_export_to_onnx(self, citrinet_model):
+        model = citrinet_model.cuda()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            filename = os.path.join(tmpdir, 'citri.onnx')
+            model.export(output=filename)
+            onnx_model = onnx.load(filename)
+            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
+            assert onnx_model.graph.input[0].name == 'audio_signal'
+            assert onnx_model.graph.input[1].name == 'length'
+            assert onnx_model.graph.output[0].name == 'logprobs'
+
+    @pytest.mark.pleasefixme
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    def test_ConformerModel_export_to_onnx(self, conformer_model):
+        model = conformer_model.cuda()
+        with tempfile.TemporaryDirectory() as tmpdir, torch.cuda.amp.autocast():
+            filename = os.path.join(tmpdir, 'conf.onnx')
+            device = next(model.parameters()).device
+            input_example = torch.randn(4, model.encoder._feat_in, 777, device=device)
+            input_example_length = torch.full(size=(input_example.shape[0],), fill_value=777, device=device)
+            model.export(
+                output=filename,
+                input_example=tuple([input_example, input_example_length]),
+                check_trace=True,
+            )
+
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    def test_EncDecCitrinetModel_limited_SE_export_to_onnx(self, citrinet_model):
+        model = citrinet_model.cuda()
+        asr_module_utils.change_conv_asr_se_context_window(model, context_window=24, update_config=False)
+
+        with tempfile.TemporaryDirectory() as tmpdir, torch.cuda.amp.autocast():
+            filename = os.path.join(tmpdir, 'citri_se.onnx')
+            model.export(
+                output=filename,
+                check_trace=True,
+            )
+            onnx_model = onnx.load(filename)
+            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
+            assert onnx_model.graph.input[0].name == 'audio_signal'
+            assert onnx_model.graph.input[1].name == 'length'
+            assert onnx_model.graph.output[0].name == 'logprobs'
+
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    def test_EncDecRNNTModel_export_to_onnx(self, citrinet_rnnt_model):
+        model = citrinet_rnnt_model.cuda()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            fn = 'citri_rnnt.onnx'
+            filename = os.path.join(tmpdir, fn)
+            files, descr = model.export(output=filename, verbose=False)
+
+            encoder_filename = os.path.join(tmpdir, 'encoder-' + fn)
+            assert files[0] == encoder_filename
+            assert os.path.exists(encoder_filename)
+            onnx_model = onnx.load(encoder_filename)
+            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
+            assert len(onnx_model.graph.input) == 2
+            assert len(onnx_model.graph.output) == 2
+            assert onnx_model.graph.input[0].name == 'audio_signal'
+            assert onnx_model.graph.input[1].name == 'length'
+            assert onnx_model.graph.output[0].name == 'outputs'
+            assert onnx_model.graph.output[1].name == 'encoded_lengths'
+
+            decoder_joint_filename = os.path.join(tmpdir, 'decoder_joint-' + fn)
+            assert files[1] == decoder_joint_filename
+            assert os.path.exists(decoder_joint_filename)
+            onnx_model = onnx.load(decoder_joint_filename)
+            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
+
+            input_examples = model.decoder.input_example()
+            assert type(input_examples[-1]) == tuple
+            num_states = len(input_examples[-1])
+            state_name = list(model.decoder.output_types.keys())[-1]
+
+            # enc_logits + (all decoder inputs - state tuple) + flattened state list
+            assert len(onnx_model.graph.input) == (1 + (len(input_examples) - 1) + num_states)
+            assert onnx_model.graph.input[0].name == 'encoder_outputs'
+            assert onnx_model.graph.input[1].name == 'targets'
+            assert onnx_model.graph.input[2].name == 'target_length'
+
+            if num_states > 0:
+                for idx, ip in enumerate(onnx_model.graph.input[3:]):
+                    assert ip.name == "input_" + state_name + '_' + str(idx + 1)
+
+            assert len(onnx_model.graph.output) == (len(input_examples) - 1) + num_states
+            assert onnx_model.graph.output[0].name == 'outputs'
+            assert onnx_model.graph.output[1].name == 'prednet_lengths'
+
+            if num_states > 0:
+                for idx, op in enumerate(onnx_model.graph.output[2:]):
+                    assert op.name == "output_" + state_name + '_' + str(idx + 1)
+
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    def test_EncDecRNNTModel_export_to_ts(self, citrinet_rnnt_model):
+        model = citrinet_rnnt_model.cuda()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            fn = 'citri_rnnt.ts'
+            filename = os.path.join(tmpdir, fn)
+            # Perform export + test with the input examples of the RNNT model.
+            files, descr = model.export(output=filename, verbose=False, check_trace=True)
+
+            encoder_filename = os.path.join(tmpdir, 'encoder-' + fn)
+            assert files[0] == encoder_filename
+            assert os.path.exists(encoder_filename)
+
+            ts_encoder = torch.jit.load(encoder_filename)
+            assert ts_encoder is not None
+
+            arguments = ts_encoder.forward.schema.arguments[1:]  # First value is `self`
+            assert arguments[0].name == 'audio_signal'
+            assert arguments[1].name == 'length'
+
+            decoder_joint_filename = os.path.join(tmpdir, 'decoder_joint-' + fn)
+            assert files[1] == decoder_joint_filename
+            assert os.path.exists(decoder_joint_filename)
+
+            ts_decoder_joint = torch.jit.load(decoder_joint_filename)
+            assert ts_decoder_joint is not None
+
+            ts_decoder_joint_args = ts_decoder_joint.forward.schema.arguments[1:]  # First value is self
+
+            input_examples = model.decoder.input_example()
+            assert type(input_examples[-1]) == tuple
+            num_states = len(input_examples[-1])
+            state_name = list(model.decoder.output_types.keys())[-1]
+
+            # enc_logits + (all decoder inputs - state tuple) + flattened state list
+            assert len(ts_decoder_joint_args) == (1 + (len(input_examples) - 1) + num_states)
+            assert ts_decoder_joint_args[0].name == 'encoder_outputs'
+            assert ts_decoder_joint_args[1].name == 'targets'
+            assert ts_decoder_joint_args[2].name == 'target_length'
+
+            if num_states > 0:
+                for idx, ip in enumerate(ts_decoder_joint_args[3:]):
+                    assert ip.name == "input_" + state_name + '_' + str(idx + 1)
+
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    def test_EncDecCTCModel_adapted_export_to_onnx(self):
+        model_config = DictConfig(
+            {
+                'preprocessor': DictConfig(self.preprocessor),
+                'encoder': DictConfig(self.encoder_dict),
+                'decoder': DictConfig(self.decoder_dict),
+            }
+        )
+
+        # support adapter in encoder
+        model_config.encoder.cls = model_config.encoder.cls + 'Adapter'  # ConvASREncoderAdapter
+
+        # load model
+        model = EncDecCTCModel(cfg=model_config)
+
+        # add adapter
+        adapter_cfg = OmegaConf.structured(
+            LinearAdapterConfig(in_features=model_config.encoder.params.jasper[0].filters, dim=32)
+        )
+        model.add_adapter('temp', cfg=adapter_cfg)
+
+        model = model.cuda()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            filename = os.path.join(tmpdir, 'qn.onnx')
+            model.export(
+                output=filename,
+                check_trace=True,
+            )
+            onnx_model = onnx.load(filename)
+            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
+            assert onnx_model.graph.input[0].name == 'audio_signal'
+            assert onnx_model.graph.output[0].name == 'logprobs'
+
+    def setup_method(self):
+        self.preprocessor = {
+            'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor',
+            'params': dict({}),
+        }
+
+        self.encoder_dict = {
+            'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+            'params': {
+                'feat_in': 64,
+                'activation': 'relu',
+                'conv_mask': True,
+                'jasper': [
+                    {
+                        'filters': 1024,
+                        'repeat': 1,
+                        'kernel': [1],
+                        'stride': [1],
+                        'dilation': [1],
+                        'dropout': 0.0,
+                        'residual': False,
+                        'separable': True,
+                        'se': True,
+                        'se_context_size': -1,
+                    }
+                ],
+            },
+        }
+
+        self.decoder_dict = {
+            'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
+            'params': {
+                'feat_in': 1024,
+                'num_classes': 28,
+                'vocabulary': [
+                    ' ',
+                    'a',
+                    'b',
+                    'c',
+                    'd',
+                    'e',
+                    'f',
+                    'g',
+                    'h',
+                    'i',
+                    'j',
+                    'k',
+                    'l',
+                    'm',
+                    'n',
+                    'o',
+                    'p',
+                    'q',
+                    'r',
+                    's',
+                    't',
+                    'u',
+                    'v',
+                    'w',
+                    'x',
+                    'y',
+                    'z',
+                    "'",
+                ],
+            },
+        }
+
+
+@pytest.fixture()
+def speech_classification_model():
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': 32,
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                }
+            ],
+        },
+    }
+
+    decoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASRDecoderClassification',
+        'params': {
+            'feat_in': 32,
+            'num_classes': 30,
+        },
+    }
+
+    modelConfig = DictConfig(
+        {
+            'preprocessor': DictConfig(preprocessor),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'labels': ListConfig(["dummy_cls_{}".format(i + 1) for i in range(30)]),
+        }
+    )
+    model = EncDecClassificationModel(cfg=modelConfig)
+    return model
+
+
+@pytest.fixture()
+def speaker_label_model():
+    preprocessor = {
+        '_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor',
+    }
+    encoder = {
+        '_target_': 'nemo.collections.asr.modules.ConvASREncoder',
+        'feat_in': 64,
+        'activation': 'relu',
+        'conv_mask': True,
+        'jasper': [
+            {
+                'filters': 512,
+                'repeat': 1,
+                'kernel': [1],
+                'stride': [1],
+                'dilation': [1],
+                'dropout': 0.0,
+                'residual': False,
+                'separable': False,
+            }
+        ],
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.SpeakerDecoder',
+        'feat_in': 512,
+        'num_classes': 2,
+        'pool_mode': 'attention',
+        'emb_sizes': [1024],
+    }
+
+    modelConfig = DictConfig(
+        {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
+    )
+    speaker_model = EncDecSpeakerLabelModel(cfg=modelConfig)
+    return speaker_model
+
+
+@pytest.fixture()
+def citrinet_model():
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 80,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': 512,
+                    'repeat': 1,
+                    'kernel': [5],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                },
+                {
+                    'filters': 512,
+                    'repeat': 5,
+                    'kernel': [11],
+                    'stride': [2],
+                    'dilation': [1],
+                    'dropout': 0.1,
+                    'residual': True,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                    'stride_last': True,
+                    'residual_mode': 'stride_add',
+                },
+                {
+                    'filters': 512,
+                    'repeat': 5,
+                    'kernel': [13],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.1,
+                    'residual': True,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                },
+                {
+                    'filters': 640,
+                    'repeat': 1,
+                    'kernel': [41],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': True,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                },
+            ],
+        },
+    }
+
+    decoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
+        'params': {'feat_in': 640, 'num_classes': 1024, 'vocabulary': list(chr(i % 28) for i in range(0, 1024))},
+    }
+
+    modelConfig = DictConfig(
+        {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
+    )
+    citri_model = EncDecCTCModel(cfg=modelConfig)
+    return citri_model
+
+
+@pytest.fixture()
+def citrinet_rnnt_model():
+    labels = list(chr(i % 28) for i in range(0, 1024))
+    model_defaults = {'enc_hidden': 640, 'pred_hidden': 256, 'joint_hidden': 320}
+
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+    encoder = {
+        '_target_': 'nemo.collections.asr.modules.ConvASREncoder',
+        'feat_in': 80,
+        'activation': 'relu',
+        'conv_mask': True,
+        'jasper': [
+            {
+                'filters': 512,
+                'repeat': 1,
+                'kernel': [5],
+                'stride': [1],
+                'dilation': [1],
+                'dropout': 0.0,
+                'residual': False,
+                'separable': True,
+                'se': True,
+                'se_context_size': -1,
+            },
+            {
+                'filters': 512,
+                'repeat': 5,
+                'kernel': [11],
+                'stride': [2],
+                'dilation': [1],
+                'dropout': 0.1,
+                'residual': True,
+                'separable': True,
+                'se': True,
+                'se_context_size': -1,
+                'stride_last': True,
+                'residual_mode': 'stride_add',
+            },
+            {
+                'filters': 512,
+                'repeat': 5,
+                'kernel': [13],
+                'stride': [1],
+                'dilation': [1],
+                'dropout': 0.1,
+                'residual': True,
+                'separable': True,
+                'se': True,
+                'se_context_size': -1,
+            },
+            {
+                'filters': 640,
+                'repeat': 1,
+                'kernel': [41],
+                'stride': [1],
+                'dilation': [1],
+                'dropout': 0.0,
+                'residual': True,
+                'separable': True,
+                'se': True,
+                'se_context_size': -1,
+            },
+        ],
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.RNNTDecoder',
+        'prednet': {'pred_hidden': 256, 'pred_rnn_layers': 1, 'dropout': 0.0},
+    }
+
+    joint = {
+        '_target_': 'nemo.collections.asr.modules.RNNTJoint',
+        'fuse_loss_wer': False,
+        'jointnet': {'joint_hidden': 320, 'activation': 'relu', 'dropout': 0.0},
+    }
+
+    decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 5}}
+
+    modelConfig = DictConfig(
+        {
+            'preprocessor': DictConfig(preprocessor),
+            'labels': labels,
+            'model_defaults': DictConfig(model_defaults),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'joint': DictConfig(joint),
+            'decoding': DictConfig(decoding),
+        }
+    )
+    citri_model = EncDecRNNTModel(cfg=modelConfig)
+    return citri_model
+
+
+@pytest.fixture()
+def conformer_model():
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConformerEncoder',
+        'params': {
+            'feat_in': 80,
+            'feat_out': -1,
+            'n_layers': 2,
+            'd_model': 256,
+            'subsampling': 'striding',
+            'subsampling_factor': 4,
+            'subsampling_conv_channels': 512,
+            'reduction': None,
+            'reduction_position': None,
+            'reduction_factor': 1,
+            'ff_expansion_factor': 4,
+            'self_attention_model': 'rel_pos',
+            'n_heads': 8,
+            'att_context_size': [-1, -1],
+            'xscaling': True,
+            'untie_biases': True,
+            'pos_emb_max_len': 500,
+            'conv_kernel_size': 31,
+            'dropout': 0.1,
+            'dropout_pre_encoder': 0.1,
+            'dropout_emb': 0.0,
+            'dropout_att': 0.1,
+        },
+    }
+
+    decoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
+        'params': {'feat_in': 256, 'num_classes': 1024, 'vocabulary': list(chr(i % 28) for i in range(0, 1024))},
+    }
+
+    modelConfig = DictConfig(
+        {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
+    )
+    conformer_model = EncDecCTCModel(cfg=modelConfig)
+    return conformer_model

tests/collections/asr/test_asr_filterbankfeatures_seq_len.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import librosa
+import numpy as np
+import pytest
+import torch
+
+from nemo.collections.asr.parts.preprocessing.features import FilterbankFeatures
+
+
+class TestFilterbankFeatures:
+    @pytest.mark.unit
+    def test_seq_len(self):
+        fb_module = FilterbankFeatures(exact_pad=False, pad_to=1)
+        test_1 = torch.randn(1, 800)
+        test_1_len = torch.tensor([800])
+        fb_spec, fb_len = fb_module(test_1, test_1_len)
+        assert fb_spec.shape[2] - 1 == fb_len[0], f"{fb_spec.shape} != {fb_len}"
+        librosa_spec = librosa.stft(test_1.cpu().detach().numpy().squeeze(), n_fft=512, hop_length=160, win_length=320)
+
+        assert librosa_spec.shape[1] == fb_spec.shape[2], f"{librosa_spec.shape} != {fb_spec.shape}"
+
+    @pytest.mark.unit
+    def test_random_stft_sizes(self):
+        for _ in range(5):
+            nfft = 2 ** np.random.randint(7, 12)
+            window_size = np.random.randint(100, nfft)
+            hop_size = np.random.randint(64, window_size)
+            fb_module = FilterbankFeatures(
+                exact_pad=False,
+                pad_to=1,
+                n_fft=nfft,
+                n_window_size=window_size,
+                n_window_stride=hop_size,
+                normalize=False,
+            )
+            audio_length = np.random.randint(nfft, 2**16)
+            test_1 = torch.randn(1, audio_length)
+            test_1_len = torch.tensor([audio_length])
+            fb_spec, fb_len = fb_module(test_1, test_1_len)
+            assert (
+                fb_spec.shape[2] - 1 == fb_len[0]
+            ), f"{fb_spec.shape} != {fb_len}: {nfft}, {window_size}, {hop_size}, {audio_length}"
+
+            librosa_spec = librosa.stft(
+                test_1.cpu().detach().numpy().squeeze(), n_fft=nfft, hop_length=hop_size, win_length=window_size
+            )
+
+            assert (
+                librosa_spec.shape[1] == fb_spec.shape[2]
+            ), f"{librosa_spec.shape} != {fb_spec.shape}: {nfft}, {window_size}, {hop_size}, {audio_length}"
+
+    @pytest.mark.unit
+    def test_random_stft_sizes_exact_pad(self):
+        for _ in range(5):
+            nfft = 2 ** np.random.randint(7, 12)
+            window_size = np.random.randint(100, nfft)
+            hop_size = np.random.randint(64, window_size)
+            if hop_size % 2 == 1:
+                hop_size = hop_size - 1
+            fb_module = FilterbankFeatures(
+                exact_pad=True,
+                pad_to=1,
+                n_fft=nfft,
+                n_window_size=window_size,
+                n_window_stride=hop_size,
+                normalize=False,
+            )
+            audio_length = np.random.randint(nfft, 2**16)
+            test_1 = torch.randn(1, audio_length)
+            test_1_len = torch.tensor([audio_length])
+            fb_spec, fb_len = fb_module(test_1, test_1_len)
+            assert (
+                fb_spec.shape[2] - 1 == fb_len[0]
+            ), f"{fb_spec.shape} != {fb_len}: {nfft}, {window_size}, {hop_size}, {audio_length}"
+
+            test_2 = test_1.cpu().detach().numpy().squeeze()
+            test_2 = np.pad(test_2, int((nfft - hop_size) // 2), mode="reflect")
+            librosa_spec = librosa.stft(
+                test_2,
+                n_fft=nfft,
+                hop_length=hop_size,
+                win_length=window_size,
+                center=False,
+            )
+
+            assert (
+                fb_spec.shape[2] == librosa_spec.shape[1]
+            ), f"{fb_spec.shape} != {librosa_spec.shape}: {nfft}, {window_size}, {hop_size}, {audio_length}"
+
+            assert (
+                fb_spec.shape[2] == audio_length // hop_size
+            ), f"{fb_spec.shape}, {nfft}, {window_size}, {hop_size}, {audio_length}, {audio_length // hop_size}"

tests/collections/asr/test_asr_hybrid_rnnt_ctc_model_bpe.py

@@ -0,0 +1,360 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import shutil
+import tempfile
+
+import pytest
+import torch
+from lhotse import CutSet, MonoCut
+from lhotse.testing.dummies import DummyManifest
+from omegaconf import DictConfig
+
+from nemo.collections.asr.data.audio_to_text_lhotse import LhotseSpeechToTextBpeDataset
+from nemo.collections.asr.models.hybrid_rnnt_ctc_bpe_models import EncDecHybridRNNTCTCBPEModel
+from nemo.collections.asr.parts.submodules import rnnt_beam_decoding as beam_decode
+from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCBPEDecoding, CTCBPEDecodingConfig
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.collections.common import tokenizers
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+
+NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cpu_is_supported(
+    __NUMBA_MINIMUM_VERSION__
+) or numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+
+
+@pytest.fixture()
+def hybrid_asr_model(test_data_dir):
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+
+    model_defaults = {'enc_hidden': 1024, 'pred_hidden': 64}
+
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': model_defaults['enc_hidden'],
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                }
+            ],
+        },
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.RNNTDecoder',
+        'prednet': {
+            'pred_hidden': model_defaults['pred_hidden'],
+            'pred_rnn_layers': 1,
+        },
+    }
+
+    joint = {
+        '_target_': 'nemo.collections.asr.modules.RNNTJoint',
+        'jointnet': {
+            'joint_hidden': 32,
+            'activation': 'relu',
+        },
+    }
+
+    decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 30}}
+
+    tokenizer = {'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'}
+
+    loss = {'loss_name': 'default', 'warprnnt_numba_kwargs': {'fastemit_lambda': 0.001}}
+
+    aux_ctc = {
+        'ctc_loss_weight': 0.3,
+        'use_cer': False,
+        'ctc_reduction': 'mean_batch',
+        'decoder': {
+            '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+            'feat_in': 1024,
+            'num_classes': -2,
+            'vocabulary': None,
+        },
+        'decoding': DictConfig(CTCBPEDecodingConfig),
+    }
+
+    modelConfig = DictConfig(
+        {
+            'preprocessor': DictConfig(preprocessor),
+            'model_defaults': DictConfig(model_defaults),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'joint': DictConfig(joint),
+            'tokenizer': DictConfig(tokenizer),
+            'decoding': DictConfig(decoding),
+            'loss': DictConfig(loss),
+            'aux_ctc': DictConfig(aux_ctc),
+        }
+    )
+
+    model_instance = EncDecHybridRNNTCTCBPEModel(cfg=modelConfig)
+    return model_instance
+
+
+class TestEncDecHybridRNNTCTCBPEModel:
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_constructor(self, hybrid_asr_model):
+        hybrid_asr_model.train()
+        # TODO: make proper config and assert correct number of weights
+        # Check to/from config_dict:
+        confdict = hybrid_asr_model.to_config_dict()
+        instance2 = EncDecHybridRNNTCTCBPEModel.from_config_dict(confdict)
+        assert isinstance(instance2, EncDecHybridRNNTCTCBPEModel)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_forward(self, hybrid_asr_model):
+        hybrid_asr_model = hybrid_asr_model.eval()
+
+        hybrid_asr_model.preprocessor.featurizer.dither = 0.0
+        hybrid_asr_model.preprocessor.featurizer.pad_to = 0
+
+        hybrid_asr_model.compute_eval_loss = False
+
+        input_signal = torch.randn(size=(4, 512))
+        length = torch.randint(low=321, high=500, size=[4])
+
+        with torch.no_grad():
+            # batch size 1
+            logprobs_instance = []
+            for i in range(input_signal.size(0)):
+                logprobs_ins, _ = hybrid_asr_model.forward(
+                    input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1]
+                )
+                logprobs_instance.append(logprobs_ins)
+            logits_instance = torch.cat(logprobs_instance, 0)
+
+            # batch size 4
+            logprobs_batch, _ = hybrid_asr_model.forward(input_signal=input_signal, input_signal_length=length)
+
+        assert logits_instance.shape == logprobs_batch.shape
+        diff = torch.mean(torch.abs(logits_instance - logprobs_batch))
+        assert diff <= 1e-6
+        diff = torch.max(torch.abs(logits_instance - logprobs_batch))
+        assert diff <= 1e-6
+
+    @pytest.mark.unit
+    def test_predict_step(self, hybrid_asr_model):
+        hybrid_asr_model = hybrid_asr_model.eval()
+        cuts = DummyManifest(CutSet, begin_id=0, end_id=1, with_data=True)
+        dataset = LhotseSpeechToTextBpeDataset(tokenizer=hybrid_asr_model.tokenizer, return_cuts=True)
+        batch = dataset[cuts]
+        outputs = hybrid_asr_model.predict_step(batch, 0)
+        assert len(outputs) == 1
+        assert len(outputs[0]) == 2
+        assert isinstance(outputs[0][0], MonoCut)
+        assert isinstance(outputs[0][1], Hypothesis)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_save_restore_artifact(self, hybrid_asr_model):
+        hybrid_asr_model.train()
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = os.path.join(tmp_dir, 'rnnt_bpe.nemo')
+            hybrid_asr_model.save_to(path)
+
+            new_model = EncDecHybridRNNTCTCBPEModel.restore_from(path)
+            assert isinstance(new_model, type(hybrid_asr_model))
+            assert new_model.vocab_path.endswith('_vocab.txt')
+
+            assert len(new_model.tokenizer.tokenizer.get_vocab()) == 128
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_save_restore_artifact_spe(self, hybrid_asr_model, test_data_dir):
+        hybrid_asr_model.train()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128")
+            hybrid_asr_model.change_vocabulary(new_tokenizer_dir=tokenizer_dir, new_tokenizer_type='bpe')
+
+            save_path = os.path.join(tmpdir, 'ctc_bpe.nemo')
+            hybrid_asr_model.train()
+            hybrid_asr_model.save_to(save_path)
+
+            new_model = EncDecHybridRNNTCTCBPEModel.restore_from(save_path)
+            assert isinstance(new_model, type(hybrid_asr_model))
+            assert isinstance(new_model.tokenizer, tokenizers.SentencePieceTokenizer)
+            assert new_model.model_path.endswith('_tokenizer.model')
+            assert new_model.vocab_path.endswith('_vocab.txt')
+            assert new_model.spe_vocab_path.endswith('_tokenizer.vocab')
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_save_restore_artifact_agg(self, hybrid_asr_model, test_data_dir):
+        tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128")
+        tok_en = {"dir": tokenizer_dir, "type": "wpe"}
+        # the below is really an english tokenizer but we pretend it is spanish
+        tok_es = {"dir": tokenizer_dir, "type": "wpe"}
+        tcfg = DictConfig({"type": "agg", "langs": {"en": tok_en, "es": tok_es}})
+        with tempfile.TemporaryDirectory() as tmpdir:
+            hybrid_asr_model.change_vocabulary(new_tokenizer_dir=tcfg, new_tokenizer_type="agg")
+
+            save_path = os.path.join(tmpdir, "ctc_agg.nemo")
+            hybrid_asr_model.train()
+            hybrid_asr_model.save_to(save_path)
+
+            new_model = EncDecHybridRNNTCTCBPEModel.restore_from(save_path)
+            assert isinstance(new_model, type(hybrid_asr_model))
+            assert isinstance(new_model.tokenizer, tokenizers.AggregateTokenizer)
+
+            # should be double
+            assert new_model.tokenizer.tokenizer.vocab_size == 264
+            assert len(new_model.tokenizer.tokenizer.get_vocab()) == 264
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_vocab_change(self, test_data_dir, hybrid_asr_model):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            old_tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
+            new_tokenizer_dir = os.path.join(tmpdir, 'tokenizer')
+
+            os.makedirs(new_tokenizer_dir, exist_ok=True)
+            shutil.copy2(old_tokenizer_dir, new_tokenizer_dir)
+
+            nw1 = hybrid_asr_model.num_weights
+            hybrid_asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe')
+            # No change
+            assert nw1 == hybrid_asr_model.num_weights
+
+            with open(os.path.join(new_tokenizer_dir, 'vocab.txt'), 'a+') as f:
+                f.write("!\n")
+                f.write('$\n')
+                f.write('@\n')
+
+            hybrid_asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe')
+
+            # rnn embedding + joint + bias
+            pred_embedding = 3 * (hybrid_asr_model.decoder.pred_hidden)
+            joint_joint = 3 * (hybrid_asr_model.joint.joint_hidden + 1)
+            ctc_decoder = 3 * (hybrid_asr_model.ctc_decoder._feat_in + 1)
+            assert hybrid_asr_model.num_weights == (nw1 + (pred_embedding + joint_joint) + ctc_decoder)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_decoding_change(self, hybrid_asr_model):
+        assert isinstance(hybrid_asr_model.decoding.decoding, greedy_decode.GreedyBatchedRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'greedy'
+        new_strategy.greedy = DictConfig({'max_symbols': 10})
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model.decoding.decoding, greedy_decode.GreedyRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'beam'
+        new_strategy.beam = DictConfig({'beam_size': 1})
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert hybrid_asr_model.decoding.decoding.search_type == "default"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'beam'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert hybrid_asr_model.decoding.decoding.search_type == "default"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'tsd'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert hybrid_asr_model.decoding.decoding.search_type == "tsd"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'alsd'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert hybrid_asr_model.decoding.decoding.search_type == "alsd"
+
+        assert hybrid_asr_model.ctc_decoding is not None
+        assert isinstance(hybrid_asr_model.ctc_decoding, CTCBPEDecoding)
+        assert hybrid_asr_model.ctc_decoding.cfg.strategy == "greedy_batch"
+        assert hybrid_asr_model.ctc_decoding.preserve_alignments is False
+        assert hybrid_asr_model.ctc_decoding.compute_timestamps is False
+
+        cfg = CTCBPEDecodingConfig(preserve_alignments=True, compute_timestamps=True)
+        hybrid_asr_model.change_decoding_strategy(cfg, decoder_type="ctc")
+
+        assert hybrid_asr_model.ctc_decoding.preserve_alignments is True
+        assert hybrid_asr_model.ctc_decoding.compute_timestamps is True
+        assert hybrid_asr_model.cur_decoder == "ctc"
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_decoding_type_change(self, hybrid_asr_model):
+        assert isinstance(hybrid_asr_model.decoding.decoding, greedy_decode.GreedyBatchedRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'greedy'
+        new_strategy.greedy = DictConfig({'max_symbols': 10})
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy, decoder_type='rnnt')
+        assert isinstance(hybrid_asr_model.decoding.decoding, greedy_decode.GreedyRNNTInfer)
+        assert hybrid_asr_model.cur_decoder == 'rnnt'
+
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy, decoder_type='ctc')
+        assert isinstance(hybrid_asr_model.ctc_decoding, CTCBPEDecoding)
+        assert hybrid_asr_model.cur_decoder == 'ctc'
+
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy, decoder_type='rnnt')
+        assert isinstance(hybrid_asr_model.decoding.decoding, greedy_decode.GreedyRNNTInfer)
+        assert hybrid_asr_model.cur_decoder == 'rnnt'

tests/collections/asr/test_asr_hybrid_rnnt_ctc_model_bpe_prompt.py

@@ -0,0 +1,566 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import shutil
+import tempfile
+
+import pytest
+import torch
+from omegaconf import DictConfig
+
+from nemo.collections.asr.models.hybrid_rnnt_ctc_bpe_models_prompt import EncDecHybridRNNTCTCBPEModelWithPrompt
+from nemo.collections.asr.parts.submodules import rnnt_beam_decoding as beam_decode
+from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCBPEDecoding, CTCBPEDecodingConfig
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.collections.common import tokenizers
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+
+NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cpu_is_supported(
+    __NUMBA_MINIMUM_VERSION__
+) or numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+
+
+@pytest.fixture()
+def hybrid_asr_model_with_prompt(test_data_dir):
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+
+    model_defaults = {
+        'enc_hidden': 1024,
+        'pred_hidden': 640,
+        'initialize_prompt_feature': True,  # Enable prompt feature initialization
+        'prompt_dictionary': {
+            'en_US': 0,
+            'es_ES': 1,
+            'fr_FR': 2,
+            'de_DE': 3,
+        },
+    }
+
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': model_defaults['enc_hidden'],
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                }
+            ],
+        },
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.RNNTDecoder',
+        'prednet': {
+            'pred_hidden': model_defaults['pred_hidden'],
+            'pred_rnn_layers': 1,
+        },
+    }
+
+    joint = {
+        '_target_': 'nemo.collections.asr.modules.RNNTJoint',
+        'jointnet': {
+            'joint_hidden': 640,
+            'activation': 'relu',
+        },
+    }
+
+    decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 30}}
+
+    tokenizer = {'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'}
+
+    loss = {'loss_name': 'default', 'warprnnt_numba_kwargs': {'fastemit_lambda': 0.001}}
+
+    aux_ctc = {
+        'ctc_loss_weight': 0.1,
+        'use_cer': False,
+        'ctc_reduction': 'mean_batch',
+        'decoder': {
+            '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+            'feat_in': 1024,
+            'num_classes': -2,
+            'vocabulary': None,
+        },
+        'decoding': DictConfig(CTCBPEDecodingConfig),
+    }
+
+    modelConfig = DictConfig(
+        {
+            'preprocessor': DictConfig(preprocessor),
+            'model_defaults': DictConfig(model_defaults),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'joint': DictConfig(joint),
+            'tokenizer': DictConfig(tokenizer),
+            'decoding': DictConfig(decoding),
+            'loss': DictConfig(loss),
+            'aux_ctc': DictConfig(aux_ctc),
+            'num_prompts': 128,
+        }
+    )
+
+    model_instance = EncDecHybridRNNTCTCBPEModelWithPrompt(cfg=modelConfig)
+    return model_instance
+
+
+class TestEncDecHybridRNNTCTCBPEModelWithPrompt:
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_constructor(self, hybrid_asr_model_with_prompt):
+        hybrid_asr_model_with_prompt.train()
+        # Check to/from config_dict:
+        confdict = hybrid_asr_model_with_prompt.to_config_dict()
+        instance2 = EncDecHybridRNNTCTCBPEModelWithPrompt.from_config_dict(confdict)
+        assert isinstance(instance2, EncDecHybridRNNTCTCBPEModelWithPrompt)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_forward_with_prompt(self, hybrid_asr_model_with_prompt):
+        """Exercise the legacy 3D one-hot ``prompt`` forward path."""
+        hybrid_asr_model_with_prompt = hybrid_asr_model_with_prompt.eval()
+
+        hybrid_asr_model_with_prompt.preprocessor.featurizer.dither = 0.0
+        hybrid_asr_model_with_prompt.preprocessor.featurizer.pad_to = 0
+
+        hybrid_asr_model_with_prompt.compute_eval_loss = False
+
+        input_signal = torch.randn(size=(4, 512))
+        length = torch.randint(low=321, high=500, size=[4])
+
+        # Calculate expected timesteps dynamically for the batch
+        with torch.no_grad():
+            # Process the entire batch to get the actual encoded timesteps
+            batch_processed, batch_processed_len = hybrid_asr_model_with_prompt.preprocessor(
+                input_signal=input_signal, length=length
+            )
+            # Run through encoder to get actual encoded length
+            encoded_sample, encoded_len_sample = hybrid_asr_model_with_prompt.encoder(
+                audio_signal=batch_processed, length=batch_processed_len
+            )
+            # Get the maximum encoded length for creating prompt tensor
+            max_encoded_timesteps = encoded_sample.shape[2]  # [B, D, T] format
+
+        # Create prompt tensor with the correct timesteps dimension
+        prompt = torch.randn(size=(4, max_encoded_timesteps, hybrid_asr_model_with_prompt.num_prompts))
+
+        with torch.no_grad():
+            # batch size 1
+            logprobs_instance = []
+            for i in range(input_signal.size(0)):
+                logprobs_ins, _ = hybrid_asr_model_with_prompt.forward(
+                    input_signal=input_signal[i : i + 1],
+                    input_signal_length=length[i : i + 1],
+                    prompt=prompt[i : i + 1],
+                )
+                logprobs_instance.append(logprobs_ins)
+            logits_instance = torch.cat(logprobs_instance, 0)
+
+            # batch size 4
+            logprobs_batch, _ = hybrid_asr_model_with_prompt.forward(
+                input_signal=input_signal, input_signal_length=length, prompt=prompt
+            )
+
+        assert logits_instance.shape == logprobs_batch.shape
+        diff = torch.mean(torch.abs(logits_instance - logprobs_batch))
+        assert diff <= 1e-6
+        diff = torch.max(torch.abs(logits_instance - logprobs_batch))
+        assert diff <= 1e-6
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_forward(self, hybrid_asr_model_with_prompt):
+        """Exercise the canonical 1D ``prompt_indices`` forward path (model builds the one-hot internally)."""
+        hybrid_asr_model_with_prompt = hybrid_asr_model_with_prompt.eval()
+
+        hybrid_asr_model_with_prompt.preprocessor.featurizer.dither = 0.0
+        hybrid_asr_model_with_prompt.preprocessor.featurizer.pad_to = 0
+
+        hybrid_asr_model_with_prompt.compute_eval_loss = False
+
+        input_signal = torch.randn(size=(4, 512))
+        length = torch.randint(low=321, high=500, size=[4])
+
+        # 1D per-sample language ids
+        prompt_indices = torch.tensor([0, 1, 2, 3], dtype=torch.long)
+
+        with torch.no_grad():
+            # batch size 1
+            logprobs_instance = []
+            for i in range(input_signal.size(0)):
+                logprobs_ins, _ = hybrid_asr_model_with_prompt.forward(
+                    input_signal=input_signal[i : i + 1],
+                    input_signal_length=length[i : i + 1],
+                    prompt_indices=prompt_indices[i : i + 1],
+                )
+                logprobs_instance.append(logprobs_ins)
+            logits_instance = torch.cat(logprobs_instance, 0)
+
+            # batch size 4
+            logprobs_batch, _ = hybrid_asr_model_with_prompt.forward(
+                input_signal=input_signal, input_signal_length=length, prompt_indices=prompt_indices
+            )
+
+        assert logits_instance.shape == logprobs_batch.shape
+        diff = torch.mean(torch.abs(logits_instance - logprobs_batch))
+        assert diff <= 1e-6
+        diff = torch.max(torch.abs(logits_instance - logprobs_batch))
+        assert diff <= 1e-6
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_forward_prompt_vs_indices_equivalence(self, hybrid_asr_model_with_prompt):
+        """A 3D one-hot ``prompt`` and the matching 1D ``prompt_indices`` must produce identical output."""
+        hybrid_asr_model_with_prompt = hybrid_asr_model_with_prompt.eval()
+
+        hybrid_asr_model_with_prompt.preprocessor.featurizer.dither = 0.0
+        hybrid_asr_model_with_prompt.preprocessor.featurizer.pad_to = 0
+
+        input_signal = torch.randn(size=(4, 512))
+        length = torch.randint(low=321, high=500, size=[4])
+        prompt_indices = torch.tensor([0, 1, 2, 3], dtype=torch.long)
+
+        # Build a 3D one-hot prompt that matches what the model would build internally.
+        with torch.no_grad():
+            processed, processed_len = hybrid_asr_model_with_prompt.preprocessor(
+                input_signal=input_signal, length=length
+            )
+            encoded_sample, _ = hybrid_asr_model_with_prompt.encoder(audio_signal=processed, length=processed_len)
+            time_steps = encoded_sample.shape[2]  # [B, D, T]
+
+        num_prompts = hybrid_asr_model_with_prompt.num_prompts
+        prompt_one_hot = torch.zeros(4, time_steps, num_prompts)
+        prompt_one_hot.scatter_(2, prompt_indices.view(4, 1, 1).expand(-1, time_steps, -1), 1.0)
+
+        with torch.no_grad():
+            out_from_prompt, _ = hybrid_asr_model_with_prompt.forward(
+                input_signal=input_signal, input_signal_length=length, prompt=prompt_one_hot
+            )
+            out_from_indices, _ = hybrid_asr_model_with_prompt.forward(
+                input_signal=input_signal, input_signal_length=length, prompt_indices=prompt_indices
+            )
+
+        assert out_from_prompt.shape == out_from_indices.shape
+        assert torch.max(torch.abs(out_from_prompt - out_from_indices)) <= 1e-6
+
+    @pytest.mark.unit
+    def test_predict_step_with_prompt(self, hybrid_asr_model_with_prompt):
+        """Exercise the legacy ``.dim() == 3`` branch of predict_step (3D one-hot prompt in the batch)."""
+        hybrid_asr_model_with_prompt = hybrid_asr_model_with_prompt.eval()
+
+        # Create a simple batch manually
+        batch_size = 1
+        seq_len = 1600
+        hidden_len = 200
+        num_prompts = 128
+
+        # Create mock batch data
+        audio_signal = torch.randn(batch_size, seq_len)
+        audio_lengths = torch.tensor([seq_len])
+        transcript = torch.randint(0, 10, (batch_size, 10))
+        transcript_lengths = torch.tensor([10])
+        prompt = torch.zeros(batch_size, hidden_len, num_prompts)
+        prompt[0, :, 0] = 1  # Set first prompt to 1
+
+        batch = (audio_signal, audio_lengths, transcript, transcript_lengths, prompt)
+
+        outputs = hybrid_asr_model_with_prompt.predict_step(batch, 0)
+        assert len(outputs) == 1
+        assert len(outputs[0]) == 2
+        assert isinstance(outputs[0][1], Hypothesis)
+
+    @pytest.mark.unit
+    def test_predict_step(self, hybrid_asr_model_with_prompt):
+        """Exercise the canonical ``.dim() == 1`` branch of predict_step (1D prompt_indices in the batch)."""
+        hybrid_asr_model_with_prompt = hybrid_asr_model_with_prompt.eval()
+
+        batch_size = 1
+        seq_len = 1600
+
+        audio_signal = torch.randn(batch_size, seq_len)
+        audio_lengths = torch.tensor([seq_len])
+        transcript = torch.randint(0, 10, (batch_size, 10))
+        transcript_lengths = torch.tensor([10])
+        # 1D tensor -> hits the prompt_indices branch in predict_step.
+        prompt_indices = torch.tensor([0], dtype=torch.long)
+
+        batch = (audio_signal, audio_lengths, transcript, transcript_lengths, prompt_indices)
+
+        outputs = hybrid_asr_model_with_prompt.predict_step(batch, 0)
+        assert len(outputs) == 1
+        assert len(outputs[0]) == 2
+        assert isinstance(outputs[0][1], Hypothesis)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_save_restore_artifact(self, hybrid_asr_model_with_prompt):
+        hybrid_asr_model_with_prompt.train()
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = os.path.join(tmp_dir, 'rnnt_bpe_prompt.nemo')
+            hybrid_asr_model_with_prompt.save_to(path)
+
+            new_model = EncDecHybridRNNTCTCBPEModelWithPrompt.restore_from(path)
+            assert isinstance(new_model, type(hybrid_asr_model_with_prompt))
+            assert new_model.vocab_path.endswith('_vocab.txt')
+
+            assert len(new_model.tokenizer.tokenizer.get_vocab()) == 128
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_save_restore_artifact_spe(self, hybrid_asr_model_with_prompt, test_data_dir):
+        hybrid_asr_model_with_prompt.train()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128")
+            hybrid_asr_model_with_prompt.change_vocabulary(new_tokenizer_dir=tokenizer_dir, new_tokenizer_type='bpe')
+
+            save_path = os.path.join(tmpdir, 'rnnt_bpe_prompt.nemo')
+            hybrid_asr_model_with_prompt.train()
+            hybrid_asr_model_with_prompt.save_to(save_path)
+
+            new_model = EncDecHybridRNNTCTCBPEModelWithPrompt.restore_from(save_path)
+            assert isinstance(new_model, type(hybrid_asr_model_with_prompt))
+            assert isinstance(new_model.tokenizer, tokenizers.SentencePieceTokenizer)
+            assert new_model.model_path.endswith('_tokenizer.model')
+            assert new_model.vocab_path.endswith('_vocab.txt')
+            assert new_model.spe_vocab_path.endswith('_tokenizer.vocab')
+
+    @pytest.mark.unit
+    def test_save_restore_artifact_agg(self, hybrid_asr_model_with_prompt, test_data_dir):
+        tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128")
+        tok_en = {"dir": tokenizer_dir, "type": "wpe"}
+        # the below is really an english tokenizer but we pretend it is spanish
+        tok_es = {"dir": tokenizer_dir, "type": "wpe"}
+        tcfg = DictConfig({"type": "agg", "langs": {"en": tok_en, "es": tok_es}})
+        with tempfile.TemporaryDirectory() as tmpdir:
+            hybrid_asr_model_with_prompt.change_vocabulary(new_tokenizer_dir=tcfg, new_tokenizer_type="agg")
+
+            save_path = os.path.join(tmpdir, "rnnt_agg_prompt.nemo")
+            hybrid_asr_model_with_prompt.train()
+            hybrid_asr_model_with_prompt.save_to(save_path)
+
+            new_model = EncDecHybridRNNTCTCBPEModelWithPrompt.restore_from(save_path)
+            assert isinstance(new_model, type(hybrid_asr_model_with_prompt))
+            assert isinstance(new_model.tokenizer, tokenizers.AggregateTokenizer)
+
+            # Both source tokenizers are the same 132-token vocab; the AggregateTokenizer
+            # deduplicates 10 shared control tokens, so total = 132 + (132 - 10) = 254.
+            assert new_model.tokenizer.tokenizer.vocab_size == 264
+            assert len(new_model.tokenizer.tokenizer.get_vocab()) == 264
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_vocab_change(self, test_data_dir, hybrid_asr_model_with_prompt):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            old_tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
+            new_tokenizer_dir = os.path.join(tmpdir, 'tokenizer')
+
+            os.makedirs(new_tokenizer_dir, exist_ok=True)
+            shutil.copy2(old_tokenizer_dir, new_tokenizer_dir)
+
+            nw1 = hybrid_asr_model_with_prompt.num_weights
+            hybrid_asr_model_with_prompt.change_vocabulary(
+                new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe'
+            )
+            # No change
+            assert nw1 == hybrid_asr_model_with_prompt.num_weights
+
+            with open(os.path.join(new_tokenizer_dir, 'vocab.txt'), 'a+') as f:
+                f.write("!\n")
+                f.write('$\n')
+                f.write('@\n')
+
+            hybrid_asr_model_with_prompt.change_vocabulary(
+                new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe'
+            )
+
+            # rnn embedding + joint + bias
+            pred_embedding = 3 * (hybrid_asr_model_with_prompt.decoder.pred_hidden)
+            joint_joint = 3 * (hybrid_asr_model_with_prompt.joint.joint_hidden + 1)
+            ctc_decoder = 3 * (hybrid_asr_model_with_prompt.ctc_decoder._feat_in + 1)
+            assert hybrid_asr_model_with_prompt.num_weights == (nw1 + (pred_embedding + joint_joint) + ctc_decoder)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_decoding_change(self, hybrid_asr_model_with_prompt):
+        assert isinstance(hybrid_asr_model_with_prompt.decoding.decoding, greedy_decode.GreedyBatchedRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'greedy'
+        new_strategy.greedy = DictConfig({'max_symbols': 10})
+        hybrid_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model_with_prompt.decoding.decoding, greedy_decode.GreedyRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'beam'
+        new_strategy.beam = DictConfig({'beam_size': 1})
+        hybrid_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model_with_prompt.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert hybrid_asr_model_with_prompt.decoding.decoding.search_type == "default"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'beam'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        hybrid_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model_with_prompt.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert hybrid_asr_model_with_prompt.decoding.decoding.search_type == "default"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'tsd'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        hybrid_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model_with_prompt.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert hybrid_asr_model_with_prompt.decoding.decoding.search_type == "tsd"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'alsd'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        hybrid_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model_with_prompt.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert hybrid_asr_model_with_prompt.decoding.decoding.search_type == "alsd"
+
+        assert hybrid_asr_model_with_prompt.ctc_decoding is not None
+        assert isinstance(hybrid_asr_model_with_prompt.ctc_decoding, CTCBPEDecoding)
+        assert hybrid_asr_model_with_prompt.ctc_decoding.cfg.strategy == "greedy_batch"
+        assert hybrid_asr_model_with_prompt.ctc_decoding.preserve_alignments is False
+        assert hybrid_asr_model_with_prompt.ctc_decoding.compute_timestamps is False
+
+        cfg = CTCBPEDecodingConfig(preserve_alignments=True, compute_timestamps=True)
+        hybrid_asr_model_with_prompt.change_decoding_strategy(cfg, decoder_type="ctc")
+
+        assert hybrid_asr_model_with_prompt.ctc_decoding.preserve_alignments is True
+        assert hybrid_asr_model_with_prompt.ctc_decoding.compute_timestamps is True
+        assert hybrid_asr_model_with_prompt.cur_decoder == "ctc"
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_decoding_type_change(self, hybrid_asr_model_with_prompt):
+        assert isinstance(hybrid_asr_model_with_prompt.decoding.decoding, greedy_decode.GreedyBatchedRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'greedy'
+        new_strategy.greedy = DictConfig({'max_symbols': 10})
+        hybrid_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy, decoder_type='rnnt')
+        assert isinstance(hybrid_asr_model_with_prompt.decoding.decoding, greedy_decode.GreedyRNNTInfer)
+        assert hybrid_asr_model_with_prompt.cur_decoder == 'rnnt'
+
+        hybrid_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy, decoder_type='ctc')
+        assert isinstance(hybrid_asr_model_with_prompt.ctc_decoding, CTCBPEDecoding)
+        assert hybrid_asr_model_with_prompt.cur_decoder == 'ctc'
+
+        hybrid_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy, decoder_type='rnnt')
+        assert isinstance(hybrid_asr_model_with_prompt.decoding.decoding, greedy_decode.GreedyRNNTInfer)
+        assert hybrid_asr_model_with_prompt.cur_decoder == 'rnnt'
+
+    @pytest.mark.unit
+    def test_input_output_types_with_prompt(self, hybrid_asr_model_with_prompt):
+        """Test that input/output types include prompt-specific types."""
+        input_types = hybrid_asr_model_with_prompt.input_types
+        output_types = hybrid_asr_model_with_prompt.output_types
+
+        # Check that prompt is included in input types
+        assert 'prompt' in input_types
+        # Check axes - neural types use tuples with symbolic names
+        prompt_axes = input_types['prompt'].axes
+        assert len(prompt_axes) == 3  # Should be 3D tensor
+
+        # Check standard input types are present
+        assert 'input_signal' in input_types
+        assert 'input_signal_length' in input_types
+
+        # Check output types
+        assert 'outputs' in output_types
+        assert 'encoded_lengths' in output_types
+
+    @pytest.mark.unit
+    def test_prompt_feature_initialization(self, hybrid_asr_model_with_prompt):
+        """Test that prompt feature initialization works correctly."""
+        # Test that the model has prompt-related attributes
+        assert hasattr(hybrid_asr_model_with_prompt, 'concat')
+        assert hasattr(hybrid_asr_model_with_prompt, 'num_prompts')
+        assert hasattr(hybrid_asr_model_with_prompt, 'prompt_kernel')
+
+        # Test that concat is enabled
+        assert hybrid_asr_model_with_prompt.concat == True
+
+        # Test prompt kernel dimensions
+        expected_input_size = (
+            hybrid_asr_model_with_prompt.num_prompts + hybrid_asr_model_with_prompt._cfg.model_defaults.enc_hidden
+        )
+        expected_output_size = hybrid_asr_model_with_prompt._cfg.model_defaults.enc_hidden
+
+        # Check first layer of prompt kernel
+        first_layer = hybrid_asr_model_with_prompt.prompt_kernel[0]
+        assert first_layer.in_features == expected_input_size
+        assert first_layer.out_features == expected_output_size * 2
+
+    @pytest.mark.unit
+    def test_prompt_truncation(self, hybrid_asr_model_with_prompt):
+        """Test that prompts are properly truncated when longer than encoded sequence."""
+        hybrid_asr_model_with_prompt.eval()
+
+        input_signal = torch.randn(size=(1, 512))  # Short signal
+        length = torch.tensor([512])
+
+        # Create a very long prompt (longer than expected encoded length)
+        long_prompt = torch.randn(size=(1, 1000, hybrid_asr_model_with_prompt.num_prompts))
+
+        with torch.no_grad():
+            encoded, encoded_len = hybrid_asr_model_with_prompt.forward(
+                input_signal=input_signal, input_signal_length=length, prompt=long_prompt
+            )
+
+        # Should not crash and should produce valid output
+        assert encoded.shape[0] == 1
+        assert encoded_len.shape[0] == 1

tests/collections/asr/test_asr_hybrid_rnnt_ctc_model_char.py

@@ -0,0 +1,780 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import copy
+from typing import Optional
+
+import pytest
+import torch
+from lhotse import CutSet, MonoCut
+from lhotse.testing.dummies import DummyManifest
+from omegaconf import DictConfig, ListConfig
+
+from nemo.collections.asr.data.audio_to_text_lhotse import LhotseSpeechToTextBpeDataset
+from nemo.collections.asr.models import EncDecHybridRNNTCTCModel
+from nemo.collections.asr.modules import RNNTDecoder, RNNTJoint, SampledRNNTJoint, StatelessTransducerDecoder
+from nemo.collections.asr.parts.submodules import rnnt_beam_decoding as beam_decode
+from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecoding, CTCDecodingConfig
+from nemo.collections.asr.parts.utils import rnnt_utils
+from nemo.collections.common.parts.preprocessing.parsers import make_parser
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+from nemo.utils.config_utils import assert_dataclass_signature_match
+
+NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cpu_is_supported(
+    __NUMBA_MINIMUM_VERSION__
+) or numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+
+
+@pytest.fixture()
+def hybrid_asr_model():
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+
+    # fmt: off
+    labels = [' ', 'a', 'b', 'c', 'd', 'e', 'f',
+              'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o',
+              'p', 'q', 'r', 's', 't', 'u', 'v', 'w',
+              'x', 'y', 'z', "'",
+              ]
+    # fmt: on
+
+    model_defaults = {'enc_hidden': 1024, 'pred_hidden': 64}
+
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': model_defaults['enc_hidden'],
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                }
+            ],
+        },
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.RNNTDecoder',
+        'prednet': {'pred_hidden': model_defaults['pred_hidden'], 'pred_rnn_layers': 1},
+    }
+
+    joint = {
+        '_target_': 'nemo.collections.asr.modules.RNNTJoint',
+        'jointnet': {'joint_hidden': 32, 'activation': 'relu'},
+    }
+
+    decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 30}}
+
+    loss = {'loss_name': 'default', 'warprnnt_numba_kwargs': {'fastemit_lambda': 0.001}}
+
+    aux_ctc = {
+        'ctc_loss_weight': 0.3,
+        'use_cer': False,
+        'ctc_reduction': 'mean_batch',
+        'decoder': {
+            '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+            'feat_in': 1024,
+            'num_classes': len(labels),
+            'vocabulary': labels,
+        },
+        'decoding': DictConfig(CTCDecodingConfig),
+    }
+
+    modelConfig = DictConfig(
+        {
+            'labels': ListConfig(labels),
+            'preprocessor': DictConfig(preprocessor),
+            'model_defaults': DictConfig(model_defaults),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'joint': DictConfig(joint),
+            'decoding': DictConfig(decoding),
+            'loss': DictConfig(loss),
+            'aux_ctc': DictConfig(aux_ctc),
+        }
+    )
+
+    model_instance = EncDecHybridRNNTCTCModel(cfg=modelConfig)
+    return model_instance
+
+
+class TestEncDecHybridRNNTCTCModel:
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_constructor(self, hybrid_asr_model):
+        hybrid_asr_model.train()
+        # TODO: make proper config and assert correct number of weights
+        # Check to/from config_dict:
+        confdict = hybrid_asr_model.to_config_dict()
+        instance2 = EncDecHybridRNNTCTCModel.from_config_dict(confdict)
+        assert isinstance(instance2, EncDecHybridRNNTCTCModel)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_forward(self, hybrid_asr_model):
+        hybrid_asr_model = hybrid_asr_model.eval()
+
+        hybrid_asr_model.preprocessor.featurizer.dither = 0.0
+        hybrid_asr_model.preprocessor.featurizer.pad_to = 0
+
+        hybrid_asr_model.compute_eval_loss = False
+
+        input_signal = torch.randn(size=(4, 512))
+        length = torch.randint(low=321, high=500, size=[4])
+
+        with torch.no_grad():
+            # batch size 1
+            logprobs_instance = []
+            for i in range(input_signal.size(0)):
+                logprobs_ins, _ = hybrid_asr_model.forward(
+                    input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1]
+                )
+                logprobs_instance.append(logprobs_ins)
+            logprobs_instance = torch.cat(logprobs_instance, 0)
+
+            # batch size 4
+            logprobs_batch, _ = hybrid_asr_model.forward(input_signal=input_signal, input_signal_length=length)
+
+        assert logprobs_instance.shape == logprobs_batch.shape
+        diff = torch.mean(torch.abs(logprobs_instance - logprobs_batch))
+        assert diff <= 1e-6
+        diff = torch.max(torch.abs(logprobs_instance - logprobs_batch))
+        assert diff <= 1e-6
+
+    @pytest.mark.unit
+    def test_predict_step(self, hybrid_asr_model):
+        token_list = [" ", "a", "b", "c"]
+        hybrid_asr_model = hybrid_asr_model.eval()
+        cuts = DummyManifest(CutSet, begin_id=0, end_id=1, with_data=True)
+        dataset = LhotseSpeechToTextBpeDataset(tokenizer=make_parser(labels=token_list), return_cuts=True)
+        batch = dataset[cuts]
+        outputs = hybrid_asr_model.predict_step(batch, 0)
+        assert len(outputs) == 1
+        assert len(outputs[0]) == 2
+        assert isinstance(outputs[0][0], MonoCut)
+        assert isinstance(outputs[0][1], rnnt_utils.Hypothesis)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_vocab_change(self, hybrid_asr_model):
+        old_vocab = copy.deepcopy(hybrid_asr_model.joint.vocabulary)
+        nw1 = hybrid_asr_model.num_weights
+        hybrid_asr_model.change_vocabulary(new_vocabulary=old_vocab)
+        # No change
+        assert nw1 == hybrid_asr_model.num_weights
+        new_vocab = copy.deepcopy(old_vocab)
+        new_vocab.append('!')
+        new_vocab.append('$')
+        new_vocab.append('@')
+        hybrid_asr_model.change_vocabulary(new_vocabulary=new_vocab)
+        # fully connected + bias
+        # rnn embedding + joint + bias
+        pred_embedding = 3 * (hybrid_asr_model.decoder.pred_hidden)
+        joint_joint = 3 * (hybrid_asr_model.joint.joint_hidden + 1)
+        ctc_decoder = 3 * (hybrid_asr_model.ctc_decoder._feat_in + 1)
+        assert hybrid_asr_model.num_weights == (nw1 + (pred_embedding + joint_joint) + ctc_decoder)
+        assert hybrid_asr_model.ctc_decoder.vocabulary == hybrid_asr_model.joint.vocabulary
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_decoding_change(self, hybrid_asr_model):
+        assert isinstance(hybrid_asr_model.decoding.decoding, greedy_decode.GreedyBatchedRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'greedy'
+        new_strategy.greedy = DictConfig({'max_symbols': 10})
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model.decoding.decoding, greedy_decode.GreedyRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'beam'
+        new_strategy.beam = DictConfig({'beam_size': 1})
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert hybrid_asr_model.decoding.decoding.search_type == "default"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'beam'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert hybrid_asr_model.decoding.decoding.search_type == "default"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'tsd'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert hybrid_asr_model.decoding.decoding.search_type == "tsd"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'alsd'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(hybrid_asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert hybrid_asr_model.decoding.decoding.search_type == "alsd"
+
+        assert hybrid_asr_model.ctc_decoding is not None
+        assert isinstance(hybrid_asr_model.ctc_decoding, CTCDecoding)
+        assert hybrid_asr_model.ctc_decoding.cfg.strategy == "greedy_batch"
+        assert hybrid_asr_model.ctc_decoding.preserve_alignments is False
+        assert hybrid_asr_model.ctc_decoding.compute_timestamps is False
+
+        cfg = CTCDecodingConfig(preserve_alignments=True, compute_timestamps=True)
+        hybrid_asr_model.change_decoding_strategy(cfg, decoder_type="ctc")
+
+        assert hybrid_asr_model.ctc_decoding.preserve_alignments is True
+        assert hybrid_asr_model.ctc_decoding.compute_timestamps is True
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_decoding_type_change(self, hybrid_asr_model):
+        assert isinstance(hybrid_asr_model.decoding.decoding, greedy_decode.GreedyBatchedRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'greedy'
+        new_strategy.greedy = DictConfig({'max_symbols': 10})
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy, decoder_type='rnnt')
+        assert isinstance(hybrid_asr_model.decoding.decoding, greedy_decode.GreedyRNNTInfer)
+        assert hybrid_asr_model.cur_decoder == 'rnnt'
+
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy, decoder_type='ctc')
+        assert isinstance(hybrid_asr_model.ctc_decoding, CTCDecoding)
+        assert hybrid_asr_model.cur_decoder == 'ctc'
+
+        hybrid_asr_model.change_decoding_strategy(decoding_cfg=new_strategy, decoder_type='rnnt')
+        assert isinstance(hybrid_asr_model.decoding.decoding, greedy_decode.GreedyRNNTInfer)
+        assert hybrid_asr_model.cur_decoder == 'rnnt'
+
+    @pytest.mark.unit
+    def test_GreedyRNNTInferConfig(self):
+        IGNORE_ARGS = [
+            'decoder_model',
+            'joint_model',
+            'blank_index',
+            'tdt_include_duration_confidence',
+            'tdt_include_token_duration',
+            'boosting_tree',
+            'boosting_tree_alpha',
+        ]
+
+        result = assert_dataclass_signature_match(
+            greedy_decode.GreedyRNNTInfer, greedy_decode.GreedyRNNTInferConfig, ignore_args=IGNORE_ARGS
+        )
+
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_GreedyBatchedRNNTInferConfig(self):
+        IGNORE_ARGS = [
+            'decoder_model',
+            'joint_model',
+            'blank_index',
+            'exclude_blank_from_confidence',
+            'tdt_include_duration_confidence',
+            'tdt_include_token_duration',
+            'ngram_lm_model',
+            'ngram_lm_alpha',
+            'boosting_tree',
+            'boosting_tree_alpha',
+            'fusion_models',
+            'fusion_models_alpha',
+        ]
+
+        result = assert_dataclass_signature_match(
+            greedy_decode.GreedyBatchedRNNTInfer, greedy_decode.GreedyBatchedRNNTInferConfig, ignore_args=IGNORE_ARGS
+        )
+
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_BeamRNNTInferConfig(self):
+        IGNORE_ARGS = [
+            'decoder_model',
+            'joint_model',
+            'blank_index',
+            'boosting_tree',
+            'boosting_tree_alpha',
+            'preserve_frame_confidence',
+            'tdt_include_duration_confidence',
+            'confidence_method_cfg',
+        ]
+
+        result = assert_dataclass_signature_match(
+            beam_decode.BeamRNNTInfer, beam_decode.BeamRNNTInferConfig, ignore_args=IGNORE_ARGS
+        )
+
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        ("greedy_class", "loop_labels"),
+        [
+            (greedy_decode.GreedyRNNTInfer, None),
+            (greedy_decode.GreedyBatchedRNNTInfer, True),
+            (greedy_decode.GreedyBatchedRNNTInfer, False),
+        ],
+    )
+    def test_greedy_decoding(self, greedy_class, loop_labels: Optional[bool]):
+        token_list = [" ", "a", "b", "c"]
+        vocab_size = len(token_list)
+
+        encoder_output_size = 4
+        decoder_output_size = 4
+        joint_output_shape = 4
+
+        prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1}
+        jointnet_cfg = {
+            'encoder_hidden': encoder_output_size,
+            'pred_hidden': decoder_output_size,
+            'joint_hidden': joint_output_shape,
+            'activation': 'relu',
+        }
+
+        decoder = RNNTDecoder(prednet_cfg, vocab_size)
+        joint_net = RNNTJoint(jointnet_cfg, vocab_size, vocabulary=token_list)
+
+        additional_decoding_kwargs = {} if loop_labels is None else {"loop_labels": loop_labels}
+        greedy = greedy_class(
+            decoder, joint_net, blank_index=len(token_list) - 1, max_symbols_per_step=5, **additional_decoding_kwargs
+        )
+
+        # (B, D, T)
+        enc_out = torch.randn(1, encoder_output_size, 30)
+        enc_len = torch.tensor([30], dtype=torch.int32)
+
+        with torch.no_grad():
+            _ = greedy(encoder_output=enc_out, encoded_lengths=enc_len)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "greedy_class",
+        [greedy_decode.GreedyRNNTInfer],
+    )
+    def test_greedy_multi_decoding(self, greedy_class):
+        token_list = [" ", "a", "b", "c"]
+        vocab_size = len(token_list)
+
+        encoder_output_size = 4
+        decoder_output_size = 4
+        joint_output_shape = 4
+
+        prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1}
+        jointnet_cfg = {
+            'encoder_hidden': encoder_output_size,
+            'pred_hidden': decoder_output_size,
+            'joint_hidden': joint_output_shape,
+            'activation': 'relu',
+        }
+
+        decoder = RNNTDecoder(prednet_cfg, vocab_size)
+        joint_net = RNNTJoint(jointnet_cfg, vocab_size, vocabulary=token_list)
+
+        greedy = greedy_class(decoder, joint_net, blank_index=len(token_list) - 1, max_symbols_per_step=5)
+
+        # (B, D, T)
+        enc_out = torch.randn(1, encoder_output_size, 30)
+        enc_len = torch.tensor([30], dtype=torch.int32)
+
+        with torch.no_grad():
+            (partial_hyp) = greedy(encoder_output=enc_out, encoded_lengths=enc_len)
+            partial_hyp = partial_hyp[0]
+            _ = greedy(encoder_output=enc_out, encoded_lengths=enc_len, partial_hypotheses=partial_hyp)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        ("greedy_class", "loop_labels"),
+        [
+            (greedy_decode.GreedyRNNTInfer, None),
+            (greedy_decode.GreedyBatchedRNNTInfer, True),
+            (greedy_decode.GreedyBatchedRNNTInfer, False),
+        ],
+    )
+    @pytest.mark.parametrize("context_size", [1, 2])
+    def test_greedy_decoding_stateless_decoder(self, greedy_class, loop_labels: Optional[bool], context_size: int):
+        token_list = [" ", "a", "b", "c"]
+        vocab_size = len(token_list)
+
+        encoder_output_size = 4
+        decoder_output_size = 4
+        joint_output_shape = 4
+
+        prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1, 'context_size': context_size}
+        jointnet_cfg = {
+            'encoder_hidden': encoder_output_size,
+            'pred_hidden': decoder_output_size,
+            'joint_hidden': joint_output_shape,
+            'activation': 'relu',
+        }
+
+        decoder = StatelessTransducerDecoder(prednet_cfg, vocab_size)
+        joint_net = RNNTJoint(jointnet_cfg, vocab_size, vocabulary=token_list)
+
+        additional_decoding_kwargs = {} if loop_labels is None else {"loop_labels": loop_labels}
+        greedy = greedy_class(
+            decoder, joint_net, blank_index=len(token_list) - 1, max_symbols_per_step=5, **additional_decoding_kwargs
+        )
+
+        # (B, D, T)
+        enc_out = torch.randn(1, encoder_output_size, 30)
+        enc_len = torch.tensor([30], dtype=torch.int32)
+
+        with torch.no_grad():
+            _ = greedy(encoder_output=enc_out, encoded_lengths=enc_len)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "greedy_class",
+        [greedy_decode.GreedyRNNTInfer],
+    )
+    def test_greedy_multi_decoding_stateless_decoder(self, greedy_class):
+        token_list = [" ", "a", "b", "c"]
+        vocab_size = len(token_list)
+
+        encoder_output_size = 4
+        decoder_output_size = 4
+        joint_output_shape = 4
+
+        prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1}
+        jointnet_cfg = {
+            'encoder_hidden': encoder_output_size,
+            'pred_hidden': decoder_output_size,
+            'joint_hidden': joint_output_shape,
+            'activation': 'relu',
+        }
+
+        decoder = StatelessTransducerDecoder(prednet_cfg, vocab_size)
+        joint_net = RNNTJoint(jointnet_cfg, vocab_size, vocabulary=token_list)
+
+        greedy = greedy_class(decoder, joint_net, blank_index=len(token_list) - 1, max_symbols_per_step=5)
+
+        # (B, D, T)
+        enc_out = torch.randn(1, encoder_output_size, 30)
+        enc_len = torch.tensor([30], dtype=torch.int32)
+
+        with torch.no_grad():
+            (partial_hyp) = greedy(encoder_output=enc_out, encoded_lengths=enc_len)
+            partial_hyp = partial_hyp[0]
+            _ = greedy(encoder_output=enc_out, encoded_lengths=enc_len, partial_hypotheses=partial_hyp)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        ("greedy_class", "loop_labels"),
+        [
+            (greedy_decode.GreedyRNNTInfer, None),
+            (greedy_decode.GreedyBatchedRNNTInfer, True),
+            (greedy_decode.GreedyBatchedRNNTInfer, False),
+        ],
+    )
+    def test_greedy_decoding_preserve_alignment(self, greedy_class, loop_labels: Optional[bool]):
+        token_list = [" ", "a", "b", "c"]
+        vocab_size = len(token_list)
+
+        encoder_output_size = 4
+        decoder_output_size = 4
+        joint_output_shape = 4
+
+        prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1}
+        jointnet_cfg = {
+            'encoder_hidden': encoder_output_size,
+            'pred_hidden': decoder_output_size,
+            'joint_hidden': joint_output_shape,
+            'activation': 'relu',
+        }
+
+        decoder = RNNTDecoder(prednet_cfg, vocab_size)
+        joint_net = RNNTJoint(jointnet_cfg, vocab_size, vocabulary=token_list)
+
+        additional_decoding_kwargs = {} if loop_labels is None else {"loop_labels": loop_labels}
+        greedy = greedy_class(
+            decoder,
+            joint_net,
+            blank_index=len(token_list) - 1,
+            preserve_alignments=True,
+            max_symbols_per_step=5,
+            **additional_decoding_kwargs,
+        )
+
+        # (B, D, T)
+        enc_out = torch.randn(1, encoder_output_size, 30)
+        enc_len = torch.tensor([30], dtype=torch.int32)
+
+        with torch.no_grad():
+            hyp = greedy(encoder_output=enc_out, encoded_lengths=enc_len)[0][0]  # type: rnnt_utils.Hypothesis
+            assert hyp.alignments is not None
+
+            for t in range(len(hyp.alignments)):
+                for u in range(len(hyp.alignments[t])):
+                    logp, label = hyp.alignments[t][u]
+                    assert torch.is_tensor(logp)
+                    assert torch.is_tensor(label)
+
+    # @pytest.mark.skipif(
+    #     not NUMBA_RNNT_LOSS_AVAILABLE,
+    #     reason='RNNTLoss has not been compiled with appropriate numba version.',
+    # )
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "greedy"},
+            {"search_type": "default", "score_norm": False, "return_best_hypothesis": False},
+            {"search_type": "alsd", "alsd_max_target_len": 20, "return_best_hypothesis": False},
+            {"search_type": "tsd", "tsd_max_sym_exp_per_step": 3, "return_best_hypothesis": False},
+            {"search_type": "maes", "maes_num_steps": 2, "maes_expansion_beta": 2, "return_best_hypothesis": False},
+            {"search_type": "maes", "maes_num_steps": 3, "maes_expansion_beta": 1, "return_best_hypothesis": False},
+        ],
+    )
+    def test_beam_decoding(self, beam_config):
+        token_list = [" ", "a", "b", "c"]
+        vocab_size = len(token_list)
+        beam_size = 1 if beam_config["search_type"] == "greedy" else 2
+
+        encoder_output_size = 4
+        decoder_output_size = 4
+        joint_output_shape = 4
+
+        prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1}
+        jointnet_cfg = {
+            'encoder_hidden': encoder_output_size,
+            'pred_hidden': decoder_output_size,
+            'joint_hidden': joint_output_shape,
+            'activation': 'relu',
+        }
+
+        decoder = RNNTDecoder(prednet_cfg, vocab_size)
+        joint_net = RNNTJoint(jointnet_cfg, vocab_size, vocabulary=token_list)
+
+        beam = beam_decode.BeamRNNTInfer(
+            decoder,
+            joint_net,
+            beam_size=beam_size,
+            **beam_config,
+        )
+
+        # (B, D, T)
+        enc_out = torch.randn(1, encoder_output_size, 30)
+        enc_len = torch.tensor([30], dtype=torch.int32)
+
+        with torch.no_grad():
+            _ = beam(encoder_output=enc_out, encoded_lengths=enc_len)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "greedy"},
+            {"search_type": "default", "score_norm": False, "return_best_hypothesis": False},
+        ],
+    )
+    def test_beam_decoding_preserve_alignments(self, beam_config):
+        token_list = [" ", "a", "b", "c"]
+        vocab_size = len(token_list)
+        beam_size = 1 if beam_config["search_type"] == "greedy" else 2
+
+        encoder_output_size = 4
+        decoder_output_size = 4
+        joint_output_shape = 4
+
+        prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1}
+        jointnet_cfg = {
+            'encoder_hidden': encoder_output_size,
+            'pred_hidden': decoder_output_size,
+            'joint_hidden': joint_output_shape,
+            'activation': 'relu',
+        }
+
+        decoder = RNNTDecoder(prednet_cfg, vocab_size)
+        joint_net = RNNTJoint(jointnet_cfg, vocab_size, vocabulary=token_list)
+
+        beam = beam_decode.BeamRNNTInfer(
+            decoder, joint_net, beam_size=beam_size, **beam_config, preserve_alignments=True
+        )
+
+        # (B, D, T)
+        enc_out = torch.randn(1, encoder_output_size, 30)
+        enc_len = torch.tensor([30], dtype=torch.int32)
+
+        with torch.no_grad():
+            hyp = beam(encoder_output=enc_out, encoded_lengths=enc_len)[0][0]  # type: rnnt_utils.Hypothesis
+
+            if isinstance(hyp, rnnt_utils.NBestHypotheses):
+                hyp = hyp.n_best_hypotheses[0]  # select top hypothesis only
+
+            assert hyp.alignments is not None
+
+            for t in range(len(hyp.alignments)):
+                for u in range(len(hyp.alignments[t])):
+                    logp, label = hyp.alignments[t][u]
+                    assert torch.is_tensor(logp)
+                    assert torch.is_tensor(label)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        ("greedy_class", "loop_labels"),
+        [
+            (greedy_decode.GreedyRNNTInfer, None),
+            (greedy_decode.GreedyBatchedRNNTInfer, True),
+            (greedy_decode.GreedyBatchedRNNTInfer, False),
+        ],
+    )
+    def test_greedy_decoding_SampledRNNTJoint(self, greedy_class, loop_labels: Optional[bool]):
+        token_list = [" ", "a", "b", "c"]
+        vocab_size = len(token_list)
+
+        encoder_output_size = 4
+        decoder_output_size = 4
+        joint_output_shape = 4
+
+        prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1}
+        jointnet_cfg = {
+            'encoder_hidden': encoder_output_size,
+            'pred_hidden': decoder_output_size,
+            'joint_hidden': joint_output_shape,
+            'activation': 'relu',
+        }
+
+        decoder = RNNTDecoder(prednet_cfg, vocab_size)
+        joint_net = SampledRNNTJoint(jointnet_cfg, vocab_size, n_samples=2, vocabulary=token_list)
+
+        additional_decoding_kwargs = {} if loop_labels is None else {"loop_labels": loop_labels}
+        greedy = greedy_class(
+            decoder, joint_net, blank_index=len(token_list) - 1, max_symbols_per_step=5, **additional_decoding_kwargs
+        )
+
+        # (B, D, T)
+        enc_out = torch.randn(1, encoder_output_size, 30)
+        enc_len = torch.tensor([30], dtype=torch.int32)
+
+        with torch.no_grad():
+            _ = greedy(encoder_output=enc_out, encoded_lengths=enc_len)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "greedy"},
+            {"search_type": "default", "score_norm": False, "return_best_hypothesis": False},
+            {"search_type": "alsd", "alsd_max_target_len": 20, "return_best_hypothesis": False},
+            {"search_type": "tsd", "tsd_max_sym_exp_per_step": 3, "return_best_hypothesis": False},
+            {"search_type": "maes", "maes_num_steps": 2, "maes_expansion_beta": 2, "return_best_hypothesis": False},
+            {"search_type": "maes", "maes_num_steps": 3, "maes_expansion_beta": 1, "return_best_hypothesis": False},
+        ],
+    )
+    def test_beam_decoding_SampledRNNTJoint(self, beam_config):
+        token_list = [" ", "a", "b", "c"]
+        vocab_size = len(token_list)
+        beam_size = 1 if beam_config["search_type"] == "greedy" else 2
+
+        encoder_output_size = 4
+        decoder_output_size = 4
+        joint_output_shape = 4
+
+        prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1}
+        jointnet_cfg = {
+            'encoder_hidden': encoder_output_size,
+            'pred_hidden': decoder_output_size,
+            'joint_hidden': joint_output_shape,
+            'activation': 'relu',
+        }
+
+        decoder = RNNTDecoder(prednet_cfg, vocab_size)
+        joint_net = SampledRNNTJoint(jointnet_cfg, vocab_size, n_samples=2, vocabulary=token_list)
+
+        beam = beam_decode.BeamRNNTInfer(
+            decoder,
+            joint_net,
+            beam_size=beam_size,
+            **beam_config,
+        )
+
+        # (B, D, T)
+        enc_out = torch.randn(1, encoder_output_size, 30)
+        enc_len = torch.tensor([30], dtype=torch.int32)
+
+        with torch.no_grad():
+            _ = beam(encoder_output=enc_out, encoded_lengths=enc_len)

tests/collections/asr/test_asr_interctc_models.py

@@ -0,0 +1,269 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Dict
+
+import lightning.pytorch as pl
+import pytest
+import torch
+from omegaconf import DictConfig, ListConfig
+
+from nemo.collections.asr.models import EncDecCTCModel, EncDecHybridRNNTCTCModel
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig
+from nemo.core.classes.mixins import AccessMixin
+
+
+def jasper_encoder_config(num_layers=1) -> Dict:
+    return {
+        '_target_': 'nemo.collections.asr.modules.ConvASREncoder',
+        'feat_in': 64,
+        'activation': 'relu',
+        'conv_mask': True,
+        'jasper': [
+            {
+                'filters': 4,
+                'repeat': 1,
+                'kernel': [1],
+                'stride': [1],
+                'dilation': [1],
+                'dropout': 0.0,
+                'residual': False,
+                'separable': True,
+                'se': True,
+                'se_context_size': -1,
+            }
+        ]
+        * num_layers,
+    }
+
+
+def conformer_encoder_config() -> Dict:
+    return {
+        '_target_': 'nemo.collections.asr.modules.ConformerEncoder',
+        'feat_in': 64,
+        'n_layers': 8,
+        'd_model': 4,
+    }
+
+
+class TestInterCTCLoss:
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "model_class",
+        [EncDecCTCModel, EncDecHybridRNNTCTCModel],
+    )
+    @pytest.mark.parametrize(
+        "encoder_config",
+        [jasper_encoder_config(num_layers=8), conformer_encoder_config()],
+    )
+    @pytest.mark.parametrize(
+        "apply_at_layers,loss_weights",
+        [
+            ([2, 4], [0.1, 0.3]),
+            ([4], [0.3]),
+            ([], []),
+            # errors
+            ([2, 4], [0.1]),
+            ([2], [0.1, 0.3]),
+            ([], [0.3]),
+        ],
+    )
+    def test_forward(self, model_class, encoder_config, apply_at_layers, loss_weights):
+        preprocessor_config = {'_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor'}
+        vocabulary = [
+            ' ',
+            'a',
+            'b',
+            'c',
+            'd',
+            'e',
+            'f',
+            'g',
+            'h',
+            'i',
+            'j',
+            'k',
+            'l',
+            'm',
+            'n',
+            'o',
+            'p',
+            'q',
+            'r',
+            's',
+            't',
+            'u',
+            'v',
+            'w',
+            'x',
+            'y',
+            'z',
+            "'",
+        ]
+        if model_class is EncDecCTCModel:
+            decoder_config = {
+                '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+                'feat_in': None,
+                'num_classes': len(vocabulary),
+                'vocabulary': vocabulary,
+            }
+            model_config = DictConfig(
+                {
+                    'compute_eval_loss': True,  # will be ignored by the model
+                    'preprocessor': DictConfig(preprocessor_config),
+                    'encoder': DictConfig(encoder_config),
+                    'decoder': DictConfig(decoder_config),
+                }
+            )
+        else:
+            decoder_config = {
+                '_target_': 'nemo.collections.asr.modules.RNNTDecoder',
+                'prednet': {'pred_hidden': 4, 'pred_rnn_layers': 1},
+            }
+            joint_config = {
+                '_target_': 'nemo.collections.asr.modules.RNNTJoint',
+                'jointnet': {'joint_hidden': 4, 'activation': 'relu'},
+            }
+            decoding_config = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 30}}
+            loss_config = {'loss_name': 'default', 'warprnnt_numba_kwargs': {'fastemit_lambda': 0.001}}
+
+            aux_ctc_config = {
+                'ctc_loss_weight': 0.3,
+                'use_cer': False,
+                'ctc_reduction': 'mean_batch',
+                'decoder': {
+                    '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+                    'feat_in': None,
+                    'num_classes': len(vocabulary),
+                    'vocabulary': vocabulary,
+                },
+                'decoding': DictConfig(CTCDecodingConfig),
+            }
+            model_config = DictConfig(
+                {
+                    'compute_eval_loss': True,
+                    'labels': ListConfig(vocabulary),
+                    'preprocessor': DictConfig(preprocessor_config),
+                    'model_defaults': DictConfig({'enc_hidden': 4, 'pred_hidden': 4}),
+                    'encoder': DictConfig(encoder_config),
+                    'decoder': DictConfig(decoder_config),
+                    'joint': DictConfig(joint_config),
+                    'decoding': DictConfig(decoding_config),
+                    'loss': DictConfig(loss_config),
+                    'aux_ctc': DictConfig(aux_ctc_config),
+                }
+            )
+        model_config.update(
+            {
+                'interctc': {'loss_weights': loss_weights, 'apply_at_layers': apply_at_layers},
+                'optim': {'name': 'adamw'},
+            }
+        )
+
+        class DummyDataset(torch.utils.data.Dataset):
+            """Simply returns a single set of values."""
+
+            def __init__(self, values):
+                self.values = values
+
+            def __len__(self):
+                return 1
+
+            def __getitem__(self, idx):
+                return self.values
+
+        # this sometimes results in all zeros in the output which breaks tests
+        # so using this only for the ptl calls in the bottom, but using
+        # processed signal directly initially to remove the chance of
+        # this edge-case
+        input_signal = torch.randn(size=(1, 512))
+        input_length = torch.randint(low=321, high=500, size=[1])
+        target = torch.randint(size=(1, input_length[0]), low=0, high=28)
+        target_length = torch.tensor([input_length[0]])
+
+        processed_signal = torch.randn(size=([1, 64, 12]))
+        processed_length = torch.tensor([8])
+
+        if len(apply_at_layers) != len(loss_weights):
+            # has to throw an error here
+            with pytest.raises(
+                ValueError, match="Length of interctc.apply_at_layers has to match interctc.loss_weights"
+            ):
+                asr_model = model_class(cfg=model_config)
+                asr_model.train()
+                logprobs, _, _ = asr_model.forward(input_signal=input_signal, input_signal_length=input_length)
+        else:
+            asr_model = model_class(cfg=model_config)
+            asr_model.train()
+            AccessMixin.set_access_enabled(access_enabled=True, guid=asr_model.model_guid)
+            logprobs, *_ = asr_model.forward(
+                processed_signal=processed_signal, processed_signal_length=processed_length
+            )
+            captured_tensors = asr_model.get_captured_interctc_tensors()
+            AccessMixin.reset_registry(asr_model)
+            assert len(captured_tensors) == len(apply_at_layers)
+            for output in captured_tensors:
+                # checking that values are not the same, if shape is the same
+                assert output[0].shape != logprobs.shape or not torch.allclose(output[0], logprobs)
+                # hybrid model returns output of encoder, so it's not expected to match
+                if model_class is EncDecCTCModel:
+                    assert output[0].shape == logprobs.shape
+
+            # Explicitly pass accelerator as cpu, since default val in PTL >= 2.0 is auto and it picks cuda
+            # which further causes an error in all reduce at: https://github.com/NVIDIA/NeMo/blob/v1.18.1/nemo/collections/asr/modules/conv_asr.py#L209
+            trainer = pl.Trainer(max_epochs=1, accelerator='cpu')
+            trainer.fit(
+                asr_model,
+                train_dataloaders=torch.utils.data.DataLoader(
+                    DummyDataset([input_signal, input_length, target, target_length]),
+                    collate_fn=lambda x: x[0],
+                ),
+                val_dataloaders=torch.utils.data.DataLoader(
+                    DummyDataset([input_signal, input_length, target, target_length]),
+                    collate_fn=lambda x: x[0],
+                ),
+            )
+            required_metrics = ['final_loss'] if len(loss_weights) > 0 else []
+            required_metrics += [f'inter_ctc_loss_l{idx}' for idx in apply_at_layers]
+            prefix = "val_"
+            required_metrics += [f'{prefix}{metric}' for metric in required_metrics]
+            required_metrics += [f'{prefix}wer'] + [f'{prefix}inter_wer_l{idx}' for idx in apply_at_layers]
+            for metric in required_metrics:
+                if 'loss' in metric and 'val_' in metric:
+                    if model_config['compute_eval_loss']:
+                        assert metric in trainer.logged_metrics
+                    else:
+                        assert metric not in trainer.logged_metrics
+                else:
+                    assert metric in trainer.logged_metrics
+
+            trainer.test(
+                asr_model,
+                dataloaders=torch.utils.data.DataLoader(
+                    DummyDataset([input_signal, input_length, target, target_length]),
+                    collate_fn=lambda x: x[0],
+                ),
+            )
+            required_metrics = [f'inter_ctc_loss_l{idx}' for idx in apply_at_layers]
+            prefix = 'test_'
+            # note that "=" is on purpose here, not "+=", since we only log test metrics
+            required_metrics = [f'{prefix}{metric}' for metric in required_metrics]
+            required_metrics += [f'{prefix}wer'] + [f'{prefix}inter_wer_l{idx}' for idx in apply_at_layers]
+            for metric in required_metrics:
+                if 'loss' in metric:
+                    if model_config['compute_eval_loss']:
+                        assert metric in trainer.logged_metrics
+                    else:
+                        assert metric not in trainer.logged_metrics
+                else:
+                    assert metric in trainer.logged_metrics

tests/collections/asr/test_asr_lhotse_dataset.py

@@ -0,0 +1,152 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import patch
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.dataset import AudioSamples
+from lhotse.testing.dummies import DummyManifest
+
+from nemo.collections.asr.data.audio_to_text_lhotse import LhotseSpeechToTextBpeDataset
+from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer, create_spt_model
+
+
+@pytest.fixture(scope="session")
+def tokenizer(tmp_path_factory) -> SentencePieceTokenizer:
+    tmpdir = tmp_path_factory.mktemp("klingon_tokens")
+    text_path = tmpdir / "text.txt"
+    text_path.write_text("\n".join(map(chr, range(ord('a'), ord('z')))))
+    model_path, vocab_path = create_spt_model(
+        text_path, vocab_size=32, sample_size=-1, do_lower_case=False, output_dir=str(tmpdir)
+    )
+    return SentencePieceTokenizer(model_path)
+
+
+def test_lhotse_asr_dataset(tokenizer):
+    # 3 cuts of duration 1s with audio and a single supervision with text 'irrelevant'
+    cuts = DummyManifest(CutSet, begin_id=0, end_id=3, with_data=True)
+
+    # cuts[0] is the default case: audio + single untokenized superivision
+
+    # cuts[1]: audio + single pre-tokenized superivision
+    cuts[1].supervisions[0].tokens = tokenizer.text_to_ids(cuts[1].supervisions[0].text)
+
+    # cuts[2]: audio + two supervisions
+    cuts[2].supervisions = [
+        SupervisionSegment(id="cuts2-sup0", recording_id=cuts[2].recording_id, start=0, duration=0.5, text="first"),
+        SupervisionSegment(id="cuts2-sup1", recording_id=cuts[2].recording_id, start=0.5, duration=0.5, text="second"),
+    ]
+
+    dataset = LhotseSpeechToTextBpeDataset(tokenizer=tokenizer)
+    batch = dataset[cuts]
+
+    assert isinstance(batch, tuple)
+    assert len(batch) == 4
+    assert all(isinstance(t, torch.Tensor) for t in batch)
+
+    audio, audio_lens, tokens, token_lens = batch
+
+    assert audio.shape == (3, 16000)
+    assert audio_lens.tolist() == [16000] * 3
+
+    assert tokens.shape == (3, 13)
+    assert tokens[0].tolist() == [1, 10, 19, 19, 6, 13, 6, 23, 2, 15, 21, 0, 0]
+    assert tokens[1].tolist() == tokens[0].tolist()
+    assert tokens[2].tolist() == [1, 7, 10, 19, 20, 21, 1, 20, 6, 4, 16, 15, 5]
+
+    assert token_lens.tolist() == [11, 11, 13]
+
+
+def test_lhotse_asr_dataset_metadata(tokenizer):
+
+    cuts = DummyManifest(CutSet, begin_id=0, end_id=2, with_data=True)
+
+    cuts[0].id = "cuts0"
+    cuts[1].id = "cuts1"
+    cuts[0].supervisions = [
+        SupervisionSegment(id="cuts0-sup0", recording_id=cuts[0].recording_id, start=0.2, duration=0.5, text="first"),
+    ]
+    cuts[1].supervisions = [
+        SupervisionSegment(id="cuts1-sup0", recording_id=cuts[1].recording_id, start=0, duration=1, text=""),
+    ]
+
+    datasets_metadata = LhotseSpeechToTextBpeDataset(tokenizer=tokenizer, return_cuts=True)
+    batch = datasets_metadata[cuts]
+    assert isinstance(batch, tuple)
+    assert len(batch) == 5
+
+    _, _, _, _, cuts_metadata = batch
+
+    assert cuts_metadata[0].supervisions[0].text == "first"
+    assert cuts_metadata[1].supervisions[0].text == ""
+    assert cuts_metadata[0].id == "cuts0"
+    assert cuts_metadata[1].id == "cuts1"
+
+    assert cuts_metadata[0].supervisions[0].duration == 0.5
+    assert cuts_metadata[0].supervisions[0].start == 0.2
+
+    assert cuts_metadata[1].supervisions[0].duration == 1
+    assert cuts_metadata[1].supervisions[0].start == 0.0
+
+
+def test_lhotse_asr_dataset_ais_batch_loading_enabled(tokenizer, monkeypatch):
+    """Test that USE_AIS_GET_BATCH=true passes use_batch_loader=True to AudioSamples."""
+    monkeypatch.setenv("USE_AIS_GET_BATCH", "true")
+
+    with patch.object(AudioSamples, "__init__", return_value=None) as mock_init:
+        mock_init.side_effect = lambda *args, **kwargs: None
+        try:
+            dataset = LhotseSpeechToTextBpeDataset(tokenizer=tokenizer)
+        except Exception:
+            pass
+        # Check that AudioSamples was called with use_batch_loader=True
+        mock_init.assert_called_with(fault_tolerant=True, use_batch_loader=True)
+
+
+def test_lhotse_asr_dataset_ais_batch_loading_disabled(tokenizer, monkeypatch):
+    """Test that without USE_AIS_GET_BATCH, use_batch_loader=False is passed to AudioSamples."""
+    monkeypatch.delenv("USE_AIS_GET_BATCH", raising=False)
+
+    with patch.object(AudioSamples, "__init__", return_value=None) as mock_init:
+        mock_init.side_effect = lambda *args, **kwargs: None
+        try:
+            dataset = LhotseSpeechToTextBpeDataset(tokenizer=tokenizer)
+        except Exception:
+            pass
+        # Check that AudioSamples was called with use_batch_loader=False
+        mock_init.assert_called_with(fault_tolerant=True, use_batch_loader=False)
+
+
+def test_lhotse_asr_dataset_ais_batch_loading_fallback(tokenizer, monkeypatch):
+    """Test fallback when Lhotse doesn't support use_batch_loader (< 1.32.0)."""
+    monkeypatch.setenv("USE_AIS_GET_BATCH", "true")
+
+    call_args = []
+
+    original_init = AudioSamples.__init__
+
+    def mock_init(self, *args, **kwargs):
+        call_args.append(kwargs.copy())
+        if "use_batch_loader" in kwargs:
+            raise TypeError("unexpected keyword argument 'use_batch_loader'")
+        return original_init(self, *args, **kwargs)
+
+    with patch.object(AudioSamples, "__init__", mock_init):
+        dataset = LhotseSpeechToTextBpeDataset(tokenizer=tokenizer)
+
+    # First call should have use_batch_loader=True, second call should not
+    assert call_args[0] == {"fault_tolerant": True, "use_batch_loader": True}
+    assert call_args[1] == {"fault_tolerant": True}

tests/collections/asr/test_asr_lhotse_speaker_dataset.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import DummyManifest
+
+from nemo.collections.asr.data.audio_to_text_lhotse_speaker import LhotseSpeechToTextSpkBpeDataset
+from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer, create_spt_model
+
+
+@pytest.fixture(scope="session")
+def tokenizer(tmp_path_factory) -> SentencePieceTokenizer:
+    tmpdir = tmp_path_factory.mktemp("klingon_tokens")
+    text_path = tmpdir / "text.txt"
+    text_path.write_text("\n".join(map(chr, range(ord('a'), ord('z')))))
+    model_path, vocab_path = create_spt_model(
+        text_path, vocab_size=32, sample_size=-1, do_lower_case=False, output_dir=str(tmpdir)
+    )
+    return SentencePieceTokenizer(model_path)
+
+
+def test_lhotse_asr_speaker_dataset(tokenizer):
+    # 3 cuts of duration 1s with audio and a single supervision with text 'irrelevant'
+    cuts = DummyManifest(CutSet, begin_id=0, end_id=2, with_data=True)
+
+    # cuts[0] is the default case: audio + single untokenized superivision
+
+    # cuts[1]: audio + two supervisions
+    cuts[1].supervisions = [
+        SupervisionSegment(
+            id="cuts1-sup0", recording_id=cuts[1].recording_id, start=0, duration=0.5, text="first", speaker="0"
+        ),
+        SupervisionSegment(
+            id="cuts1-sup1", recording_id=cuts[1].recording_id, start=0.5, duration=0.5, text="second", speaker="1"
+        ),
+    ]
+
+    dataset = LhotseSpeechToTextSpkBpeDataset(cfg={}, tokenizer=tokenizer)
+    batch = dataset[cuts]
+
+    assert isinstance(batch, tuple)
+    assert len(batch) == 6
+    assert all(isinstance(t, torch.Tensor) for t in batch)
+
+    audio, audio_lens, tokens, token_lens, spk_targets, bg_spk_targets = batch
+
+    assert audio.shape == (2, 16000)
+    assert audio_lens.tolist() == [16000] * 2
+
+    assert tokens.shape == (2, 11)
+    assert tokens[0].tolist() == [1, 10, 19, 19, 6, 13, 6, 23, 2, 15, 21]
+    assert tokens[1].tolist() == [1, 20, 6, 4, 16, 15, 5, 0, 0, 0, 0] or tokens[1].tolist() == [
+        1,
+        7,
+        10,
+        19,
+        20,
+        21,
+        0,
+        0,
+        0,
+        0,
+        0,
+    ]
+    assert token_lens.tolist() == [11, 7] or token_lens.tolist() == [11, 6]
+
+    assert spk_targets.shape == (2, 13)
+    assert spk_targets[0].long().tolist() == [1] * 13
+    assert spk_targets[1].long().sum().item() in [6, 7]
+
+    assert bg_spk_targets.shape == (2, 13)
+    assert bg_spk_targets[0].long().tolist() == [0] * 13
+    assert bg_spk_targets[1].long().sum().item() in [6, 7]
+
+    assert (spk_targets[1] + bg_spk_targets[1]).long().tolist() == [1] * 13

tests/collections/asr/test_asr_local_attn.py

@@ -0,0 +1,197 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import shutil
+import tempfile
+
+import lightning.pytorch as pl
+import pytest
+import torch
+from omegaconf import DictConfig
+
+from nemo.collections.asr.models import ASRModel, EncDecCTCModel
+
+
+def getattr2(object, attr):
+    if not '.' in attr:
+        return getattr(object, attr)
+    else:
+        arr = attr.split('.')
+        return getattr2(getattr(object, arr[0]), '.'.join(arr[1:]))
+
+
+class TestASRLocalAttention:
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_forward(self):
+        asr_model = ASRModel.from_pretrained("stt_en_conformer_ctc_small")
+        asr_model = asr_model.eval()
+
+        len = 16000 * 60 * 30  # 30 minutes, OOM without local attention
+        input_signal_long = torch.randn(size=(1, len), device=asr_model.device)
+        length_long = torch.tensor([len], device=asr_model.device)
+
+        # switch to local attn
+        asr_model.change_attention_model(self_attention_model="rel_pos_local_attn", att_context_size=(64, 64))
+        with torch.no_grad():
+            asr_model.forward(input_signal=input_signal_long, input_signal_length=length_long)
+
+        # switch context size only (keep local)
+        asr_model.change_attention_model(att_context_size=(192, 192))
+        with torch.no_grad():
+            asr_model.forward(input_signal=input_signal_long, input_signal_length=length_long)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_change_save_restore(self):
+
+        model = ASRModel.from_pretrained("stt_en_conformer_ctc_small")
+        model.change_attention_model(self_attention_model="rel_pos_local_attn", att_context_size=(64, 64))
+        attr_for_eq_check = ["encoder.self_attention_model", "encoder.att_context_size"]
+
+        with tempfile.TemporaryDirectory() as restore_folder:
+            with tempfile.TemporaryDirectory() as save_folder:
+                save_folder_path = save_folder
+                # Where model will be saved
+                model_save_path = os.path.join(save_folder, f"{model.__class__.__name__}.nemo")
+                model.save_to(save_path=model_save_path)
+                # Where model will be restored from
+                model_restore_path = os.path.join(restore_folder, f"{model.__class__.__name__}.nemo")
+                shutil.copy(model_save_path, model_restore_path)
+            # at this point save_folder should not exist
+            assert save_folder_path is not None and not os.path.exists(save_folder_path)
+            assert not os.path.exists(model_save_path)
+            assert os.path.exists(model_restore_path)
+            # attempt to restore
+            model_copy = model.__class__.restore_from(
+                restore_path=model_restore_path,
+                map_location=None,
+                strict=True,
+                return_config=False,
+                override_config_path=None,
+            )
+
+            assert model.num_weights == model_copy.num_weights
+            if attr_for_eq_check is not None and len(attr_for_eq_check) > 0:
+                for attr in attr_for_eq_check:
+                    assert getattr2(model, attr) == getattr2(model_copy, attr)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "global_tokens",
+        [0, 1, 4],
+    )
+    @pytest.mark.parametrize(
+        "global_tokens_spacing",
+        [1, 4],
+    )
+    def test_train(self, global_tokens, global_tokens_spacing):
+        preprocessor_config = {'_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor'}
+        vocabulary = [
+            ' ',
+            'a',
+            'b',
+            'c',
+            'd',
+            'e',
+            'f',
+            'g',
+            'h',
+            'i',
+            'j',
+            'k',
+            'l',
+            'm',
+            'n',
+            'o',
+            'p',
+            'q',
+            'r',
+            's',
+            't',
+            'u',
+            'v',
+            'w',
+            'x',
+            'y',
+            'z',
+            "'",
+        ]
+        encoder_config = {
+            '_target_': 'nemo.collections.asr.modules.ConformerEncoder',
+            'feat_in': 64,
+            'n_layers': 8,
+            'd_model': 4,
+            'self_attention_model': 'rel_pos_local_attn',
+            'att_context_size': [128, 128],
+            'global_tokens': global_tokens,
+            'global_tokens_spacing': global_tokens_spacing,
+        }
+        decoder_config = {
+            '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+            'feat_in': None,
+            'num_classes': len(vocabulary),
+            'vocabulary': vocabulary,
+        }
+        model_config = DictConfig(
+            {
+                'preprocessor': DictConfig(preprocessor_config),
+                'encoder': DictConfig(encoder_config),
+                'decoder': DictConfig(decoder_config),
+                'optim': {'name': 'adamw'},
+            }
+        )
+
+        class DummyDataset(torch.utils.data.Dataset):
+            """Simply returns a single set of values."""
+
+            def __init__(self, values):
+                self.values = values
+
+            def __len__(self):
+                return 1
+
+            def __getitem__(self, idx):
+                return self.values
+
+        input_signal = torch.randn(size=(1, 960000))
+        input_length = torch.tensor([960000])
+        target = torch.randint(size=(1, 280), low=0, high=28)
+        target_length = torch.tensor([280])
+
+        asr_model = EncDecCTCModel(cfg=model_config)
+        asr_model.train()
+        _ = asr_model.forward(input_signal=input_signal, input_signal_length=input_length)
+        # Explicitly pass accelerator as cpu, since default val in PTL >= 2.0 is auto and it picks cuda
+        # which further causes an error in all reduce at: https://github.com/NVIDIA/NeMo/blob/v1.18.1/nemo/collections/asr/modules/conformer_encoder.py#L462
+        # and in ConvASREncoder where device is CPU
+        trainer = pl.Trainer(max_epochs=1, accelerator='cpu')
+        trainer.fit(
+            asr_model,
+            train_dataloaders=torch.utils.data.DataLoader(
+                DummyDataset([input_signal, input_length, target, target_length]),
+                collate_fn=lambda x: x[0],
+            ),
+            val_dataloaders=torch.utils.data.DataLoader(
+                DummyDataset([input_signal, input_length, target, target_length]),
+                collate_fn=lambda x: x[0],
+            ),
+        )
+        trainer.test(
+            asr_model,
+            dataloaders=torch.utils.data.DataLoader(
+                DummyDataset([input_signal, input_length, target, target_length]),
+                collate_fn=lambda x: x[0],
+            ),
+        )

tests/collections/asr/test_asr_modules.py

@@ -0,0 +1,386 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+from omegaconf import OmegaConf
+
+from nemo.collections.asr import modules
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+from nemo.utils import config_utils, logging
+
+
+class TestASRModulesBasicTests:
+    @pytest.mark.unit
+    def test_AudioToMelSpectrogramPreprocessor_config(self):
+        # Test that dataclass matches signature of module
+        result = config_utils.assert_dataclass_signature_match(
+            modules.AudioToMelSpectrogramPreprocessor,
+            modules.audio_preprocessing.AudioToMelSpectrogramPreprocessorConfig,
+        )
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_AudioToMelSpectrogramPreprocessor_batch(self):
+        # Test 1 that should test the pure stft implementation as much as possible
+        instance1 = modules.AudioToMelSpectrogramPreprocessor(normalize="per_feature", dither=0, pad_to=0)
+
+        # Ensure that the two functions behave similarily
+        for _ in range(10):
+            input_signal, length = instance1.input_example(4, 512, 321)
+
+            with torch.no_grad():
+                # batch size 1
+                res_instance, length_instance = [], []
+                for i in range(input_signal.size(0)):
+                    res_ins, length_ins = instance1(input_signal=input_signal[i : i + 1], length=length[i : i + 1])
+                    res_instance.append(res_ins)
+                    length_instance.append(length_ins)
+
+                res_instance = torch.cat(res_instance, 0)
+                length_instance = torch.cat(length_instance, 0)
+
+                # batch size 4
+                res_batch, length_batch = instance1(input_signal=input_signal, length=length)
+
+            assert res_instance.shape == res_batch.shape
+            assert length_instance.shape == length_batch.shape
+            diff = torch.mean(torch.abs(res_instance - res_batch))
+            assert diff <= 1e-3
+            diff = torch.max(torch.abs(res_instance - res_batch))
+            assert diff <= 1e-3
+
+    @pytest.mark.run_only_on('GPU')
+    def test_AudioToMelSpectrogramPreprocessor_gpu(self):
+        instance0 = modules.AudioToMelSpectrogramPreprocessor().to("cuda")
+        input_signal, length = instance0.input_example()
+
+        with torch.no_grad():
+            processed_signal, _ = instance0(input_signal=input_signal, length=length)
+
+        assert processed_signal.device == input_signal.device
+
+    @pytest.mark.unit
+    def test_SpectrogramAugmentationr_legacy(self):
+        # Make sure constructor works
+        instance1 = modules.SpectrogramAugmentation(
+            freq_masks=10, time_masks=3, rect_masks=3, use_numba_spec_augment=False, use_vectorized_spec_augment=False
+        )
+        assert isinstance(instance1, modules.SpectrogramAugmentation)
+
+        # Make sure forward doesn't throw with expected input
+        instance0 = modules.AudioToMelSpectrogramPreprocessor(dither=0)
+        input_signal, length = instance0.input_example(4, 512, 321)
+        res0 = instance0(input_signal=input_signal, length=length)
+        res = instance1(input_spec=res0[0], length=length)
+
+        assert res.shape == res0[0].shape
+
+    @pytest.mark.unit
+    @pytest.mark.run_only_on('GPU')
+    def test_SpectrogramAugmentationr_vectorized(self):
+        # Make sure constructor works
+        instance1 = modules.SpectrogramAugmentation(
+            freq_masks=10, time_masks=3, rect_masks=3, use_numba_spec_augment=False, use_vectorized_spec_augment=True
+        )
+        assert isinstance(instance1, modules.SpectrogramAugmentation)
+
+        # Make sure forward doesn't throw with expected input
+        instance0 = modules.AudioToMelSpectrogramPreprocessor(dither=0)
+        input_signal, length = instance0.input_example(4, 512, 321)
+        res0 = instance0(input_signal=input_signal, length=length)
+        res = instance1(input_spec=res0[0], length=length)
+
+        assert res.shape == res0[0].shape
+
+    @pytest.mark.unit
+    @pytest.mark.run_only_on('GPU')
+    def test_SpectrogramAugmentationr_numba_kernel(self, caplog):
+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
+
+        logging._logger.propagate = True
+        original_verbosity = logging.get_verbosity()
+        logging.set_verbosity(logging.DEBUG)
+        caplog.set_level(logging.DEBUG)
+
+        # Make sure constructor works
+        instance1 = modules.SpectrogramAugmentation(
+            freq_masks=10, time_masks=3, rect_masks=3, use_numba_spec_augment=True, use_vectorized_spec_augment=False
+        )
+        assert isinstance(instance1, modules.SpectrogramAugmentation)
+
+        # Make sure forward doesn't throw with expected input
+        instance0 = modules.AudioToMelSpectrogramPreprocessor(dither=0)
+        input_signal, length = instance0.input_example(8, 512, 321)
+        res0 = instance0(input_signal=input_signal, length=length)
+        res = instance1(input_spec=res0[0], length=length)
+
+        assert res.shape == res0[0].shape
+
+        # check tha numba kernel debug message indicates that it is available for use
+        assert """Numba SpecAugment kernel is available""" in caplog.text
+
+        logging._logger.propagate = False
+        logging.set_verbosity(original_verbosity)
+
+    @pytest.mark.unit
+    def test_SpectrogramAugmentationr_config(self):
+        # Test that dataclass matches signature of module
+        result = config_utils.assert_dataclass_signature_match(
+            modules.SpectrogramAugmentation,
+            modules.audio_preprocessing.SpectrogramAugmentationConfig,
+        )
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_CropOrPadSpectrogramAugmentation(self):
+        # Make sure constructor works
+        audio_length = 128
+        instance1 = modules.CropOrPadSpectrogramAugmentation(audio_length=audio_length)
+        assert isinstance(instance1, modules.CropOrPadSpectrogramAugmentation)
+
+        # Make sure forward doesn't throw with expected input
+        instance0 = modules.AudioToMelSpectrogramPreprocessor(dither=0)
+        input_signal, length = instance0.input_example(4, 512, 321)
+        res0 = instance0(input_signal=input_signal, length=length)
+        res, new_length = instance1(input_signal=res0[0], length=length)
+
+        assert res.shape == torch.Size([4, 64, audio_length])
+        assert all(new_length == torch.tensor([128] * 4))
+
+    @pytest.mark.unit
+    def test_CropOrPadSpectrogramAugmentation_config(self):
+        # Test that dataclass matches signature of module
+        result = config_utils.assert_dataclass_signature_match(
+            modules.CropOrPadSpectrogramAugmentation,
+            modules.audio_preprocessing.CropOrPadSpectrogramAugmentationConfig,
+        )
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_MaskedPatchAugmentation(self):
+        # Make sure constructor works
+        audio_length = 128
+        instance1 = modules.MaskedPatchAugmentation(patch_size=16, mask_patches=0.5, freq_masks=2, freq_width=10)
+        assert isinstance(instance1, modules.MaskedPatchAugmentation)
+
+        # Make sure forward doesn't throw with expected input
+        instance0 = modules.AudioToMelSpectrogramPreprocessor(dither=0)
+        input_signal, length = instance0.input_example(4, 512, 321)
+        res0 = instance0(input_signal=input_signal, length=length)
+        res = instance1(input_spec=res0[0], length=length)
+
+        assert res.shape == res0[0].shape
+
+    @pytest.mark.unit
+    def test_MaskedPatchAugmentation_config(self):
+        # Test that dataclass matches signature of module
+        result = config_utils.assert_dataclass_signature_match(
+            modules.MaskedPatchAugmentation,
+            modules.audio_preprocessing.MaskedPatchAugmentationConfig,
+        )
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_RNNTDecoder(self):
+        vocab = list(range(10))
+        vocab = [str(x) for x in vocab]
+        vocab_size = len(vocab)
+
+        pred_config = OmegaConf.create(
+            {
+                '_target_': 'nemo.collections.asr.modules.RNNTDecoder',
+                'prednet': {
+                    'pred_hidden': 32,
+                    'pred_rnn_layers': 1,
+                },
+                'vocab_size': vocab_size,
+                'blank_as_pad': True,
+            }
+        )
+
+        prednet = modules.RNNTDecoder.from_config_dict(pred_config)
+
+        # num params
+        pred_hidden = pred_config.prednet.pred_hidden
+        embed = (vocab_size + 1) * pred_hidden  # embedding with blank
+        rnn = (
+            2 * 4 * (pred_hidden * pred_hidden + pred_hidden)
+        )  # (ih + hh) * (ifco gates) * (indim * hiddendim + bias)
+        assert prednet.num_weights == (embed + rnn)
+
+        # State initialization
+        x_ = torch.zeros(4, dtype=torch.float32)
+        states = prednet.initialize_state(x_)
+
+        for state_i in states:
+            assert state_i.dtype == x_.dtype
+            assert state_i.device == x_.device
+            assert state_i.shape[1] == len(x_)
+
+        # Blank hypotheses test
+        blank = vocab_size
+        hyp = Hypothesis(score=0.0, y_sequence=[blank])
+        cache = {}
+        pred, states, _ = prednet.score_hypothesis(hyp, cache)
+
+        assert pred.shape == torch.Size([1, 1, pred_hidden])
+        assert len(states) == 2
+        for state_i in states:
+            assert state_i.dtype == pred.dtype
+            assert state_i.device == pred.device
+            assert state_i.shape[1] == len(pred)
+
+        # Blank stateless predict
+        g, states = prednet.predict(y=None, state=None, add_sos=False, batch_size=1)
+
+        assert g.shape == torch.Size([1, 1, pred_hidden])
+        assert len(states) == 2
+        for state_i in states:
+            assert state_i.dtype == g.dtype
+            assert state_i.device == g.device
+            assert state_i.shape[1] == len(g)
+
+        # Blank stateful predict
+        g, states2 = prednet.predict(y=None, state=states, add_sos=False, batch_size=1)
+
+        assert g.shape == torch.Size([1, 1, pred_hidden])
+        assert len(states2) == 2
+        for state_i, state_j in zip(states, states2):
+            assert (state_i - state_j).square().sum().sqrt() > 0.0
+
+        # Predict with token and state
+        token = torch.full([1, 1], fill_value=0, dtype=torch.long)
+        g, states = prednet.predict(y=token, state=states2, add_sos=False, batch_size=None)
+
+        assert g.shape == torch.Size([1, 1, pred_hidden])
+        assert len(states) == 2
+
+        # Predict with blank token and no state
+        token = torch.full([1, 1], fill_value=blank, dtype=torch.long)
+        g, states = prednet.predict(y=token, state=None, add_sos=False, batch_size=None)
+
+        assert g.shape == torch.Size([1, 1, pred_hidden])
+        assert len(states) == 2
+
+    @pytest.mark.unit
+    def test_RNNTJoint(self):
+        vocab = list(range(10))
+        vocab = [str(x) for x in vocab]
+        vocab_size = len(vocab)
+
+        batchsize = 4
+        encoder_hidden = 64
+        pred_hidden = 32
+        joint_hidden = 16
+
+        joint_cfg = OmegaConf.create(
+            {
+                '_target_': 'nemo.collections.asr.modules.RNNTJoint',
+                'num_classes': vocab_size,
+                'vocabulary': vocab,
+                'jointnet': {
+                    'encoder_hidden': encoder_hidden,
+                    'pred_hidden': pred_hidden,
+                    'joint_hidden': joint_hidden,
+                    'activation': 'relu',
+                },
+            }
+        )
+
+        jointnet = modules.RNNTJoint.from_config_dict(joint_cfg)
+
+        enc = torch.zeros(batchsize, encoder_hidden, 48)  # [B, D1, T]
+        dec = torch.zeros(batchsize, pred_hidden, 24)  # [B, D2, U]
+
+        # forward call test
+        out = jointnet(encoder_outputs=enc, decoder_outputs=dec)
+        assert out.shape == torch.Size([batchsize, 48, 24, vocab_size + 1])  # [B, T, U, V + 1]
+
+        # joint() step test
+        enc2 = enc.transpose(1, 2)  # [B, T, D1]
+        dec2 = dec.transpose(1, 2)  # [B, U, D2]
+        out2 = jointnet.joint(enc2, dec2)  # [B, T, U, V + 1]
+        assert (out - out2).abs().sum() <= 1e-5
+
+        # assert vocab size
+        assert jointnet.num_classes_with_blank == vocab_size + 1
+
+    @pytest.mark.unit
+    def test_HATJoint(self):
+        vocab = list(range(10))
+        vocab = [str(x) for x in vocab]
+        vocab_size = len(vocab)
+
+        batchsize = 4
+        encoder_hidden = 64
+        pred_hidden = 32
+        joint_hidden = 16
+
+        joint_cfg = OmegaConf.create(
+            {
+                '_target_': 'nemo.collections.asr.modules.HATJoint',
+                'num_classes': vocab_size,
+                'vocabulary': vocab,
+                'jointnet': {
+                    'encoder_hidden': encoder_hidden,
+                    'pred_hidden': pred_hidden,
+                    'joint_hidden': joint_hidden,
+                    'activation': 'relu',
+                },
+            }
+        )
+
+        jointnet = modules.HATJoint.from_config_dict(joint_cfg)
+
+        enc = torch.zeros(batchsize, encoder_hidden, 48)  # [B, D1, T]
+        dec = torch.zeros(batchsize, pred_hidden, 24)  # [B, D2, U]
+
+        # forward call test
+        out = jointnet(encoder_outputs=enc, decoder_outputs=dec)
+        assert out.shape == torch.Size([batchsize, 48, 24, vocab_size + 1])  # [B, T, U, V + 1]
+
+        # joint() step test
+        enc2 = enc.transpose(1, 2)  # [B, T, D1]
+        dec2 = dec.transpose(1, 2)  # [B, U, D2]
+        out2 = jointnet.joint(enc2, dec2)  # [B, T, U, V + 1]
+        assert (out - out2).abs().sum() <= 1e-5
+
+        # joint() step test for internal LM subtraction
+        jointnet.return_hat_ilm = True
+        hat_output = jointnet.joint(enc2, dec2)  # HATJointOutput dataclass
+        out3, ilm = hat_output.hat_logprobs, hat_output.ilm_logprobs  # [B, T, U, V + 1] and [B, 1, U, V]
+        assert (out - out3).abs().sum() <= 1e-5
+        assert ilm.shape == torch.Size([batchsize, 1, 24, vocab_size])  # [B, 1, U, V] without blank simbol
+
+        # assert vocab size
+        assert jointnet.num_classes_with_blank == vocab_size + 1

tests/collections/asr/test_asr_multitalker_models.py

@@ -0,0 +1,198 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import pytest
+import torch
+from omegaconf import DictConfig
+
+from nemo.collections.asr.models.multitalker_asr_models import EncDecMultiTalkerRNNTBPEModel
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+
+NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cpu_is_supported(
+    __NUMBA_MINIMUM_VERSION__
+) or numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+
+
+@pytest.fixture()
+def asr_model(test_data_dir):
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+
+    model_defaults = {'enc_hidden': 1024, 'pred_hidden': 64}
+    spk_kernel_type = "ff"
+    spk_kernel_layers = [0]
+    add_bg_spk_kernel = True
+
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConformerEncoder',
+        'params': {
+            'feat_in': 64,
+            'n_layers': 1,
+            'd_model': model_defaults['enc_hidden'],  # Required by SpeakerKernelMixin
+            'subsampling': 'dw_striding',
+            'subsampling_factor': 2,
+            'ff_expansion_factor': 4,
+            'self_attention_model': 'rel_pos',
+            'n_heads': 4,
+            'conv_kernel_size': 7,
+            'dropout': 0.1,
+        },
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.RNNTDecoder',
+        'prednet': {
+            'pred_hidden': model_defaults['pred_hidden'],
+            'pred_rnn_layers': 1,
+        },
+    }
+
+    joint = {
+        '_target_': 'nemo.collections.asr.modules.RNNTJoint',
+        'jointnet': {
+            'joint_hidden': 32,
+            'activation': 'relu',
+        },
+    }
+
+    decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 30}}
+
+    tokenizer = {'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'}
+
+    loss = {'loss_name': 'default', 'warprnnt_numba_kwargs': {'fastemit_lambda': 0.001}}
+
+    modelConfig = DictConfig(
+        {
+            'preprocessor': DictConfig(preprocessor),
+            'model_defaults': DictConfig(model_defaults),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'joint': DictConfig(joint),
+            'tokenizer': DictConfig(tokenizer),
+            'decoding': DictConfig(decoding),
+            'loss': DictConfig(loss),
+            'spk_kernel_type': spk_kernel_type,
+            'spk_kernel_layers': spk_kernel_layers,
+            'add_bg_spk_kernel': add_bg_spk_kernel,
+        }
+    )
+
+    model_instance = EncDecMultiTalkerRNNTBPEModel(cfg=modelConfig)
+    return model_instance
+
+
+class TestEncDecMultiTalkerRNNTBPEModel:
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_constructor(self, asr_model):
+        """Test model constructor and speaker kernel initialization."""
+        asr_model.train()
+
+        # Check that it's the correct type
+        assert isinstance(asr_model, EncDecMultiTalkerRNNTBPEModel)
+
+        # Check speaker kernel configuration
+        assert hasattr(asr_model, 'spk_kernel_type')
+        assert hasattr(asr_model, 'spk_kernel_layers')
+        assert hasattr(asr_model, 'add_bg_spk_kernel')
+
+        # Check speaker kernel initialization
+        assert asr_model.spk_kernel_type == "ff"
+        assert asr_model.spk_kernel_layers == [0]
+        assert asr_model.add_bg_spk_kernel is True
+
+        # Check speaker kernels exist
+        assert hasattr(asr_model, 'spk_kernels')
+        if asr_model.add_bg_spk_kernel:
+            assert hasattr(asr_model, 'bg_spk_kernels')
+
+        # Test config dict conversion
+        confdict = asr_model.to_config_dict()
+        instance2 = EncDecMultiTalkerRNNTBPEModel.from_config_dict(confdict)
+        assert isinstance(instance2, EncDecMultiTalkerRNNTBPEModel)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_forward(self, asr_model):
+        """Test forward pass functionality."""
+        asr_model = asr_model.eval()
+
+        asr_model.preprocessor.featurizer.dither = 0.0
+        asr_model.preprocessor.featurizer.pad_to = 0
+        asr_model.compute_eval_loss = False
+
+        input_signal = torch.randn(size=(4, 512))
+        length = torch.randint(low=321, high=500, size=[4])
+
+        # Create mock speaker targets
+        batch_size = input_signal.size(0)
+        target_length = 32  # Typical encoder output length for test
+        spk_targets = torch.randint(0, 2, (batch_size, target_length), dtype=torch.float32)
+        bg_spk_targets = torch.randint(0, 2, (batch_size, target_length), dtype=torch.float32)
+
+        # Set speaker targets
+        asr_model.set_speaker_targets(spk_targets, bg_spk_targets)
+
+        with torch.no_grad():
+            # batch size 1
+            logprobs_instance = []
+            for i in range(input_signal.size(0)):
+                # Set individual speaker targets for each sample
+                asr_model.set_speaker_targets(spk_targets[i : i + 1], bg_spk_targets[i : i + 1])
+                logprobs_ins, _ = asr_model.forward(
+                    input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1]
+                )
+                logprobs_instance.append(logprobs_ins)
+            logits_instance = torch.cat(logprobs_instance, 0)
+
+            # batch size 4
+            asr_model.set_speaker_targets(spk_targets, bg_spk_targets)
+            logprobs_batch, _ = asr_model.forward(input_signal=input_signal, input_signal_length=length)
+
+        assert logits_instance.shape == logprobs_batch.shape
+        diff = torch.mean(torch.abs(logits_instance - logprobs_batch))
+        assert diff <= 1e-5  # Allow slightly higher tolerance for speaker processing
+        diff = torch.max(torch.abs(logits_instance - logprobs_batch))
+        assert diff <= 1e-5
+
+    @pytest.mark.unit
+    def test_speaker_target_setting(self, asr_model):
+        """Test speaker target setting functionality."""
+        batch_size = 2
+        target_length = 32
+
+        spk_targets = torch.randint(0, 2, (batch_size, target_length), dtype=torch.float32)
+        bg_spk_targets = torch.randint(0, 2, (batch_size, target_length), dtype=torch.float32)
+
+        # Test setting speaker targets
+        asr_model.set_speaker_targets(spk_targets, bg_spk_targets)
+        assert torch.equal(asr_model.spk_targets, spk_targets)
+        if asr_model.add_bg_spk_kernel:
+            assert torch.equal(asr_model.bg_spk_targets, bg_spk_targets)
+
+        # Test clearing speaker targets
+        asr_model.set_speaker_targets(None, None)
+        assert asr_model.spk_targets is None
+        if asr_model.add_bg_spk_kernel:
+            assert asr_model.bg_spk_targets is None

tests/collections/asr/test_asr_parts_submodules_batchnorm.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import OrderedDict
+
+import pytest
+import torch
+import torch.nn as nn
+
+from nemo.collections.asr.parts.submodules.batchnorm import (
+    FusedBatchNorm1d,
+    replace_bn_with_fused_bn,
+    replace_bn_with_fused_bn_all,
+)
+
+
+class TestFusedBatchNorm1d:
+    @pytest.mark.unit
+    def test_constructor(self):
+        num_features = 10
+        fused_bn = FusedBatchNorm1d(num_features=num_features)
+        assert fused_bn.weight.shape[0] == num_features
+        assert fused_bn.bias.shape[0] == num_features
+        # check initialization: weight is ones, bias is zeros (identity)
+        assert torch.allclose(fused_bn.weight, torch.ones(num_features))
+        assert torch.allclose(fused_bn.bias, torch.zeros(num_features))
+
+    @pytest.mark.unit
+    def test_from_batchnorm(self):
+        num_features = 10
+
+        # construct batchnorm
+        bn = nn.BatchNorm1d(num_features=num_features)
+
+        # update bn stats
+        bn.train()
+        batch_size = 4
+        for _ in range(10):
+            _ = bn(torch.rand(batch_size, num_features))
+
+        # test eval mode is equivalent
+        fused_bn = FusedBatchNorm1d.from_batchnorm(bn)
+        bn.eval()
+
+        sample_2d = torch.rand(batch_size, num_features)
+        assert torch.allclose(bn(sample_2d), fused_bn(sample_2d))
+
+        sample_3d = torch.rand(batch_size, num_features, 5)
+        assert torch.allclose(bn(sample_3d), fused_bn(sample_3d))
+
+
+class TestReplaceBNWithFusedBN:
+    @pytest.mark.unit
+    def test_replace_bn_with_fused_bn(self):
+        model = nn.Sequential(
+            OrderedDict(
+                [
+                    ("linear1", nn.Linear(1, 10)),
+                    ("bn1", nn.BatchNorm1d(10)),
+                    ("relu1", nn.ReLU()),
+                    ("linear2", nn.Linear(10, 11)),
+                    ("bn2", nn.BatchNorm1d(11)),
+                    (
+                        "submodule1",
+                        nn.Sequential(OrderedDict([("linear3", nn.Linear(11, 12)), ("bn3", nn.BatchNorm1d(12))])),
+                    ),
+                ]
+            )
+        )
+        replace_bn_with_fused_bn(model, "submodule1.bn3")
+        assert isinstance(model.bn1, nn.BatchNorm1d)
+        assert isinstance(model.bn2, nn.BatchNorm1d)
+        assert isinstance(model.submodule1.bn3, FusedBatchNorm1d)
+
+    @pytest.mark.unit
+    def test_replace_bn_with_fused_bn_all(self):
+        model = nn.Sequential(
+            OrderedDict(
+                [
+                    ("linear1", nn.Linear(1, 10)),
+                    ("bn1", nn.BatchNorm1d(10)),
+                    ("relu1", nn.ReLU()),
+                    ("linear2", nn.Linear(10, 11)),
+                    ("bn2", nn.BatchNorm1d(11)),
+                    (
+                        "submodule1",
+                        nn.Sequential(OrderedDict([("linear3", nn.Linear(11, 12)), ("bn3", nn.BatchNorm1d(12))])),
+                    ),
+                ]
+            )
+        )
+        replace_bn_with_fused_bn_all(model)
+        assert isinstance(model.bn1, FusedBatchNorm1d)
+        assert isinstance(model.bn2, FusedBatchNorm1d)
+        assert isinstance(model.submodule1.bn3, FusedBatchNorm1d)

tests/collections/asr/test_asr_regression_model.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import pytest
+from omegaconf import DictConfig
+
+from nemo.collections.asr.models.classification_models import EncDecRegressionModel
+
+
+@pytest.fixture()
+def speech_regression_model():
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': 32,
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                }
+            ],
+        },
+    }
+
+    decoder = {
+        'cls': 'nemo.collections.asr.modules.conv_asr.ConvASRDecoderClassification',
+        'params': {'feat_in': 32, 'return_logits': True, 'num_classes': 1},
+    }
+
+    modelConfig = DictConfig(
+        {
+            'preprocessor': DictConfig(preprocessor),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'labels': None,
+            'is_regression_task': True,
+        }
+    )
+    model = EncDecRegressionModel(cfg=modelConfig)
+    return model
+
+
+class TestEncDecRegressionModel:
+    @pytest.mark.unit
+    def test_constructor(self, speech_regression_model):
+        asr_model = speech_regression_model.train()
+
+        conv_cnt = (64 * 32 * 1 + 32) + (64 * 1 * 1 + 32)  # separable kernel + bias + pointwise kernel + bias
+        bn_cnt = (4 * 32) * 2  # 2 * moving averages
+        dec_cnt = 32 * 1 + 1  # fc + bias
+
+        param_count = conv_cnt + bn_cnt + dec_cnt
+        assert asr_model.num_weights == param_count
+
+        # Check to/from config_dict:
+        confdict = asr_model.to_config_dict()
+        instance2 = EncDecRegressionModel.from_config_dict(confdict)
+
+        assert isinstance(instance2, EncDecRegressionModel)
+
+    @pytest.mark.unit
+    def test_transcription(self, speech_regression_model, test_data_dir):
+
+        audio_filenames = ['an22-flrp-b.wav', 'an90-fbbh-b.wav']
+        audio_paths = [os.path.join(test_data_dir, "asr", "train", "an4", "wav", fp) for fp in audio_filenames]
+
+        model = speech_regression_model.eval()
+
+        # Test Top 1 classification transcription
+        results = model.transcribe(audio_paths, batch_size=2)
+        assert len(results) == 2

tests/collections/asr/test_asr_rnnt_encoder_model_bpe.py

@@ -0,0 +1,372 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import shutil
+import tempfile
+
+import pytest
+import torch
+from lhotse import CutSet, MonoCut
+from lhotse.testing.dummies import DummyManifest
+from omegaconf import DictConfig
+
+from nemo.collections.asr.data.audio_to_text_lhotse import LhotseSpeechToTextBpeDataset
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.models.rnnt_bpe_models import EncDecRNNTBPEModel
+from nemo.collections.asr.parts.submodules import rnnt_beam_decoding as beam_decode
+from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.collections.common import tokenizers
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+
+NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cpu_is_supported(
+    __NUMBA_MINIMUM_VERSION__
+) or numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+
+
+@pytest.fixture()
+def asr_model(test_data_dir):
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+
+    model_defaults = {'enc_hidden': 1024, 'pred_hidden': 64}
+
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': model_defaults['enc_hidden'],
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                }
+            ],
+        },
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.RNNTDecoder',
+        'prednet': {
+            'pred_hidden': model_defaults['pred_hidden'],
+            'pred_rnn_layers': 1,
+        },
+    }
+
+    joint = {
+        '_target_': 'nemo.collections.asr.modules.RNNTJoint',
+        'jointnet': {
+            'joint_hidden': 32,
+            'activation': 'relu',
+        },
+    }
+
+    decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 30}}
+
+    tokenizer = {'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'}
+
+    loss = {'loss_name': 'default', 'warprnnt_numba_kwargs': {'fastemit_lambda': 0.001}}
+
+    modelConfig = DictConfig(
+        {
+            'preprocessor': DictConfig(preprocessor),
+            'model_defaults': DictConfig(model_defaults),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'joint': DictConfig(joint),
+            'tokenizer': DictConfig(tokenizer),
+            'decoding': DictConfig(decoding),
+            'loss': DictConfig(loss),
+        }
+    )
+
+    model_instance = EncDecRNNTBPEModel(cfg=modelConfig)
+    return model_instance
+
+
+class NestedRNNTModel(ASRModel):
+    def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None):
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        if 'inner_model' in self.cfg:
+            self.register_nemo_submodule(
+                "inner_model", config_field="inner_model", model=EncDecRNNTBPEModel(self.cfg.inner_model)
+            )
+        else:
+            # Restore a model from pretrained checkpoint
+            self.register_nemo_submodule(
+                "inner_model",
+                config_field="inner_model",
+                model=ASRModel.from_pretrained('stt_en_conformer_transducer_small', map_location='cpu'),
+            )
+
+        self.linear = torch.nn.Linear(
+            self.inner_model.tokenizer.vocab_size + 1, self.inner_model.tokenizer.vocab_size + 1
+        )
+        self.inner_model.freeze()
+
+    setup_training_data = lambda *args, **kwargs: None
+    setup_validation_data = lambda *args, **kwargs: None
+    transcribe = lambda *args, **kwargs: []
+
+
+class TestEncDecRNNTBPEModel:
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_constructor(self, asr_model):
+        asr_model.train()
+        # TODO: make proper config and assert correct number of weights
+        # Check to/from config_dict:
+        confdict = asr_model.to_config_dict()
+        instance2 = EncDecRNNTBPEModel.from_config_dict(confdict)
+        assert isinstance(instance2, EncDecRNNTBPEModel)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_forward(self, asr_model):
+        asr_model = asr_model.eval()
+
+        asr_model.preprocessor.featurizer.dither = 0.0
+        asr_model.preprocessor.featurizer.pad_to = 0
+
+        asr_model.compute_eval_loss = False
+
+        input_signal = torch.randn(size=(4, 512))
+        length = torch.randint(low=321, high=500, size=[4])
+
+        with torch.no_grad():
+            # batch size 1
+            logprobs_instance = []
+            for i in range(input_signal.size(0)):
+                logprobs_ins, _ = asr_model.forward(
+                    input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1]
+                )
+                logprobs_instance.append(logprobs_ins)
+            logits_instance = torch.cat(logprobs_instance, 0)
+
+            # batch size 4
+            logprobs_batch, _ = asr_model.forward(input_signal=input_signal, input_signal_length=length)
+
+        assert logits_instance.shape == logprobs_batch.shape
+        diff = torch.mean(torch.abs(logits_instance - logprobs_batch))
+        assert diff <= 1e-6
+        diff = torch.max(torch.abs(logits_instance - logprobs_batch))
+        assert diff <= 1e-6
+
+    @pytest.mark.unit
+    def test_predict_step(self, asr_model):
+        asr_model = asr_model.eval()
+        cuts = DummyManifest(CutSet, begin_id=0, end_id=1, with_data=True)
+        dataset = LhotseSpeechToTextBpeDataset(tokenizer=asr_model.tokenizer, return_cuts=True)
+        batch = dataset[cuts]
+        outputs = asr_model.predict_step(batch, 0)
+        assert len(outputs) == 1
+        assert len(outputs[0]) == 2
+        assert isinstance(outputs[0][0], MonoCut)
+        assert isinstance(outputs[0][1], Hypothesis)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_save_restore_artifact(self, asr_model):
+        asr_model.train()
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = os.path.join(tmp_dir, 'rnnt_bpe.nemo')
+            asr_model.save_to(path)
+
+            new_model = EncDecRNNTBPEModel.restore_from(path)
+            assert isinstance(new_model, type(asr_model))
+            assert new_model.vocab_path.endswith('_vocab.txt')
+
+            assert len(new_model.tokenizer.tokenizer.get_vocab()) == 128
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_save_restore_artifact_spe(self, asr_model, test_data_dir):
+        asr_model.train()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128")
+            asr_model.change_vocabulary(new_tokenizer_dir=tokenizer_dir, new_tokenizer_type='bpe')
+
+            save_path = os.path.join(tmpdir, 'ctc_bpe.nemo')
+            asr_model.train()
+            asr_model.save_to(save_path)
+
+            new_model = EncDecRNNTBPEModel.restore_from(save_path)
+            assert isinstance(new_model, type(asr_model))
+            assert isinstance(new_model.tokenizer, tokenizers.SentencePieceTokenizer)
+            assert new_model.model_path.endswith('_tokenizer.model')
+            assert new_model.vocab_path.endswith('_vocab.txt')
+            assert new_model.spe_vocab_path.endswith('_tokenizer.vocab')
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_save_restore_artifact_agg(self, asr_model, test_data_dir):
+        tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128")
+        tok_en = {"dir": tokenizer_dir, "type": "wpe"}
+        # the below is really an english tokenizer but we pretend it is spanish
+        tok_es = {"dir": tokenizer_dir, "type": "wpe"}
+        tcfg = DictConfig({"type": "agg", "langs": {"en": tok_en, "es": tok_es}})
+        with tempfile.TemporaryDirectory() as tmpdir:
+            asr_model.change_vocabulary(new_tokenizer_dir=tcfg, new_tokenizer_type="agg")
+
+            save_path = os.path.join(tmpdir, "ctc_agg.nemo")
+            asr_model.train()
+            asr_model.save_to(save_path)
+
+            new_model = EncDecRNNTBPEModel.restore_from(save_path)
+            assert isinstance(new_model, type(asr_model))
+            assert isinstance(new_model.tokenizer, tokenizers.AggregateTokenizer)
+
+            # should be double
+            assert new_model.tokenizer.tokenizer.vocab_size == 264
+            assert len(new_model.tokenizer.tokenizer.get_vocab()) == 264
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_vocab_change(self, test_data_dir, asr_model):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            old_tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
+            new_tokenizer_dir = os.path.join(tmpdir, 'tokenizer')
+
+            os.makedirs(new_tokenizer_dir, exist_ok=True)
+            shutil.copy2(old_tokenizer_dir, new_tokenizer_dir)
+
+            nw1 = asr_model.num_weights
+            asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe')
+            # No change
+            assert nw1 == asr_model.num_weights
+
+            with open(os.path.join(new_tokenizer_dir, 'vocab.txt'), 'a+') as f:
+                f.write("!\n")
+                f.write('$\n')
+                f.write('@\n')
+
+            asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe')
+
+            # rnn embedding + joint + bias
+            pred_embedding = 3 * (asr_model.decoder.pred_hidden)
+            joint_joint = 3 * (asr_model.joint.joint_hidden + 1)
+            assert asr_model.num_weights == (nw1 + (pred_embedding + joint_joint))
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_decoding_change(self, asr_model):
+        assert isinstance(asr_model.decoding.decoding, greedy_decode.GreedyBatchedRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'greedy'
+        new_strategy.greedy = DictConfig({'max_symbols': 10})
+        asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(asr_model.decoding.decoding, greedy_decode.GreedyRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'beam'
+        new_strategy.beam = DictConfig({'beam_size': 1})
+        asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert asr_model.decoding.decoding.search_type == "default"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'beam'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert asr_model.decoding.decoding.search_type == "default"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'tsd'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert asr_model.decoding.decoding.search_type == "tsd"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'alsd'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert asr_model.decoding.decoding.search_type == "alsd"
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    def test_save_restore_nested_model(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model = NestedRNNTModel(cfg=DictConfig({}), trainer=None)
+            path = os.path.join(tmp_dir, 'rnnt_bpe.nemo')
+            model.save_to(path)
+
+            new_model = NestedRNNTModel.restore_from(path, map_location='cpu')
+            assert model.__class__.__name__ == NestedRNNTModel.__name__
+            assert new_model.__class__.__name__ == NestedRNNTModel.__name__
+            assert isinstance(new_model, type(model))
+
+            assert new_model.inner_model.vocab_path.endswith('_vocab.txt')
+            assert len(new_model.inner_model.tokenizer.tokenizer.get_vocab()) == 1024
+
+            # Unpack the nemo file
+            NestedRNNTModel._save_restore_connector._unpack_nemo_file(path, tmp_dir)
+
+            # Check size of the checkpoint, which contains weights from pretrained model + linear layer
+            fp_weights = os.path.join(tmp_dir, 'model_weights.ckpt')
+            assert os.path.getsize(fp_weights) > 50 * (2**20)  # Assert the weights are more than 50 MB
+
+            # Check if param after restoration is exact match
+            original_state_dict = model.inner_model.state_dict()
+            new_state_dict = new_model.inner_model.state_dict()
+
+            for (old_name, old_param), (new_name, new_param) in zip(
+                original_state_dict.items(), new_state_dict.items()
+            ):
+                assert old_name == new_name
+                assert (old_param - new_param).float().abs().mean() < 1e-6

tests/collections/asr/test_asr_rnnt_encoder_model_bpe_prompt.py

@@ -0,0 +1,409 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import shutil
+import tempfile
+
+import pytest
+import torch
+from omegaconf import DictConfig
+
+from nemo.collections.asr.models.rnnt_bpe_models_prompt import EncDecRNNTBPEModelWithPrompt
+from nemo.collections.asr.parts.submodules import rnnt_beam_decoding as beam_decode
+from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode
+from nemo.collections.common import tokenizers
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+
+NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cpu_is_supported(
+    __NUMBA_MINIMUM_VERSION__
+) or numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+
+
+@pytest.fixture()
+def rnnt_asr_model_with_prompt(test_data_dir):
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+
+    model_defaults = {
+        'enc_hidden': 1024,
+        'pred_hidden': 640,
+        'initialize_prompt_feature': True,  # Enable prompt feature initialization
+        'num_prompts': 128,
+        'prompt_dictionary': {
+            'en_US': 0,
+            'es_ES': 1,
+            'fr_FR': 2,
+            'de_DE': 3,
+        },
+    }
+
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': model_defaults['enc_hidden'],
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                }
+            ],
+        },
+    }
+
+    decoder = {
+        '_target_': 'nemo.collections.asr.modules.RNNTDecoder',
+        'prednet': {
+            'pred_hidden': model_defaults['pred_hidden'],
+            'pred_rnn_layers': 1,
+        },
+    }
+
+    joint = {
+        '_target_': 'nemo.collections.asr.modules.RNNTJoint',
+        'jointnet': {
+            'joint_hidden': 640,
+            'activation': 'relu',
+        },
+    }
+
+    decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 30}}
+
+    tokenizer = {'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'}
+
+    loss = {'loss_name': 'default', 'warprnnt_numba_kwargs': {'fastemit_lambda': 0.001}}
+
+    modelConfig = DictConfig(
+        {
+            'preprocessor': DictConfig(preprocessor),
+            'model_defaults': DictConfig(model_defaults),
+            'encoder': DictConfig(encoder),
+            'decoder': DictConfig(decoder),
+            'joint': DictConfig(joint),
+            'tokenizer': DictConfig(tokenizer),
+            'decoding': DictConfig(decoding),
+            'loss': DictConfig(loss),
+        }
+    )
+
+    model_instance = EncDecRNNTBPEModelWithPrompt(cfg=modelConfig)
+    return model_instance
+
+
+class TestEncDecRNNTBPEModelWithPrompt:
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_constructor(self, rnnt_asr_model_with_prompt):
+        rnnt_asr_model_with_prompt.train()
+        # Check to/from config_dict:
+        confdict = rnnt_asr_model_with_prompt.to_config_dict()
+        instance2 = EncDecRNNTBPEModelWithPrompt.from_config_dict(confdict)
+        assert isinstance(instance2, EncDecRNNTBPEModelWithPrompt)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_forward(self, rnnt_asr_model_with_prompt):
+        rnnt_asr_model_with_prompt = rnnt_asr_model_with_prompt.eval()
+
+        rnnt_asr_model_with_prompt.preprocessor.featurizer.dither = 0.0
+        rnnt_asr_model_with_prompt.preprocessor.featurizer.pad_to = 0
+
+        rnnt_asr_model_with_prompt.compute_eval_loss = False
+
+        input_signal = torch.randn(size=(4, 512))
+        length = torch.randint(low=321, high=500, size=[4])
+
+        # 1D prompt indices: one language id per sample in the batch.
+        prompt_indices = torch.tensor([0, 1, 2, 3], dtype=torch.long)
+
+        with torch.no_grad():
+            # batch size 1
+            logprobs_instance = []
+            for i in range(input_signal.size(0)):
+                logprobs_ins, _ = rnnt_asr_model_with_prompt.forward(
+                    input_signal=input_signal[i : i + 1],
+                    input_signal_length=length[i : i + 1],
+                    prompt_indices=prompt_indices[i : i + 1],
+                )
+                logprobs_instance.append(logprobs_ins)
+            logits_instance = torch.cat(logprobs_instance, 0)
+
+            # batch size 4
+            logprobs_batch, _ = rnnt_asr_model_with_prompt.forward(
+                input_signal=input_signal, input_signal_length=length, prompt_indices=prompt_indices
+            )
+
+        assert logits_instance.shape == logprobs_batch.shape
+        diff = torch.mean(torch.abs(logits_instance - logprobs_batch))
+        assert diff <= 1e-6
+        diff = torch.max(torch.abs(logits_instance - logprobs_batch))
+        assert diff <= 1e-6
+
+    @pytest.mark.unit
+    def test_predict_step(self, rnnt_asr_model_with_prompt):
+        rnnt_asr_model_with_prompt = rnnt_asr_model_with_prompt.eval()
+
+        # Create a simple batch manually
+        batch_size = 1
+        seq_len = 1600
+
+        # Create mock batch data
+        audio_signal = torch.randn(batch_size, seq_len)
+        audio_lengths = torch.tensor([seq_len])
+        transcript = torch.randint(0, 10, (batch_size, 10))
+        transcript_lengths = torch.tensor([10])
+        prompt_indices = torch.tensor([0], dtype=torch.long)  # language id 0
+
+        batch = (audio_signal, audio_lengths, transcript, transcript_lengths, prompt_indices)
+
+        outputs = rnnt_asr_model_with_prompt.predict_step(batch, 0)
+        assert len(outputs) == 1
+        # predict_step returns list of (sample_id, hyp_or_text) pairs
+        assert len(outputs[0]) == 2
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_save_restore_artifact(self, rnnt_asr_model_with_prompt):
+        rnnt_asr_model_with_prompt.train()
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = os.path.join(tmp_dir, 'rnnt_bpe_prompt.nemo')
+            rnnt_asr_model_with_prompt.save_to(path)
+
+            # restore_from is intentionally delegated to the parent EncDecRNNTBPEModel
+            # to avoid subclass-substitution that would hang on missing prompt_dictionary.
+            new_model = EncDecRNNTBPEModelWithPrompt.restore_from(path)
+            assert new_model.vocab_path.endswith('_vocab.txt')
+
+            assert len(new_model.tokenizer.tokenizer.get_vocab()) == 128
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_save_restore_artifact_spe(self, rnnt_asr_model_with_prompt, test_data_dir):
+        rnnt_asr_model_with_prompt.train()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128")
+            rnnt_asr_model_with_prompt.change_vocabulary(new_tokenizer_dir=tokenizer_dir, new_tokenizer_type='bpe')
+
+            save_path = os.path.join(tmpdir, 'rnnt_bpe_prompt.nemo')
+            rnnt_asr_model_with_prompt.train()
+            rnnt_asr_model_with_prompt.save_to(save_path)
+
+            new_model = EncDecRNNTBPEModelWithPrompt.restore_from(save_path)
+            assert isinstance(new_model.tokenizer, tokenizers.SentencePieceTokenizer)
+            assert new_model.model_path.endswith('_tokenizer.model')
+            assert new_model.vocab_path.endswith('_vocab.txt')
+            assert new_model.spe_vocab_path.endswith('_tokenizer.vocab')
+
+    @pytest.mark.unit
+    def test_save_restore_artifact_agg(self, rnnt_asr_model_with_prompt, test_data_dir):
+        tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128")
+        tok_en = {"dir": tokenizer_dir, "type": "wpe"}
+        # the below is really an english tokenizer but we pretend it is spanish
+        tok_es = {"dir": tokenizer_dir, "type": "wpe"}
+        tcfg = DictConfig({"type": "agg", "langs": {"en": tok_en, "es": tok_es}})
+        with tempfile.TemporaryDirectory() as tmpdir:
+            rnnt_asr_model_with_prompt.change_vocabulary(new_tokenizer_dir=tcfg, new_tokenizer_type="agg")
+
+            save_path = os.path.join(tmpdir, "rnnt_agg_prompt.nemo")
+            rnnt_asr_model_with_prompt.train()
+            rnnt_asr_model_with_prompt.save_to(save_path)
+
+            new_model = EncDecRNNTBPEModelWithPrompt.restore_from(save_path)
+            assert isinstance(new_model.tokenizer, tokenizers.AggregateTokenizer)
+
+            # Both source tokenizers are the same 132-token vocab; the AggregateTokenizer
+            # deduplicates 10 shared control tokens, so total = 132 + (132 - 10) = 254.
+            assert new_model.tokenizer.tokenizer.vocab_size == 264
+            assert len(new_model.tokenizer.tokenizer.get_vocab()) == 264
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_vocab_change(self, test_data_dir, rnnt_asr_model_with_prompt):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            old_tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
+            new_tokenizer_dir = os.path.join(tmpdir, 'tokenizer')
+
+            os.makedirs(new_tokenizer_dir, exist_ok=True)
+            shutil.copy2(old_tokenizer_dir, new_tokenizer_dir)
+
+            nw1 = rnnt_asr_model_with_prompt.num_weights
+            rnnt_asr_model_with_prompt.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe')
+            # No change
+            assert nw1 == rnnt_asr_model_with_prompt.num_weights
+
+            with open(os.path.join(new_tokenizer_dir, 'vocab.txt'), 'a+') as f:
+                f.write("!\n")
+                f.write('$\n')
+                f.write('@\n')
+
+            rnnt_asr_model_with_prompt.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe')
+
+            # rnn embedding + joint + bias (no CTC decoder in RNNT-only model)
+            pred_embedding = 3 * (rnnt_asr_model_with_prompt.decoder.pred_hidden)
+            joint_joint = 3 * (rnnt_asr_model_with_prompt.joint.joint_hidden + 1)
+            assert rnnt_asr_model_with_prompt.num_weights == (nw1 + (pred_embedding + joint_joint))
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_decoding_change(self, rnnt_asr_model_with_prompt):
+        assert isinstance(rnnt_asr_model_with_prompt.decoding.decoding, greedy_decode.GreedyBatchedRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'greedy'
+        new_strategy.greedy = DictConfig({'max_symbols': 10})
+        rnnt_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(rnnt_asr_model_with_prompt.decoding.decoding, greedy_decode.GreedyRNNTInfer)
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'beam'
+        new_strategy.beam = DictConfig({'beam_size': 1})
+        rnnt_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(rnnt_asr_model_with_prompt.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert rnnt_asr_model_with_prompt.decoding.decoding.search_type == "default"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'beam'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        rnnt_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(rnnt_asr_model_with_prompt.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert rnnt_asr_model_with_prompt.decoding.decoding.search_type == "default"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'tsd'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        rnnt_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(rnnt_asr_model_with_prompt.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert rnnt_asr_model_with_prompt.decoding.decoding.search_type == "tsd"
+
+        new_strategy = DictConfig({})
+        new_strategy.strategy = 'alsd'
+        new_strategy.beam = DictConfig({'beam_size': 2})
+        rnnt_asr_model_with_prompt.change_decoding_strategy(decoding_cfg=new_strategy)
+        assert isinstance(rnnt_asr_model_with_prompt.decoding.decoding, beam_decode.BeamRNNTInfer)
+        assert rnnt_asr_model_with_prompt.decoding.decoding.search_type == "alsd"
+
+    @pytest.mark.unit
+    def test_input_output_types_with_prompt(self, rnnt_asr_model_with_prompt):
+        """Test that input/output types include prompt-specific types."""
+        input_types = rnnt_asr_model_with_prompt.input_types
+        output_types = rnnt_asr_model_with_prompt.output_types
+
+        # Check that prompt_indices is included in input types (1D, per-sample language id)
+        assert 'prompt_indices' in input_types
+        prompt_axes = input_types['prompt_indices'].axes
+        assert len(prompt_axes) == 1  # 1D tensor [B]
+
+        # Check standard input types are present
+        assert 'input_signal' in input_types
+        assert 'input_signal_length' in input_types
+
+        # Check output types
+        assert 'outputs' in output_types
+        assert 'encoded_lengths' in output_types
+
+    @pytest.mark.unit
+    def test_prompt_feature_initialization(self, rnnt_asr_model_with_prompt):
+        """Test that prompt feature initialization works correctly."""
+        # Test that the model has prompt-related attributes
+        assert hasattr(rnnt_asr_model_with_prompt, 'concat')
+        assert hasattr(rnnt_asr_model_with_prompt, 'num_prompts')
+        assert hasattr(rnnt_asr_model_with_prompt, 'prompt_kernel')
+
+        # Test that concat is enabled
+        assert rnnt_asr_model_with_prompt.concat == True
+
+        # Test prompt kernel dimensions
+        expected_input_size = (
+            rnnt_asr_model_with_prompt.num_prompts + rnnt_asr_model_with_prompt._cfg.model_defaults.enc_hidden
+        )
+        expected_output_size = rnnt_asr_model_with_prompt._cfg.model_defaults.enc_hidden
+
+        # Check first layer of prompt kernel
+        first_layer = rnnt_asr_model_with_prompt.prompt_kernel[0]
+        assert first_layer.in_features == expected_input_size
+        assert first_layer.out_features == expected_output_size * 2
+
+    @pytest.mark.unit
+    def test_set_inference_prompt(self, rnnt_asr_model_with_prompt):
+        """Test that set_inference_prompt accepts known languages and rejects unknown ones."""
+        # Known language from the prompt_dictionary fixture
+        rnnt_asr_model_with_prompt.set_inference_prompt('en_US')
+        assert rnnt_asr_model_with_prompt._inference_prompt_index == 0
+
+        rnnt_asr_model_with_prompt.set_inference_prompt('de_DE')
+        assert rnnt_asr_model_with_prompt._inference_prompt_index == 3
+
+        # Unknown language should raise
+        with pytest.raises(ValueError):
+            rnnt_asr_model_with_prompt.set_inference_prompt('zz_ZZ')
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_forward_different_prompts_produce_different_outputs(self, rnnt_asr_model_with_prompt):
+        """Different prompt_indices should yield different encoded outputs (prompt actually conditions)."""
+        rnnt_asr_model_with_prompt = rnnt_asr_model_with_prompt.eval()
+        rnnt_asr_model_with_prompt.preprocessor.featurizer.dither = 0.0
+        rnnt_asr_model_with_prompt.preprocessor.featurizer.pad_to = 0
+
+        input_signal = torch.randn(size=(2, 512))
+        length = torch.tensor([512, 512])
+
+        with torch.no_grad():
+            out_a, _ = rnnt_asr_model_with_prompt.forward(
+                input_signal=input_signal,
+                input_signal_length=length,
+                prompt_indices=torch.tensor([0, 0], dtype=torch.long),
+            )
+            out_b, _ = rnnt_asr_model_with_prompt.forward(
+                input_signal=input_signal,
+                input_signal_length=length,
+                prompt_indices=torch.tensor([1, 1], dtype=torch.long),
+            )
+
+        assert out_a.shape == out_b.shape
+        # The prompt projection should make outputs differ for different language ids.
+        assert torch.max(torch.abs(out_a - out_b)) > 0

tests/collections/asr/test_asr_rope.py

@@ -0,0 +1,581 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.modules.conformer_encoder import ConformerEncoder
+from nemo.collections.asr.parts.submodules.multi_head_attention import RoPEMultiHeadAttention, RotaryPositionalEncoding
+
+
+def _build_encoder(
+    self_attention_model='rope',
+    n_layers=2,
+    d_model=64,
+    n_heads=4,
+    use_pytorch_sdpa=False,
+    use_pytorch_sdpa_backends=None,
+    rotary_fraction=1.0,
+    rope_base=10000.0,
+    pos_emb_max_len=256,
+):
+    return ConformerEncoder(
+        feat_in=80,
+        n_layers=n_layers,
+        d_model=d_model,
+        n_heads=n_heads,
+        self_attention_model=self_attention_model,
+        subsampling_factor=4,
+        subsampling_conv_channels=32,
+        pos_emb_max_len=pos_emb_max_len,
+        rotary_fraction=rotary_fraction,
+        rope_base=rope_base,
+        use_pytorch_sdpa=use_pytorch_sdpa,
+        use_pytorch_sdpa_backends=use_pytorch_sdpa_backends,
+        dropout=0.0,
+        dropout_att=0.0,
+        dropout_emb=0.0,
+        dropout_pre_encoder=0.0,
+    ).eval()
+
+
+class TestRotaryPositionalEncoding:
+    @pytest.mark.unit
+    def test_rejects_invalid_rotary_fraction(self):
+        with pytest.raises(ValueError):
+            RotaryPositionalEncoding(d_k=16, rotary_fraction=0.0)
+        with pytest.raises(ValueError):
+            RotaryPositionalEncoding(d_k=16, rotary_fraction=1.5)
+
+    @pytest.mark.unit
+    def test_rejects_odd_effective_dim(self):
+        # d_k * rotary_fraction = 16 * 0.1875 = 3, which is odd
+        with pytest.raises(ValueError):
+            RotaryPositionalEncoding(d_k=16, rotary_fraction=0.1875)
+
+    @pytest.mark.unit
+    def test_extend_pe_grows_buffers(self):
+        pe = RotaryPositionalEncoding(d_k=16, max_len=128)
+        pe.extend_pe(64, device=torch.device('cpu'), dtype=torch.float32)
+        assert pe.cos.shape == (64, 16)
+        pe.extend_pe(128, device=torch.device('cpu'), dtype=torch.float32)
+        assert pe.cos.shape == (128, 16)
+        # No-op when buffer is already large enough.
+        prev = pe.cos.data_ptr()
+        pe.extend_pe(64, device=torch.device('cpu'), dtype=torch.float32)
+        assert pe.cos.data_ptr() == prev
+
+    @pytest.mark.unit
+    def test_forward_first_token_is_identity(self):
+        # Position 0 has zero phase, so cos=1, sin=0 -> rotation is identity.
+        pe = RotaryPositionalEncoding(d_k=16, rotary_fraction=1.0)
+        pe.extend_pe(32, device=torch.device('cpu'), dtype=torch.float32)
+        q = torch.randn(2, 4, 8, 16)
+        k = torch.randn(2, 4, 8, 16)
+        q_rot, k_rot = pe(q, k)
+        assert q_rot.shape == q.shape
+        assert k_rot.shape == k.shape
+        assert torch.allclose(q_rot[:, :, 0, :], q[:, :, 0, :], atol=1e-6)
+        assert torch.allclose(k_rot[:, :, 0, :], k[:, :, 0, :], atol=1e-6)
+
+    @pytest.mark.unit
+    def test_partial_rotation_leaves_tail_unchanged(self):
+        pe = RotaryPositionalEncoding(d_k=16, rotary_fraction=0.5)
+        pe.extend_pe(32, device=torch.device('cpu'), dtype=torch.float32)
+        q = torch.randn(2, 4, 8, 16)
+        k = torch.randn(2, 4, 8, 16)
+        q_rot, k_rot = pe(q, k)
+        # The last (d_k - d_k_rot) = 8 dims of each head must pass through untouched.
+        assert torch.allclose(q_rot[..., pe.d_k_rot :], q[..., pe.d_k_rot :])
+        assert torch.allclose(k_rot[..., pe.d_k_rot :], k[..., pe.d_k_rot :])
+
+    @pytest.mark.unit
+    def test_dot_product_translation_invariance(self):
+        # The defining property of RoPE: for the same q and k content, <q_m, k_n>
+        # depends only on the position difference (m - n). Pick two (m, n) pairs
+        # that share the same difference and assert the dot products agree.
+        pe = RotaryPositionalEncoding(d_k=16, rotary_fraction=1.0)
+        pe.extend_pe(64, device=torch.device('cpu'), dtype=torch.float32)
+
+        torch.manual_seed(0)
+        q_content = torch.randn(1, 1, 1, 16)
+        k_content = torch.randn(1, 1, 1, 16)
+
+        def dot_at(m, n):
+            cos_q = pe.cos[m : m + 1].view(1, 1, 1, 16)
+            sin_q = pe.sin[m : m + 1].view(1, 1, 1, 16)
+            cos_k = pe.cos[n : n + 1].view(1, 1, 1, 16)
+            sin_k = pe.sin[n : n + 1].view(1, 1, 1, 16)
+            q_r = pe._apply_rotary(q_content, cos_q, sin_q)
+            k_r = pe._apply_rotary(k_content, cos_k, sin_k)
+            return (q_r * k_r).sum()
+
+        # Three (m, n) pairs with the same difference n - m = 3.
+        d_a = dot_at(2, 5)
+        d_b = dot_at(10, 13)
+        d_c = dot_at(40, 43)
+        assert torch.allclose(d_a, d_b, atol=1e-5)
+        assert torch.allclose(d_a, d_c, atol=1e-5)
+
+        # Sanity: a different position difference must yield a different dot product
+        # (otherwise the rotation is a no-op or degenerate).
+        d_diff = dot_at(2, 7)  # difference 5
+        assert not torch.allclose(d_a, d_diff, atol=1e-3)
+
+    @pytest.mark.unit
+    def test_rotation_is_not_identity(self):
+        # Confirm RoPE actually mutates Q/K at non-zero positions.
+        pe = RotaryPositionalEncoding(d_k=16, rotary_fraction=1.0)
+        pe.extend_pe(32, device=torch.device('cpu'), dtype=torch.float32)
+        q = torch.randn(1, 1, 8, 16)
+        k = torch.randn(1, 1, 8, 16)
+        q_rot, k_rot = pe(q, k)
+        # Tokens after position 0 must change.
+        assert not torch.allclose(q_rot[:, :, 1:, :], q[:, :, 1:, :], atol=1e-3)
+        assert not torch.allclose(k_rot[:, :, 1:, :], k[:, :, 1:, :], atol=1e-3)
+
+    @pytest.mark.unit
+    def test_norm_preservation(self):
+        # Rotation is unitary: ||q_rot[..., t, :]||_2 == ||q[..., t, :]||_2 per (batch, head, t).
+        # Catches scaling bugs in _apply_rotary.
+        pe = RotaryPositionalEncoding(d_k=16, rotary_fraction=1.0)
+        pe.extend_pe(64, device=torch.device('cpu'), dtype=torch.float32)
+        q = torch.randn(2, 4, 16, 16)
+        k = torch.randn(2, 4, 16, 16)
+        q_rot, k_rot = pe(q, k)
+        q_norm_in = torch.linalg.norm(q, dim=-1)
+        q_norm_out = torch.linalg.norm(q_rot, dim=-1)
+        k_norm_in = torch.linalg.norm(k, dim=-1)
+        k_norm_out = torch.linalg.norm(k_rot, dim=-1)
+        assert torch.allclose(q_norm_in, q_norm_out, atol=1e-5)
+        assert torch.allclose(k_norm_in, k_norm_out, atol=1e-5)
+
+    @pytest.mark.unit
+    def test_reference_equivalence(self):
+        # Slow split-half RoPE reference written in explicit-2D-rotation form
+        # (no _rotate_half trick, no cat-duplicated cos/sin). Same math as the
+        # production code expressed via a disjoint code path, so a bug in either
+        # _rotate_half or the cos/sin layout would surface here.
+        d_k = 16
+        pe = RotaryPositionalEncoding(d_k=d_k, rotary_fraction=1.0)
+        pe.extend_pe(32, device=torch.device('cpu'), dtype=torch.float32)
+
+        torch.manual_seed(0)
+        q = torch.randn(1, 1, 8, d_k)
+        k = torch.randn(1, 1, 8, d_k)
+        q_rot, k_rot = pe(q, k)
+
+        d_half = d_k // 2
+        positions = torch.arange(8, dtype=torch.float32)
+        theta = positions[:, None] * pe.inv_freq[None, :]  # (T, d_half)
+        c = theta.cos()
+        s = theta.sin()
+
+        def rope_ref(x):
+            # Rotate each (x[..., i], x[..., i + d_half]) pair by angle theta[t, i].
+            x_a = x[..., :d_half]
+            x_b = x[..., d_half:]
+            y_a = x_a * c - x_b * s
+            y_b = x_a * s + x_b * c
+            return torch.cat((y_a, y_b), dim=-1)
+
+        assert torch.allclose(q_rot, rope_ref(q), atol=1e-6)
+        assert torch.allclose(k_rot, rope_ref(k), atol=1e-6)
+
+    @pytest.mark.unit
+    def test_extend_preserves_existing_positions(self):
+        # Extending the cos/sin buffers must not change the values at previously
+        # covered positions, otherwise streaming forward calls would silently
+        # produce different rotations across the extension boundary.
+        pe = RotaryPositionalEncoding(d_k=16, max_len=64)
+        pe.extend_pe(64, device=torch.device('cpu'), dtype=torch.float32)
+        cos_before = pe.cos[:64].clone()
+        sin_before = pe.sin[:64].clone()
+        pe.extend_pe(256, device=torch.device('cpu'), dtype=torch.float32)
+        assert torch.equal(pe.cos[:64], cos_before)
+        assert torch.equal(pe.sin[:64], sin_before)
+
+    @pytest.mark.unit
+    def test_non_contiguous_inputs(self):
+        # Real-world callers may pass non-contiguous Q/K (e.g. from .transpose()).
+        # The rotation must produce the same result as on the contiguous version.
+        pe = RotaryPositionalEncoding(d_k=16, rotary_fraction=1.0)
+        pe.extend_pe(32, device=torch.device('cpu'), dtype=torch.float32)
+
+        # Build (B, T, H, D) and transpose to (B, H, T, D) -> non-contiguous.
+        q_btnd = torch.randn(2, 8, 4, 16)
+        k_btnd = torch.randn(2, 8, 4, 16)
+        q_nc = q_btnd.transpose(1, 2)
+        k_nc = k_btnd.transpose(1, 2)
+        assert not q_nc.is_contiguous() and not k_nc.is_contiguous()
+
+        q_rot_nc, k_rot_nc = pe(q_nc, k_nc)
+        q_rot_c, k_rot_c = pe(q_nc.contiguous(), k_nc.contiguous())
+        assert torch.allclose(q_rot_nc, q_rot_c, atol=1e-6)
+        assert torch.allclose(k_rot_nc, k_rot_c, atol=1e-6)
+
+
+class TestRoPEMultiHeadAttention:
+    @pytest.mark.unit
+    def test_rejects_pos_enc_with_wrong_d_k(self):
+        # n_feat / n_head = 64 / 4 = 16, but pos_enc was built with d_k=32.
+        bad_pe = RotaryPositionalEncoding(d_k=32, max_len=64)
+        with pytest.raises(ValueError):
+            RoPEMultiHeadAttention(n_head=4, n_feat=64, dropout_rate=0.0, pos_enc=bad_pe)
+
+    @pytest.mark.unit
+    def test_v_unchanged_by_rotation(self):
+        # Confirm the rotation hook is called only with (q, k); V must never reach
+        # the positional encoder. Catches a future regression where someone adds
+        # V to the rotation hook signature.
+        pe = RotaryPositionalEncoding(d_k=16, max_len=32)
+        pe.extend_pe(32, device=torch.device('cpu'), dtype=torch.float32)
+        attn = RoPEMultiHeadAttention(n_head=4, n_feat=64, dropout_rate=0.0, pos_enc=pe).eval()
+
+        call_args = []
+        original_forward = pe.forward
+
+        def spy(q, k):
+            call_args.append((q.shape, k.shape))
+            return original_forward(q, k)
+
+        attn.pos_enc.forward = spy
+        x = torch.randn(2, 16, 64)
+        with torch.no_grad():
+            _ = attn(query=x, key=x, value=x, mask=None)
+        assert len(call_args) == 1
+        q_shape, k_shape = call_args[0]
+        # Both tensors have the same length (16); the layout is (B, H, T, d_k).
+        assert q_shape == (2, 4, 16, 16)
+        assert k_shape == (2, 4, 16, 16)
+
+    @pytest.mark.unit
+    def test_backward_smoke(self):
+        # Forward → loss → backward → every learnable param has a non-NaN, non-zero
+        # gradient. Mirrors test_transformer_encoder.py::test_backward_pass.
+        pe = RotaryPositionalEncoding(d_k=16, max_len=32)
+        pe.extend_pe(32, device=torch.device('cpu'), dtype=torch.float32)
+        attn = RoPEMultiHeadAttention(n_head=4, n_feat=64, dropout_rate=0.0, pos_enc=pe).train()
+        x = torch.randn(2, 8, 64, requires_grad=True)
+        out = attn(query=x, key=x, value=x, mask=None)
+        loss = out.sum()
+        loss.backward()
+        for name, param in attn.named_parameters():
+            assert param.grad is not None, f"No gradient for {name}"
+            assert not torch.isnan(param.grad).any(), f"NaN gradient for {name}"
+            assert (param.grad != 0).any(), f"All-zero gradient for {name}"
+
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    @pytest.mark.parametrize("backend", ['MATH', 'FLASH_ATTENTION', 'EFFICIENT_ATTENTION', 'CUDNN_ATTENTION'])
+    def test_sdpa_backend_smoke_gpu(self, backend):
+        # Each SDPA backend must run with RoPE pre-rotation under bf16 autocast
+        # (the production training path) without falling back or crashing on
+        # shape/dtype constraints. FLASH/EFFICIENT/CUDNN require fp16/bf16;
+        # bf16 satisfies all four.
+        pe = RotaryPositionalEncoding(d_k=16, max_len=32).to("cuda")
+        pe.extend_pe(32, device=torch.device('cuda'), dtype=torch.float32)
+        attn = (
+            RoPEMultiHeadAttention(
+                n_head=4,
+                n_feat=64,
+                dropout_rate=0.0,
+                pos_enc=pe,
+                use_pytorch_sdpa=True,
+                use_pytorch_sdpa_backends=[backend],
+            )
+            .to("cuda")
+            .eval()
+        )
+        x = torch.randn(2, 16, 64, device='cuda')
+        with torch.no_grad(), torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
+            out = attn(query=x, key=x, value=x, mask=None)
+        assert out.shape == (2, 16, 64)
+        assert torch.isfinite(out).all()
+
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    def test_autocast_gpu(self):
+        # Mixed-precision forward (CUDA autocast in bf16) must produce finite output.
+        # Exercises the interaction between RoPE's .to(q.dtype) cast and the
+        # avoid_float16_autocast_context wrapper in the base MHA.
+        pe = RotaryPositionalEncoding(d_k=16, max_len=32).to("cuda")
+        pe.extend_pe(32, device=torch.device('cuda'), dtype=torch.float32)
+        attn = RoPEMultiHeadAttention(n_head=4, n_feat=64, dropout_rate=0.0, pos_enc=pe).to("cuda").eval()
+        x = torch.randn(2, 16, 64, device='cuda')
+        with torch.no_grad(), torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
+            out = attn(query=x, key=x, value=x, mask=None)
+        assert torch.isfinite(out).all()
+
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "dtype,atol",
+        [(torch.float32, 1e-5), (torch.bfloat16, 5e-2), (torch.float16, 1e-2)],
+    )
+    def test_dtype_stability_gpu(self, dtype, atol):
+        # Forward in low precision must stay close to the fp32 reference.
+        pe = RotaryPositionalEncoding(d_k=16, max_len=32).to("cuda")
+        pe.extend_pe(32, device=torch.device('cuda'), dtype=torch.float32)
+        attn = RoPEMultiHeadAttention(n_head=4, n_feat=64, dropout_rate=0.0, pos_enc=pe).to("cuda").eval()
+        torch.manual_seed(0)
+        x = torch.randn(2, 16, 64, device='cuda')
+        with torch.no_grad():
+            out_ref = attn(query=x, key=x, value=x, mask=None)
+
+        # `attn.to(dtype=...)` converts every buffer including pos_enc.cos/sin to `dtype`.
+        attn_dt = attn.to(dtype=dtype)
+        x_dt = x.to(dtype=dtype)
+        with torch.no_grad():
+            out_dt = attn_dt(query=x_dt, key=x_dt, value=x_dt, mask=None)
+        assert torch.isfinite(out_dt).all()
+        assert torch.allclose(out_dt.float(), out_ref, atol=atol, rtol=atol)
+
+    @pytest.mark.unit
+    def test_streaming_matches_offline(self):
+        # The load-bearing test for the cache_len offset logic. Feeding the last
+        # `new_len` tokens with the first `cache_len` tokens as KV cache must
+        # reproduce the corresponding slice of the offline forward, because RoPE
+        # depends only on the (m - n) position difference and the cache layout
+        # preserves that.
+        pe = RotaryPositionalEncoding(d_k=16, max_len=64)
+        pe.extend_pe(32, device=torch.device('cpu'), dtype=torch.float32)
+        attn = RoPEMultiHeadAttention(n_head=4, n_feat=64, dropout_rate=0.0, pos_enc=pe).eval()
+        attn.cache_drop_size = 0  # required by update_cache
+
+        torch.manual_seed(7)
+        full_seq = torch.randn(1, 12, 64)
+        cache_len = 8
+
+        with torch.no_grad():
+            offline_out = attn(query=full_seq, key=full_seq, value=full_seq, mask=None)
+            new_query = full_seq[:, cache_len:]
+            cache = full_seq[:, :cache_len]
+            streaming_out, _ = attn(query=new_query, key=new_query, value=new_query, mask=None, cache=cache)
+
+        assert torch.allclose(streaming_out, offline_out[:, cache_len:], atol=1e-5)
+
+
+class TestConformerEncoderRoPE:
+    @pytest.mark.unit
+    def test_pos_enc_shared_across_layers(self):
+        # Critical: every layer must hold the same pos_enc instance so that the
+        # encoder's set_max_audio_length / extend_pe grows the buffers used by
+        # every layer (not just the first).
+        enc = _build_encoder()
+        assert all(layer.self_attn.pos_enc is enc.pos_enc for layer in enc.layers)
+        # And exercising the shared-extend path: growing the buffer once must be
+        # visible from every layer.
+        enc.pos_enc.extend_pe(512, device=torch.device('cpu'), dtype=torch.float32)
+        assert all(layer.self_attn.pos_enc.cos.size(0) >= 512 for layer in enc.layers)
+
+    @pytest.mark.unit
+    def test_sdpa_matches_manual(self):
+        # CPU fp32: SDPA falls back to MATH; verify it matches the manual matmul
+        # path so RoPE pre-rotation is applied consistently across both code paths.
+        enc_manual = _build_encoder(use_pytorch_sdpa=False)
+        enc_sdpa = _build_encoder(use_pytorch_sdpa=True)
+        enc_sdpa.load_state_dict(enc_manual.state_dict(), strict=False)
+        x = torch.randn(2, 80, 200)
+        lens = torch.tensor([200, 150])
+        with torch.no_grad():
+            o_manual, _ = enc_manual(audio_signal=x, length=lens)
+            o_sdpa, _ = enc_sdpa(audio_signal=x, length=lens)
+        assert torch.allclose(o_manual, o_sdpa, atol=1e-4, rtol=1e-4)
+
+    @pytest.mark.run_only_on('GPU')
+    @pytest.mark.unit
+    @pytest.mark.parametrize("backend", ['MATH', 'EFFICIENT_ATTENTION', 'CUDNN_ATTENTION'])
+    def test_sdpa_backend_matches_manual_gpu(self, backend):
+        # Forward + backward parity vs the manual path under bf16 autocast.
+        # RoPE applies to Q/K before the SDPA call, so each backend sees the
+        # same rotated tensors and must agree on outputs and gradients within
+        # bf16 tolerance. FLASH_ATTENTION is excluded because PyTorch rejects
+        # any non-null `attn_mask` on the Flash kernel and the encoder always
+        # emits a padding mask; CUDNN and EFFICIENT both accept bool masks.
+        # The MHA-level smoke test covers FLASH with mask=None.
+        enc_manual = _build_encoder(use_pytorch_sdpa=False).to("cuda")
+        enc_sdpa = _build_encoder(use_pytorch_sdpa=True, use_pytorch_sdpa_backends=[backend]).to("cuda")
+        enc_sdpa.load_state_dict(enc_manual.state_dict(), strict=False)
+
+        torch.manual_seed(0)
+        x_base = torch.randn(2, 80, 200, device='cuda')
+        x_manual = x_base.clone().requires_grad_(True)
+        x_sdpa = x_base.clone().requires_grad_(True)
+        lens = torch.tensor([200, 150], device='cuda')
+
+        with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
+            o_manual, _ = enc_manual(audio_signal=x_manual, length=lens)
+            o_sdpa, _ = enc_sdpa(audio_signal=x_sdpa, length=lens)
+
+        # Forward parity.
+        assert torch.allclose(o_manual.float(), o_sdpa.float(), atol=5e-2, rtol=5e-2)
+
+        # Backward parity: same loss, compare input grads and weight grads.
+        o_manual.sum().backward()
+        o_sdpa.sum().backward()
+        assert torch.allclose(x_manual.grad.float(), x_sdpa.grad.float(), atol=5e-2, rtol=5e-2)
+        for (n1, p1), (n2, p2) in zip(enc_manual.named_parameters(), enc_sdpa.named_parameters()):
+            assert n1 == n2
+            assert p1.grad is not None and p2.grad is not None, f"missing grad for {n1}"
+            assert torch.allclose(p1.grad.float(), p2.grad.float(), atol=5e-2, rtol=5e-2), f"grad mismatch for {n1}"
+
+    @pytest.mark.unit
+    def test_padding_does_not_leak(self):
+        # Output for the valid prefix must be invariant to the values in the
+        # padded suffix.
+        enc = _build_encoder()
+        x = torch.randn(1, 80, 200)
+        valid_len = 120
+        x1 = x.clone()
+        x1[0, :, valid_len:] = torch.randn(80, 200 - valid_len)
+        x2 = x.clone()
+        x2[0, :, valid_len:] = torch.randn(80, 200 - valid_len)
+        lens = torch.tensor([valid_len])
+        with torch.no_grad():
+            o1, _ = enc(audio_signal=x1, length=lens)
+            o2, _ = enc(audio_signal=x2, length=lens)
+        valid_out_len = valid_len // 4
+        assert torch.allclose(o1[..., :valid_out_len], o2[..., :valid_out_len], atol=1e-5)
+
+    @pytest.mark.unit
+    def test_change_attention_model_to_rope(self):
+        # Build a rel_pos encoder, swap to rope, run forward.
+        enc = _build_encoder(self_attention_model='rel_pos')
+        enc._cfg = OmegaConf.create(
+            {
+                'd_model': 64,
+                'n_heads': 4,
+                'dropout': 0.0,
+                'dropout_att': 0.0,
+                'dropout_emb': 0.0,
+                'pos_emb_max_len': 256,
+                'rope_base': 10000.0,
+                'rotary_fraction': 1.0,
+            }
+        )
+        enc.change_attention_model('rope')
+        assert isinstance(enc.pos_enc, RotaryPositionalEncoding)
+        assert all(layer.self_attn.pos_enc is enc.pos_enc for layer in enc.layers)
+        x = torch.randn(2, 80, 200)
+        lens = torch.tensor([200, 150])
+        out, _ = enc(audio_signal=x, length=lens)
+        assert torch.isfinite(out).all()
+
+    @pytest.mark.unit
+    def test_change_attention_model_preserves_use_bias_false(self):
+        # Regression: the swap loop in change_attention_model was building the new
+        # attention without forwarding use_bias, so a use_bias=False model silently
+        # gained randomly-initialised bias parameters after a swap.
+        from nemo.collections.asr.modules.conformer_encoder import ConformerEncoder
+
+        enc = ConformerEncoder(
+            feat_in=80,
+            n_layers=2,
+            d_model=64,
+            n_heads=4,
+            self_attention_model='rel_pos',
+            subsampling_factor=4,
+            subsampling_conv_channels=32,
+            pos_emb_max_len=256,
+            use_bias=False,
+            dropout=0.0,
+            dropout_att=0.0,
+            dropout_emb=0.0,
+            dropout_pre_encoder=0.0,
+        ).eval()
+        enc._cfg = OmegaConf.create(
+            {
+                'd_model': 64,
+                'n_heads': 4,
+                'dropout': 0.0,
+                'dropout_att': 0.0,
+                'dropout_emb': 0.0,
+                'pos_emb_max_len': 256,
+                'rope_base': 10000.0,
+                'rotary_fraction': 1.0,
+                'use_bias': False,
+            }
+        )
+        # Pre-condition: rel_pos attention has no biases.
+        for layer in enc.layers:
+            assert layer.self_attn.linear_q.bias is None
+
+        enc.change_attention_model('rope')
+
+        # Post-condition: still no biases — use_bias preserved through the swap.
+        for layer in enc.layers:
+            assert layer.self_attn.linear_q.bias is None
+            assert layer.self_attn.linear_k.bias is None
+            assert layer.self_attn.linear_v.bias is None
+            assert layer.self_attn.linear_out.bias is None
+
+    @pytest.mark.unit
+    def test_change_attention_model_preserves_cfg_on_partial_update(self):
+        # Regression: ASRModuleMixin.change_attention_model used to write the *raw*
+        # kwargs into self.cfg.encoder, so a partial update like
+        # `change_attention_model(rotary_fraction=0.5)` left
+        # cfg.encoder.self_attention_model = None (corrupting the saved config) and
+        # skipped writing the rope fields entirely.
+        from omegaconf import DictConfig
+
+        from nemo.collections.asr.models import EncDecCTCModel
+
+        encoder_cfg = {
+            '_target_': 'nemo.collections.asr.modules.ConformerEncoder',
+            'feat_in': 64,
+            'n_layers': 2,
+            'd_model': 64,
+            'n_heads': 4,
+            'self_attention_model': 'rope',
+            'subsampling_factor': 4,
+            'subsampling_conv_channels': 32,
+            'pos_emb_max_len': 256,
+            'rope_base': 10000.0,
+            'rotary_fraction': 1.0,
+            'dropout': 0.0,
+            'dropout_att': 0.0,
+            'dropout_emb': 0.0,
+            'dropout_pre_encoder': 0.0,
+        }
+        decoder_cfg = {
+            '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+            'feat_in': None,
+            'num_classes': 28,
+            'vocabulary': list("abcdefghijklmnopqrstuvwxyz '"),
+        }
+        preproc_cfg = {'_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor'}
+        model = EncDecCTCModel(
+            cfg=DictConfig(
+                {
+                    'preprocessor': preproc_cfg,
+                    'encoder': encoder_cfg,
+                    'decoder': decoder_cfg,
+                    'optim': {'name': 'adamw'},
+                }
+            )
+        )
+
+        # Partial update: only rotary_fraction is being changed.
+        model.change_attention_model(rotary_fraction=0.5)
+
+        # cfg.encoder must reflect the resolved values, not the None kwargs.
+        assert model.cfg.encoder.self_attention_model == 'rope'
+        assert model.cfg.encoder.att_context_size is not None
+        assert model.cfg.encoder.rotary_fraction == 0.5
+        assert model.cfg.encoder.rope_base == 10000.0
+        # Live encoder agrees.
+        assert model.encoder.self_attention_model == 'rope'
+        assert model.encoder.pos_enc.d_k_rot == 8  # d_k=16, fraction=0.5

tests/collections/asr/test_asr_samplers.py

@@ -0,0 +1,157 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import tempfile
+
+import numpy as np
+import pytest
+import soundfile as sf
+import torch
+
+from nemo.collections.asr.data import audio_to_text
+from nemo.collections.asr.parts.utils.asr_batching import SemiSortBatchSampler
+from nemo.collections.asr.parts.utils.manifest_utils import write_manifest
+
+
+class TestASRSamplers:
+    labels = [
+        " ",
+        "a",
+        "b",
+        "c",
+        "d",
+        "e",
+        "f",
+        "g",
+        "h",
+        "i",
+        "j",
+        "k",
+        "l",
+        "m",
+        "n",
+        "o",
+        "p",
+        "q",
+        "r",
+        "s",
+        "t",
+        "u",
+        "v",
+        "w",
+        "x",
+        "y",
+        "z",
+        "'",
+    ]
+
+    @pytest.mark.unit
+    def test_ssb_sampler(self):
+        # Generate random signals
+        data_min_duration = 0.1
+        data_max_duration = 16.7
+
+        random_seed = 42
+        sample_rate = 16000
+
+        _rng = np.random.default_rng(seed=random_seed)
+
+        def generate_samples(num_examples: int) -> list:
+            data_duration = np.round(_rng.uniform(low=data_min_duration, high=data_max_duration, size=num_examples), 3)
+            data_duration_samples = np.floor(data_duration * sample_rate).astype(int)
+            samples = []
+            for data_duration_sample in data_duration_samples:
+                samples.append(_rng.uniform(low=-0.5, high=0.5, size=(data_duration_sample)))
+            return samples
+
+        with tempfile.TemporaryDirectory() as test_dir:
+            # Build metadata for manifest
+            metadata = []
+
+            # Test size of dataloader with and without ssb
+            for num_samples in np.concatenate([np.array([1, 2]), _rng.integers(3, 10, 2), _rng.integers(10, 1000, 2)]):
+                samples = generate_samples(num_samples)
+
+                for n, sample in enumerate(samples):
+                    meta = dict()
+                    signal_filename = f'{n:04d}.wav'
+                    # write audio files
+                    sf.write(os.path.join(test_dir, signal_filename), sample, sample_rate)
+                    # update metadata
+                    meta['audio_filepath'] = os.path.join(test_dir, signal_filename)
+                    meta['duration'] = len(sample) / sample_rate
+                    meta['text'] = 'non empty'
+                    metadata.append(meta)
+
+                # Save manifest
+                manifest_filepath = os.path.join(test_dir, 'manifest.json')
+                write_manifest(manifest_filepath, metadata)
+
+                # Make dataset
+                dataset = audio_to_text.AudioToCharDataset(
+                    manifest_filepath=manifest_filepath,
+                    labels=self.labels,
+                    sample_rate=sample_rate,
+                    max_duration=data_max_duration,
+                    min_duration=data_min_duration,
+                )
+                durations = [sample.duration for sample in dataset.manifest_processor.collection.data]
+
+                # Compare two dataloader
+                for batch_size in _rng.integers(1, n + 20, 5):
+                    batch_size = int(batch_size)
+                    drop_last = True if _rng.integers(0, 2) else False
+                    sampler = SemiSortBatchSampler(
+                        global_rank=0,
+                        world_size=1,
+                        durations=durations,
+                        batch_size=batch_size,
+                        batch_shuffle=True,
+                        drop_last=drop_last,
+                        randomization_factor=0.1,
+                        seed=random_seed,
+                    )
+                    dataloader_with_ssb = torch.utils.data.DataLoader(
+                        dataset=dataset,
+                        batch_size=None,
+                        sampler=sampler,
+                        batch_sampler=None,
+                        collate_fn=lambda x: audio_to_text._speech_collate_fn(x, pad_id=0),
+                    )
+                    dataloader = torch.utils.data.DataLoader(
+                        dataset=dataset,
+                        batch_size=batch_size,
+                        collate_fn=lambda x: audio_to_text._speech_collate_fn(x, pad_id=0),
+                        drop_last=drop_last,
+                        shuffle=True,
+                    )
+
+                    assert abs(len(dataloader) - len(dataloader_with_ssb)) == 0, (
+                        "Different num of batches with batch! Num of batches with ssb is "
+                        f"{len(dataloader_with_ssb)} and without ssb is {len(dataloader)}!"
+                    )
+
+                    dataloader_with_ssb_exception, dataloader_exception = False, False
+
+                    try:
+                        list(dataloader_with_ssb)
+                    except:
+                        dataloader_with_ssb_exception = True
+
+                    try:
+                        list(dataloader)
+                    except:
+                        dataloader_exception = True
+
+                    assert dataloader_with_ssb_exception == dataloader_exception

tests/collections/asr/test_asr_subsampling.py

@@ -0,0 +1,61 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import pytest
+import torch
+
+from nemo.collections.asr.models import ASRModel
+
+
+class TestASRSubsamplingConvChunking:
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_forward(self):
+        asr_model = ASRModel.from_pretrained("stt_en_fastconformer_ctc_large")
+        asr_model = asr_model.eval()
+        asr_model.preprocessor.featurizer.dither = 0.0
+        asr_model.preprocessor.featurizer.pad_to = 0
+
+        len = 512
+
+        input_signal_batch1 = torch.randn(size=(1, len), device=asr_model.device)
+        length_batch1 = torch.randint(low=321, high=500, size=[1], device=asr_model.device)
+
+        input_signal_batch4 = torch.randn(size=(4, len), device=asr_model.device)
+        length_batch4 = torch.randint(low=321, high=500, size=[4], device=asr_model.device)
+
+        with torch.inference_mode():
+            # regular inference
+            logprobs_batch1_nosplit, _, _ = asr_model.forward(
+                input_signal=input_signal_batch1, input_signal_length=length_batch1
+            )
+            logprobs_batch4_nosplit, _, _ = asr_model.forward(
+                input_signal=input_signal_batch4, input_signal_length=length_batch4
+            )
+
+            # force chunking to 2
+            asr_model.change_subsampling_conv_chunking_factor(subsampling_conv_chunking_factor=2)
+
+            # chunked inference by channels as batch is 1
+            logprobs_batch1_split, _, _ = asr_model.forward(
+                input_signal=input_signal_batch1, input_signal_length=length_batch1
+            )
+            # chunked inference by batch as it is 4 [> 1]
+            logprobs_batch4_split, _, _ = asr_model.forward(
+                input_signal=input_signal_batch4, input_signal_length=length_batch4
+            )
+
+        diff = torch.mean(torch.abs(logprobs_batch1_split - logprobs_batch1_nosplit))
+        assert diff <= 0.2
+        diff = torch.mean(torch.abs(logprobs_batch4_split - logprobs_batch4_nosplit))
+        assert diff <= 0.2

tests/collections/asr/test_boosting_tree.py

@@ -0,0 +1,230 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+from lightning.pytorch import Trainer
+from torch.nn.utils.rnn import pad_sequence
+
+from nemo.collections.asr.models import EncDecCTCModelBPE
+from nemo.collections.asr.parts.context_biasing.boosting_graph_batched import (
+    BoostingTreeModelConfig,
+    GPUBoostingTreeModel,
+)
+from nemo.collections.asr.parts.context_biasing.context_graph_universal import ContextGraph
+
+DEVICES = [torch.device("cpu")]
+
+if torch.cuda.is_available():
+    DEVICES.append(torch.device("cuda"))
+
+
+@pytest.fixture(scope="module")
+def test_context_graph():
+    phrases = ["abc", "abd", "c"]
+    phrases_ids = [[1, 2, 3], [1, 2, 4], [3]]
+    scores = [0.0, 0.0, 0.0]
+    context_graph = ContextGraph(context_score=1.0, depth_scaling=1.0)
+    context_graph.build(token_ids=phrases_ids, phrases=phrases, scores=scores, uniform_weights=False)
+    return context_graph
+
+
+@pytest.fixture(scope="module")
+def test_boosting_tree(test_context_graph):
+    boosting_tree = GPUBoostingTreeModel.from_context_graph(
+        context_graph=test_context_graph,
+        vocab_size=5,
+        unk_score=0.0,
+        final_eos_score=0.0,
+        use_triton=True,
+        uniform_weights=False,
+    )
+    return boosting_tree
+
+
+@pytest.fixture(scope="module")
+def conformer_ctc_bpe_model():
+    model = EncDecCTCModelBPE.from_pretrained(model_name="stt_en_conformer_ctc_small")
+    model.set_trainer(Trainer(devices=1, accelerator="cpu"))
+    model = model.eval()
+    return model
+
+
+class TestGPUBoostingTreeModel:
+    @pytest.mark.unit
+    def test_building_context_graph(self, test_context_graph):
+        """Test initial python-based context graph"""
+        context_graph = test_context_graph
+        assert context_graph.num_nodes == 5
+        # end nodes
+        assert context_graph.root.next[1].next[2].next[3].is_end
+        assert context_graph.root.next[1].next[2].next[4].is_end
+        assert context_graph.root.next[3].is_end
+        # words in the end nodes
+        assert context_graph.root.next[1].next[2].next[3].phrase == "abc"
+        assert context_graph.root.next[1].next[2].next[4].phrase == "abd"
+        assert context_graph.root.next[3].phrase == "c"
+        # fail links
+        assert context_graph.root.next[1].next[2].next[3].fail.token == 3
+        assert context_graph.root.next[1].next[2].next[4].fail.token == -1  # root
+        assert context_graph.root.next[3].fail.token == -1  # root
+        # node scores
+        assert round(context_graph.root.next[1].next[2].next[3].node_score, 2) == 4.79
+        assert round(context_graph.root.next[1].next[2].next[4].node_score, 2) == 4.79
+        assert round(context_graph.root.next[3].node_score, 2) == 1.0
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("batch_size", [1, 3, 8])
+    def test_advance_method(self, test_boosting_tree, device, batch_size):
+        """Test advance method with different batch sizes"""
+        test_boosting_tree.to(device)
+        # Test with initial states
+        init_states = test_boosting_tree.get_init_states(batch_size=batch_size, bos=True)
+        scores, next_states = test_boosting_tree.advance(init_states)
+
+        assert scores.shape == (batch_size, 5)  # vocab_size=5
+        assert next_states.shape == (batch_size, 5)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_get_final_method(self, test_boosting_tree, device):
+        """Test get_final method for EOS scoring"""
+        test_boosting_tree.to(device)
+        # Test with various states
+        states = torch.tensor([0, 1, 2], dtype=torch.long, device=device)
+        final_scores = test_boosting_tree.get_final(states)
+
+        assert final_scores.shape == (3,)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_boosting_tree_inference(self, test_boosting_tree, device):
+        """Test boosting tree inference with predefined sentences"""
+        test_boosting_tree.to(device)
+
+        sentences_ids = [[1, 2, 3, 2, 1], [2, 2, 1, 2, 4], [3, 1, 2, 1], []]  # ['abcba', 'bbabd', 'caba', '']
+        boosting_scores = test_boosting_tree(
+            labels=pad_sequence([torch.LongTensor(sentence) for sentence in sentences_ids], batch_first=True).to(
+                device
+            ),
+            labels_lengths=torch.LongTensor([len(sentence) for sentence in sentences_ids]).to(device),
+            bos=False,
+            eos=False,
+        )
+        correct_answer = torch.tensor(
+            [
+                [1.0000, 1.6931, 2.0986, 0.0000, 1.0000],
+                [0.0000, 0.0000, 1.0000, 1.6931, 2.0986],
+                [1.0000, 1.0000, 1.6931, -1.6931, 0.0000],
+                [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
+            ],
+            device=device,
+        )
+        assert torch.allclose(boosting_scores, correct_answer, atol=1e-4)
+
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_triton_vs_pytorch_consistency(self, test_context_graph):
+        """Compare Triton vs PyTorch implementations"""
+        device = torch.device("cuda")
+
+        # Create two identical models with different implementations
+        boosting_tree_triton = GPUBoostingTreeModel.from_context_graph(
+            context_graph=test_context_graph, vocab_size=5, use_triton=True
+        ).to(device)
+
+        boosting_tree_pytorch = GPUBoostingTreeModel.from_context_graph(
+            context_graph=test_context_graph, vocab_size=5, use_triton=False
+        ).to(device)
+
+        # Test with same input
+        sentences_ids = [[1, 2, 3, 2, 1], [2, 2, 1, 2, 4]]
+        labels = pad_sequence([torch.LongTensor(s) for s in sentences_ids], batch_first=True).to(device)
+        lengths = torch.LongTensor([len(s) for s in sentences_ids]).to(device)
+
+        scores_triton = boosting_tree_triton(labels=labels, labels_lengths=lengths, bos=False, eos=False)
+        scores_pytorch = boosting_tree_pytorch(labels=labels, labels_lengths=lengths, bos=False, eos=False)
+
+        assert torch.allclose(scores_triton, scores_pytorch, atol=1e-5)
+
+    @pytest.mark.unit
+    def test_eos_handling(self, test_context_graph):
+        """Test EOS token handling (important for AED models)"""
+        boosting_tree = GPUBoostingTreeModel.from_context_graph(
+            context_graph=test_context_graph, vocab_size=5, unk_score=0.0, final_eos_score=1.0
+        )
+
+        # Test advance with EOS
+        init_states = torch.tensor([1, 2], dtype=torch.long)
+        scores, next_states = boosting_tree.advance(init_states, eos_id=0)
+
+        # state 2 in the 1st batch should have final_eos_score value
+        assert (
+            round(scores[0, 0].item(), 2) == 1.69
+        )  # (1.69+0): 1.69 as max score for state 1 and 0 because it is not final state
+        assert scores[1, 0] == 2.0  # (1+1): 1 as max score for state 2 and 1 because it is final state
+
+    @pytest.mark.unit
+    # I need to test that the boosting tree model is built correctly from the config using model_path, key_phrases_file, key_phrases_list
+    def test_boosting_tree_model_from_config(self, conformer_ctc_bpe_model, tmp_path):
+        """Test that the boosting tree model is built correctly from the config using model_path, key_phrases_file, key_phrases_list"""
+
+        # 1. build boosting tree model from model path
+        boosting_tree_cfg = BoostingTreeModelConfig()
+        phrases = ["abc", "abd", "c"]
+        phrases_ids = [conformer_ctc_bpe_model.tokenizer.text_to_ids(phrase) for phrase in phrases]
+        scores = [0.0, 0.0, 0.0]
+        context_graph = ContextGraph(
+            context_score=boosting_tree_cfg.context_score, depth_scaling=boosting_tree_cfg.depth_scaling
+        )
+        context_graph.build(
+            token_ids=phrases_ids, phrases=phrases, scores=scores, uniform_weights=boosting_tree_cfg.uniform_weights
+        )
+        test_boosting_tree = GPUBoostingTreeModel.from_context_graph(
+            context_graph=context_graph,
+            vocab_size=conformer_ctc_bpe_model.tokenizer.vocab_size,
+            unk_score=boosting_tree_cfg.unk_score,
+            final_eos_score=boosting_tree_cfg.final_eos_score,
+            use_triton=boosting_tree_cfg.use_triton,
+            uniform_weights=boosting_tree_cfg.uniform_weights,
+        )
+
+        test_boosting_tree.save_to(tmp_path / "test_boosting_tree.nemo")
+        boosting_tree_cfg = BoostingTreeModelConfig(model_path=tmp_path / "test_boosting_tree.nemo")
+        boosting_tree_from_model_path = GPUBoostingTreeModel.from_config(
+            boosting_tree_cfg, tokenizer=conformer_ctc_bpe_model.tokenizer
+        )
+
+        # 2. build boosting tree model from key phrases file
+        with open(tmp_path / "test_boosting_tree.txt", "w") as f:
+            f.write("abc\nabd\nc")
+        boosting_tree_cfg = BoostingTreeModelConfig(key_phrases_file=tmp_path / "test_boosting_tree.txt")
+        boosting_tree_from_key_phrases_file = GPUBoostingTreeModel.from_config(
+            boosting_tree_cfg, tokenizer=conformer_ctc_bpe_model.tokenizer
+        )
+
+        # 3. build boosting tree model from key phrases list
+        boosting_tree_cfg = BoostingTreeModelConfig(key_phrases_list=["abc", "abd", "c"])
+        boosting_tree_from_key_phrases_list = GPUBoostingTreeModel.from_config(
+            boosting_tree_cfg, tokenizer=conformer_ctc_bpe_model.tokenizer
+        )
+
+        # check that the boosting tree models are the same
+        assert torch.allclose(
+            boosting_tree_from_model_path.arcs_weights, boosting_tree_from_key_phrases_file.arcs_weights
+        )
+        assert torch.allclose(
+            boosting_tree_from_model_path.arcs_weights, boosting_tree_from_key_phrases_list.arcs_weights
+        )

tests/collections/asr/test_conformer_encoder.py

@@ -0,0 +1,224 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import numpy as np
+import pytest
+import torch
+
+from nemo.collections.asr.modules.conformer_encoder import ConformerEncoder
+
+
+class TestStochasticDepth:
+    """Testing stochastic depth functionality."""
+
+    def test_stochastic_depth_model_creation(self):
+        """Testing basic model creation and the drop probs are correctly assigned."""
+        n_layers = 4
+        model = ConformerEncoder(feat_in=10, n_layers=n_layers, d_model=4, feat_out=8)
+
+        # checking that by default SD is disabled
+        assert model.layer_drop_probs == [0.0] * n_layers
+
+        # linear mode
+        for drop_prob in [0.3, 0.5, 0.9]:
+            for start_layer in [1, 3]:
+                model = ConformerEncoder(
+                    feat_in=10,
+                    n_layers=n_layers,
+                    d_model=4,
+                    feat_out=8,
+                    stochastic_depth_drop_prob=drop_prob,
+                    stochastic_depth_start_layer=start_layer,
+                )
+                L = n_layers - start_layer
+                assert model.layer_drop_probs == [0.0] * start_layer + [drop_prob * l / L for l in range(1, L + 1)]
+
+        # uniform mode
+        for drop_prob in [0.3, 0.5, 0.9]:
+            model = ConformerEncoder(
+                feat_in=10,
+                n_layers=n_layers,
+                d_model=4,
+                feat_out=8,
+                stochastic_depth_drop_prob=drop_prob,
+                stochastic_depth_mode="uniform",
+                stochastic_depth_start_layer=start_layer,
+            )
+            L = n_layers - start_layer
+            assert model.layer_drop_probs == [0.0] * start_layer + [drop_prob] * L
+
+        # checking for errors
+        for drop_prob in [-1.0, 1.0]:
+            with pytest.raises(ValueError, match="stochastic_depth_drop_prob has to be in"):
+                ConformerEncoder(
+                    feat_in=10,
+                    n_layers=n_layers,
+                    d_model=4,
+                    feat_out=8,
+                    stochastic_depth_drop_prob=drop_prob,
+                    stochastic_depth_mode="uniform",
+                )
+
+        with pytest.raises(ValueError, match="stochastic_depth_mode has to be one of"):
+            ConformerEncoder(feat_in=10, n_layers=n_layers, d_model=4, feat_out=8, stochastic_depth_mode="weird")
+
+        for start_layer in [-1, 0, 5]:
+            with pytest.raises(ValueError, match="stochastic_depth_start_layer has to be in"):
+                ConformerEncoder(
+                    feat_in=10,
+                    n_layers=n_layers,
+                    d_model=4,
+                    feat_out=8,
+                    stochastic_depth_start_layer=start_layer,
+                )
+
+    @pytest.mark.pleasefixme
+    def test_stochastic_depth_forward(self):
+        """Testing that forward works and we get randomness during training, but not during eval."""
+        random_input = torch.rand((1, 2, 2))
+        random_length = torch.tensor([2], dtype=torch.int64)
+
+        model = ConformerEncoder(
+            feat_in=2,
+            n_layers=3,
+            d_model=4,
+            feat_out=4,
+            stochastic_depth_drop_prob=0.8,
+            dropout=0.0,
+            dropout_pre_encoder=0.0,
+            dropout_emb=0.0,
+            conv_norm_type="layer_norm",
+            conv_kernel_size=3,
+        )
+        model.train()
+        outputs = [None] * 5
+        for i in range(5):
+            outputs[i] = model(audio_signal=random_input, length=random_length)[0]
+        # checking that not all outputs are the same
+        num_diff = 0
+        for i in range(1, 5):
+            if not torch.allclose(outputs[i], outputs[0]):
+                num_diff += 1
+        assert num_diff > 0
+
+        model.eval()
+        outputs = [None] * 5
+        for i in range(5):
+            outputs[i] = model(audio_signal=random_input, length=random_length)[0]
+        # checking that not all outputs are the same
+        num_diff = 0
+        for i in range(1, 5):
+            if not torch.allclose(outputs[i], outputs[0]):
+                num_diff += 1
+        assert num_diff == 0
+
+
+class TestBypassPreEncode:
+    """Testing bypass pre-encode functionality."""
+
+    def test_bypass_pre_encode_forward(self):
+        """Testing that forward works with "bypass pre-encode" mode."""
+        # For pre-encoded embeddings, the shape is (batch_size, n_frames, emb_dim)
+        batch_size = 2
+        n_frames, emb_dim, feat_out = 17, 16, 8
+        random_input = torch.rand((batch_size, n_frames, emb_dim))
+        random_length = torch.tensor([n_frames], dtype=torch.int64)
+
+        model = ConformerEncoder(
+            feat_in=10,
+            n_layers=3,
+            d_model=emb_dim,
+            feat_out=feat_out,
+            stochastic_depth_drop_prob=0.0,
+            dropout=0.0,
+            dropout_pre_encoder=0.0,
+            dropout_emb=0.0,
+            conv_norm_type="layer_norm",
+            conv_kernel_size=3,
+        )
+        model.train()
+        fwd_outputs = model(audio_signal=random_input, length=random_length, bypass_pre_encode=True)[0]
+        assert fwd_outputs.shape == (batch_size, feat_out, n_frames)
+
+        model.eval()
+        fwd_outputs = model(audio_signal=random_input, length=random_length, bypass_pre_encode=True)[0]
+        assert fwd_outputs.shape == (batch_size, feat_out, n_frames)
+
+    def test_error_shape_invalid_bypass_pre_encode_forward(self):
+        """
+        Testing that error messages are correctly triggered regarding "bypass pre-encode" mode.
+        Both correct samples and wrongs samples are tested.
+
+        (1) bypass_pre_encode = False (default):
+            `audio_signal` must be a tensor containing audio features.
+            Shape: (batch, self._feat_in, n_frames)
+        (2) bypass_pre_encode = True:
+            `audio_signal` must be a tensor containing pre-encoded embeddings.
+            Shape: (batch, n_frame, self.d_model)
+        """
+        batch_size = 2
+        n_frames, emb_dim, feat_in, feat_out = 17, 16, 10, 8
+
+        pre_encode_input = torch.rand((batch_size, n_frames, emb_dim))
+        feat_input = torch.rand((batch_size, feat_in, n_frames))
+        input_length = torch.tensor([n_frames], dtype=torch.int64)
+
+        model = ConformerEncoder(
+            feat_in=feat_in,
+            n_layers=3,
+            d_model=emb_dim,
+            feat_out=feat_out,
+            stochastic_depth_drop_prob=0.0,
+            dropout=0.0,
+            dropout_pre_encoder=0.0,
+            dropout_emb=0.0,
+            conv_norm_type="layer_norm",
+            conv_kernel_size=3,
+        )
+        sub_sampled_n_frames = np.ceil(n_frames / model.subsampling_factor)
+
+        # Test with bypass_pre_encode = True, should be pre_encode_input but given feat_input.
+        model.train()
+        with pytest.raises(ValueError):
+            model(audio_signal=feat_input, length=input_length, bypass_pre_encode=True)
+
+        model.eval()
+        with pytest.raises(ValueError):
+            model(audio_signal=feat_input, length=input_length, bypass_pre_encode=True)
+
+        # Test with bypass_pre_encode = True, given the correct input pre_encode_input.
+        model.train()
+        fwd_outputs = model(audio_signal=pre_encode_input, length=input_length, bypass_pre_encode=True)[0]
+        assert fwd_outputs.shape == (batch_size, feat_out, n_frames)
+
+        model.eval()
+        fwd_outputs = model(audio_signal=pre_encode_input, length=input_length, bypass_pre_encode=True)[0]
+        assert fwd_outputs.shape == (batch_size, feat_out, n_frames)
+
+        # Test with bypass_pre_encode = False, should be feat_input but given pre_encode_input.
+        model.train()
+        with pytest.raises(ValueError):
+            model(audio_signal=pre_encode_input, length=input_length, bypass_pre_encode=False)
+
+        model.eval()
+        with pytest.raises(ValueError):
+            model(audio_signal=pre_encode_input, length=input_length, bypass_pre_encode=False)
+
+        # Test with bypass_pre_encode = False, given the correct input feat_input.
+        model.train()
+        fwd_outputs = model(audio_signal=feat_input, length=input_length, bypass_pre_encode=False)[0]
+        assert fwd_outputs.shape == (batch_size, feat_out, sub_sampled_n_frames)
+
+        model.eval()
+        fwd_outputs = model(audio_signal=feat_input, length=input_length, bypass_pre_encode=False)[0]
+        assert fwd_outputs.shape == (batch_size, feat_out, sub_sampled_n_frames)

tests/collections/asr/test_custom_tokenizer.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import Mock
+
+import pytest
+import sentencepiece as spm
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.parts.mixins import ASRBPEMixin
+from nemo.collections.common.tokenizers.canary_tokenizer import DEFAULT_TOKENS, CanaryTokenizer
+from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer, create_spt_model
+from nemo.core import Serialization
+
+
+@pytest.fixture(scope="session")
+def special_tokenizer_path(tmp_path_factory) -> str:
+    tokens = ["asr", "ast", "en", "de", "fr", "es"]
+    tmpdir = tmp_path_factory.mktemp("spl_tokens")
+    CanaryTokenizer.build_special_tokenizer(tokens, tmpdir)
+    return str(tmpdir)
+
+
+@pytest.fixture(scope="session")
+def lang_tokenizer_path(tmp_path_factory) -> str:
+    tmpdir = tmp_path_factory.mktemp("klingon_tokens")
+    text_path = tmpdir / "text.txt"
+    text_path.write_text("a\nb\nc\nd\n")
+    create_spt_model(text_path, vocab_size=8, sample_size=-1, do_lower_case=False, output_dir=str(tmpdir))
+    return str(tmpdir)
+
+
+def test_canary_tokenizer_build_special_tokenizer(tmp_path):
+    tokens = ["asr", "ast", "en", "de", "fr", "es"]
+    tokenizer = CanaryTokenizer.build_special_tokenizer(tokens, tmp_path)
+    expected_tokens = DEFAULT_TOKENS + [f"<|{t}|>" for t in tokens] + ["▁", "<unk>"]
+    tokens = []
+    for i in range(tokenizer.tokenizer.vocab_size()):
+        tokens.append(tokenizer.tokenizer.IdToPiece(i))
+    expected_tokens.sort(), tokens.sort()
+    print(expected_tokens, tokens)
+    assert expected_tokens == tokens
+
+
+def test_canary_tokenizer_init_from_cfg(special_tokenizer_path, lang_tokenizer_path):
+    class DummyModel(ASRBPEMixin, Serialization):
+        pass
+
+    model = DummyModel()
+    model.register_artifact = Mock(side_effect=lambda self, x: x)
+    config = OmegaConf.create(
+        {
+            "type": "agg",
+            "dir": None,
+            "langs": {
+                "spl_tokens": {"dir": special_tokenizer_path, "type": "bpe"},
+                "en": {"dir": lang_tokenizer_path, "type": "bpe"},
+            },
+            "custom_tokenizer": {
+                "_target_": "nemo.collections.common.tokenizers.canary_tokenizer.CanaryTokenizer",
+            },
+        }
+    )
+    model._setup_aggregate_tokenizer(config)
+    tokenizer = model.tokenizer
+
+    assert isinstance(tokenizer, CanaryTokenizer)
+    assert len(tokenizer.tokenizers_dict) == 2
+    assert set(tokenizer.tokenizers_dict.keys()) == {"spl_tokens", "en"}
+
+    assert isinstance(tokenizer.tokenizers_dict["spl_tokens"], SentencePieceTokenizer)
+    assert tokenizer.tokenizers_dict["spl_tokens"].vocab_size == 14
+
+    assert isinstance(tokenizer.tokenizers_dict["en"], SentencePieceTokenizer)
+    assert tokenizer.tokenizers_dict["en"].vocab_size == 6
+
+    assert tokenizer.text_to_ids("<|startoftranscript|><|en|><|asr|><|en|><|pnc|>", lang_id="spl_tokens") == [
+        4,
+        9,
+        7,
+        9,
+        5,
+    ]
+    assert tokenizer.text_to_ids("a", lang_id="en") == [14 + 1, 14 + 2]

tests/collections/asr/test_jasper_block.py

@@ -0,0 +1,398 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.collections.asr.parts.submodules import jasper
+
+
+class TestJasperBlock:
+    @staticmethod
+    def jasper_base_config(**kwargs):
+        base = dict(
+            inplanes=16,
+            planes=8,
+            kernel_size=[11],
+            repeat=1,
+            stride=[1],
+            dilation=[1],
+            activation="relu",
+            conv_mask=True,
+            separable=False,
+            se=False,
+        )
+        base.update(kwargs)
+        return base
+
+    def check_module_exists(self, module, cls):
+        global _MODULE_EXISTS
+        _MODULE_EXISTS = 0
+
+        def _traverse(m):
+            if isinstance(m, cls):
+                global _MODULE_EXISTS
+                _MODULE_EXISTS += 1
+
+        module.apply(_traverse)
+        assert _MODULE_EXISTS > 0
+
+    @pytest.mark.unit
+    def test_basic_block(self):
+        config = self.jasper_base_config(residual=False)
+        act = jasper.jasper_activations.get(config.pop('activation'))()
+
+        block = jasper.JasperBlock(**config, activation=act)
+
+        x = torch.randn(1, 16, 131)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        assert isinstance(block, jasper.JasperBlock)
+        assert y[0].shape == torch.Size([1, config['planes'], 131])
+        assert ylen[0] == 131
+
+    @pytest.mark.unit
+    def test_residual_block(self):
+        config = self.jasper_base_config(residual=True)
+        act = jasper.jasper_activations.get(config.pop('activation'))()
+
+        block = jasper.JasperBlock(**config, activation=act)
+
+        x = torch.randn(1, 16, 131)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        assert isinstance(block, jasper.JasperBlock)
+        assert y[0].shape == torch.Size([1, config['planes'], 131])
+        assert ylen[0] == 131
+
+    @pytest.mark.unit
+    def test_basic_block_repeat(self):
+        config = self.jasper_base_config(residual=False, repeat=3)
+        act = jasper.jasper_activations.get(config.pop('activation'))()
+
+        block = jasper.JasperBlock(**config, activation=act)
+
+        x = torch.randn(1, 16, 131)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        assert isinstance(block, jasper.JasperBlock)
+        assert y[0].shape == torch.Size([1, config['planes'], 131])
+        assert ylen[0] == 131
+        assert len(block.mconv) == 3 * 3 + 1  # (3 repeats x {1 conv + 1 norm + 1 dropout} + final conv)
+
+    @pytest.mark.unit
+    def test_basic_block_repeat_stride(self):
+        config = self.jasper_base_config(residual=False, repeat=3, stride=[2])
+        act = jasper.jasper_activations.get(config.pop('activation'))()
+
+        block = jasper.JasperBlock(**config, activation=act)
+
+        x = torch.randn(1, 16, 131)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        assert isinstance(block, jasper.JasperBlock)
+        assert y[0].shape == torch.Size([1, config['planes'], 17])  # 131 // (stride ^ repeats)
+        assert ylen[0] == 17  # 131 // (stride ^ repeats)
+        assert len(block.mconv) == 3 * 3 + 1  # (3 repeats x {1 conv + 1 norm + 1 dropout} + final conv)
+
+    @pytest.mark.unit
+    def test_basic_block_repeat_stride_last(self):
+        config = self.jasper_base_config(residual=False, repeat=3, stride=[2], stride_last=True)
+        act = jasper.jasper_activations.get(config.pop('activation'))()
+
+        block = jasper.JasperBlock(**config, activation=act)
+
+        x = torch.randn(1, 16, 131)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        assert isinstance(block, jasper.JasperBlock)
+        assert y[0].shape == torch.Size([1, config['planes'], 66])  # 131 // stride
+        assert ylen[0] == 66  # 131 // stride
+        assert len(block.mconv) == 3 * 3 + 1  # (3 repeats x {1 conv + 1 norm + 1 dropout} + final conv)
+
+    @pytest.mark.unit
+    def test_basic_block_repeat_separable(self):
+        config = self.jasper_base_config(residual=False, repeat=3, separable=True)
+        act = jasper.jasper_activations.get(config.pop('activation'))()
+
+        block = jasper.JasperBlock(**config, activation=act)
+
+        x = torch.randn(1, 16, 131)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        assert isinstance(block, jasper.JasperBlock)
+        assert y[0].shape == torch.Size([1, config['planes'], 131])
+        assert ylen[0] == 131
+        assert len(block.mconv) == 3 * 4 + 1  # (3 repeats x {1 dconv + 1 pconv + 1 norm + 1 dropout} + final conv)
+
+    @pytest.mark.unit
+    def test_basic_block_stride(self):
+        config = self.jasper_base_config(stride=[2], residual=False)
+        act = jasper.jasper_activations.get(config.pop('activation'))()
+
+        print(config)
+        block = jasper.JasperBlock(**config, activation=act)
+
+        x = torch.randn(1, 16, 131)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        assert isinstance(block, jasper.JasperBlock)
+        assert y[0].shape == torch.Size([1, config['planes'], 66])
+        assert ylen[0] == 66
+
+    @pytest.mark.unit
+    def test_residual_block_stride(self):
+        config = self.jasper_base_config(stride=[2], residual=True, residual_mode='stride_add')
+        act = jasper.jasper_activations.get(config.pop('activation'))()
+
+        print(config)
+        block = jasper.JasperBlock(**config, activation=act)
+
+        x = torch.randn(1, 16, 131)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        assert isinstance(block, jasper.JasperBlock)
+        assert y[0].shape == torch.Size([1, config['planes'], 66])
+        assert ylen[0] == 66
+
+    @pytest.mark.unit
+    def test_residual_block_activations(self):
+        for activation in jasper.jasper_activations.keys():
+            config = self.jasper_base_config(activation=activation)
+            act = jasper.jasper_activations.get(config.pop('activation'))()
+
+            block = jasper.JasperBlock(**config, activation=act)
+
+            x = torch.randn(1, 16, 131)
+            xlen = torch.tensor([131])
+            y, ylen = block(([x], xlen))
+
+            self.check_module_exists(block, act.__class__)
+            assert isinstance(block, jasper.JasperBlock)
+            assert y[0].shape == torch.Size([1, config['planes'], 131])
+            assert ylen[0] == 131
+
+    @pytest.mark.unit
+    def test_residual_block_normalizations(self):
+        NORMALIZATIONS = ["batch", "layer", "group"]
+        for normalization in NORMALIZATIONS:
+            config = self.jasper_base_config(normalization=normalization)
+            act = jasper.jasper_activations.get(config.pop('activation'))()
+
+            block = jasper.JasperBlock(**config, activation=act)
+
+            x = torch.randn(1, 16, 131)
+            xlen = torch.tensor([131])
+            y, ylen = block(([x], xlen))
+
+            assert isinstance(block, jasper.JasperBlock)
+            assert y[0].shape == torch.Size([1, config['planes'], 131])
+            assert ylen[0] == 131
+
+    @pytest.mark.unit
+    def test_residual_block_se(self):
+        config = self.jasper_base_config(se=True, se_reduction_ratio=8)
+        act = jasper.jasper_activations.get(config.pop('activation'))()
+
+        block = jasper.JasperBlock(**config, activation=act)
+
+        x = torch.randn(1, 16, 131)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        self.check_module_exists(block, jasper.SqueezeExcite)
+        assert isinstance(block, jasper.JasperBlock)
+        assert y[0].shape == torch.Size([1, config['planes'], 131])
+        assert ylen[0] == 131
+
+    @pytest.mark.unit
+    def test_residual_block_asymmetric_pad_future_contexts(self):
+        # test future contexts at various values
+        # 0 = no future context
+        # 2 = limited future context
+        # 5 = symmetric context
+        # 8 = excess future context (more future context than present or past context)
+        future_contexts = [0, 2, 5, 8]
+        for future_context in future_contexts:
+            print(future_context)
+            config = self.jasper_base_config(future_context=future_context)
+            act = jasper.jasper_activations.get(config.pop('activation'))()
+
+            block = jasper.JasperBlock(**config, activation=act)
+
+            x = torch.randn(1, 16, 131)
+            xlen = torch.tensor([131])
+            y, ylen = block(([x], xlen))
+
+            self.check_module_exists(block, torch.nn.ConstantPad1d)
+            self.check_module_exists(block, jasper.MaskedConv1d)
+
+            assert isinstance(block, jasper.JasperBlock)
+            assert y[0].shape == torch.Size([1, config['planes'], 131])
+            assert ylen[0] == 131
+            assert block.mconv[0].pad_layer is not None
+            assert block.mconv[0]._padding == (config['kernel_size'][0] - 1 - future_context, future_context)
+
+    @pytest.mark.unit
+    def test_residual_block_asymmetric_pad_future_context_fallback(self):
+        # test future contexts at various values
+        # 15 = K < FC; fall back to symmetric context
+        future_context = 15
+        print(future_context)
+        config = self.jasper_base_config(future_context=future_context)
+        act = jasper.jasper_activations.get(config.pop('activation'))()
+
+        block = jasper.JasperBlock(**config, activation=act)
+
+        x = torch.randn(1, 16, 131)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        self.check_module_exists(block, jasper.MaskedConv1d)
+
+        assert isinstance(block, jasper.JasperBlock)
+        assert y[0].shape == torch.Size([1, config['planes'], 131])
+        assert ylen[0] == 131
+        assert block.mconv[0].pad_layer is None
+        assert block.mconv[0]._padding == config['kernel_size'][0] // 2
+
+    @pytest.mark.unit
+    def test_padding_size_conv1d(self):
+        input_channels = 1
+        output_channels = 1
+        kernel_sizes = [3, 7, 11]
+        dilation_sizes = [2, 3, 4]
+        stride = 1
+        inp = torch.rand(2, 1, 40)
+
+        for kernel_size in kernel_sizes:
+            for dilation_size in dilation_sizes:
+                padding = jasper.get_same_padding(kernel_size, stride, dilation_size)
+
+                conv = torch.nn.Conv1d(
+                    input_channels, output_channels, kernel_size=kernel_size, dilation=dilation_size, padding=padding
+                )
+
+                out = conv(inp)
+                assert out.shape == inp.shape
+
+
+class TestParallelBlock:
+    @staticmethod
+    def contrust_jasper_block(**config_kwargs):
+        config = TestJasperBlock.jasper_base_config(**config_kwargs)
+        act = jasper.jasper_activations.get(config.pop('activation'))()
+        block = jasper.JasperBlock(**config, activation=act)
+        return block
+
+    @pytest.mark.unit
+    def test_blocks_with_same_input_output_channels_sum_residual(self):
+        blocks = []
+        in_planes = 8
+        out_planes = 8
+        for _ in range(2):
+            blocks.append(self.contrust_jasper_block(inplanes=in_planes, planes=out_planes))
+
+        block = jasper.ParallelBlock(blocks, residual_mode='sum')
+        x = torch.randn(1, in_planes, 140)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        assert y[0].shape == torch.Size([1, out_planes, 140])
+        assert ylen[0] == 131
+
+    @pytest.mark.unit
+    def test_blocks_with_different_input_output_channels_sum_residual(self):
+        blocks = []
+        in_planes = 8
+        out_planes = 16
+        for _ in range(2):
+            blocks.append(self.contrust_jasper_block(inplanes=in_planes, planes=out_planes))
+
+        block = jasper.ParallelBlock(blocks, residual_mode='sum')
+        x = torch.randn(1, in_planes, 140)
+        xlen = torch.tensor([131])
+
+        with pytest.raises(RuntimeError):
+            block(([x], xlen))
+
+    @pytest.mark.unit
+    def test_blocks_with_same_input_output_channels_conv_residual(self):
+        blocks = []
+        in_planes = 8
+        out_planes = 8
+        for _ in range(2):
+            blocks.append(self.contrust_jasper_block(inplanes=in_planes, planes=out_planes))
+
+        block = jasper.ParallelBlock(blocks, residual_mode='conv', in_filters=in_planes, out_filters=out_planes)
+        x = torch.randn(1, in_planes, 140)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        assert y[0].shape == torch.Size([1, out_planes, 140])
+        assert ylen[0] == 131
+
+    @pytest.mark.unit
+    def test_blocks_with_different_input_output_channels_conv_residual(self):
+        blocks = []
+        in_planes = 8
+        out_planes = 16
+        for _ in range(2):
+            blocks.append(self.contrust_jasper_block(inplanes=in_planes, planes=out_planes))
+
+        block = jasper.ParallelBlock(blocks, residual_mode='conv', in_filters=in_planes, out_filters=out_planes)
+        x = torch.randn(1, in_planes, 140)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        assert y[0].shape == torch.Size([1, out_planes, 140])
+        assert ylen[0] == 131
+
+    @pytest.mark.unit
+    def test_single_block(self):
+        in_planes = 8
+        out_planes = 16
+        blocks = [self.contrust_jasper_block(inplanes=in_planes, planes=out_planes)]
+
+        block = jasper.ParallelBlock(blocks)
+        x = torch.randn(1, in_planes, 140)
+        xlen = torch.tensor([131])
+        y, ylen = block(([x], xlen))
+
+        assert y[0].shape == torch.Size([1, out_planes, 140])
+        assert ylen[0] == 131
+
+    @pytest.mark.unit
+    def test_tower_dropout(self):
+        blocks = []
+        in_planes = 8
+        out_planes = 8
+        for _ in range(2):
+            blocks.append(self.contrust_jasper_block(inplanes=in_planes, planes=out_planes))
+
+        block = jasper.ParallelBlock(blocks, aggregation_mode='dropout', block_dropout_prob=1.0)
+        x = torch.randn(1, in_planes, 140)
+        xlen = torch.tensor([131])
+        y, _ = block(([x], xlen))
+
+        # Tower dropout is 1.0, meaning that all towers have to be dropped, so only residual remains.
+        torch.testing.assert_close(y[0], x)

tests/collections/asr/test_label_datasets.py

@@ -0,0 +1,162 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import os
+import tempfile
+
+import numpy as np
+import pytest
+import soundfile as sf
+import torch
+
+from nemo.collections.asr.data.audio_to_label import AudioToMultiLabelDataset, TarredAudioToClassificationLabelDataset
+from nemo.collections.asr.data.feature_to_label import FeatureToLabelDataset, FeatureToSeqSpeakerLabelDataset
+from nemo.collections.asr.parts.preprocessing.feature_loader import ExternalFeatureLoader
+from nemo.collections.asr.parts.preprocessing.features import WaveformFeaturizer
+
+
+class TestASRDatasets:
+    labels = ["fash", "fbbh", "fclc"]
+    unique_labels_in_seq = ['0', '1', '2', '3', "zero", "one", "two", "three"]
+
+    @pytest.mark.unit
+    def test_tarred_dataset(self, test_data_dir):
+        manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/tarred_audio_manifest.json'))
+
+        # Test braceexpand loading
+        tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar'))
+        featurizer = WaveformFeaturizer(sample_rate=16000, int_values=False, augmentor=None)
+        ds_braceexpand = TarredAudioToClassificationLabelDataset(
+            audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, featurizer=featurizer
+        )
+
+        assert len(ds_braceexpand) == 32
+        count = 0
+        for _ in ds_braceexpand:
+            count += 1
+        assert count == 32
+
+        # Test loading via list
+        tarpath = [os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar')) for i in range(2)]
+        ds_list_load = TarredAudioToClassificationLabelDataset(
+            audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, featurizer=featurizer
+        )
+        count = 0
+        for _ in ds_list_load:
+            count += 1
+        assert count == 32
+
+    @pytest.mark.unit
+    def test_tarred_dataset_duplicate_name(self, test_data_dir):
+        manifest_path = os.path.abspath(
+            os.path.join(test_data_dir, 'asr/tarred_an4/tarred_duplicate_audio_manifest.json')
+        )
+
+        # Test braceexpand loading
+        tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar'))
+        featurizer = WaveformFeaturizer(sample_rate=16000, int_values=False, augmentor=None)
+        ds_braceexpand = TarredAudioToClassificationLabelDataset(
+            audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, featurizer=featurizer
+        )
+
+        assert len(ds_braceexpand) == 6
+        count = 0
+        for _ in ds_braceexpand:
+            count += 1
+        assert count == 6
+
+        # Test loading via list
+        tarpath = [os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar')) for i in range(2)]
+        ds_list_load = TarredAudioToClassificationLabelDataset(
+            audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, featurizer=featurizer
+        )
+        count = 0
+        for _ in ds_list_load:
+            count += 1
+        assert count == 6
+
+    @pytest.mark.unit
+    def test_feat_seqlabel_dataset(self, test_data_dir):
+        manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/feat/emb.json'))
+        feature_loader = ExternalFeatureLoader(augmentor=None)
+        ds_braceexpand = FeatureToSeqSpeakerLabelDataset(
+            manifest_filepath=manifest_path, labels=self.unique_labels_in_seq, feature_loader=feature_loader
+        )
+        # fmt: off
+        correct_label = torch.tensor(
+            [0.0, 1.0, 2.0, 2.0, 1.0, 2.0, 1.0, 2.0, 2.0, 1.0, 2.0, 2.0, 3.0, 1.0, 2.0, 2.0, 2.0, 0.0, 2.0, 1.0, 1.0, 2.0, 2.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 2.0, 2.0, 2.0, 1.0, 2.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 2.0, 1.0, 2.0, 1.0,]
+        )
+        # fmt: on
+        correct_label_length = torch.tensor(50)
+
+        assert ds_braceexpand[0][0].shape == (50, 32)
+        assert torch.equal(ds_braceexpand[0][2], correct_label)
+        assert torch.equal(ds_braceexpand[0][3], correct_label_length)
+
+        count = 0
+        for _ in ds_braceexpand:
+            count += 1
+        assert count == 2
+
+    @pytest.mark.unit
+    def test_feat_label_dataset(self):
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manifest_path = os.path.join(tmpdir, 'manifest_input.json')
+            with open(manifest_path, 'w', encoding='utf-8') as fp:
+                for i in range(2):
+                    feat_file = os.path.join(tmpdir, f"feat_{i}.pt")
+                    torch.save(torch.randn(80, 5), feat_file)
+                    entry = {'feature_file': feat_file, 'duration': 100000, 'label': '0'}
+                    fp.write(json.dumps(entry) + '\n')
+
+            dataset = FeatureToLabelDataset(manifest_filepath=manifest_path, labels=self.unique_labels_in_seq)
+
+            correct_label = torch.tensor(self.unique_labels_in_seq.index('0'))
+            correct_label_length = torch.tensor(1)
+
+            assert dataset[0][0].shape == (80, 5)
+            assert torch.equal(dataset[0][2], correct_label)
+            assert torch.equal(dataset[0][3], correct_label_length)
+
+            count = 0
+            for _ in dataset:
+                count += 1
+            assert count == 2
+
+    @pytest.mark.unit
+    def test_audio_multilabel_dataset(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manifest_path = os.path.join(tmpdir, 'manifest_input.json')
+            with open(manifest_path, 'w', encoding='utf-8') as fp:
+                for i in range(2):
+                    audio_file = os.path.join(tmpdir, f"audio_{i}.wav")
+                    data = np.random.normal(0, 1, 16000 * 10)
+                    sf.write(audio_file, data, 16000)
+                    entry = {'audio_filepath': audio_file, 'duration': 10, 'label': '0 1 0 1'}
+                    fp.write(json.dumps(entry) + '\n')
+
+            dataset = AudioToMultiLabelDataset(manifest_filepath=manifest_path, sample_rate=16000, labels=['0', '1'])
+
+            correct_label = torch.tensor([0, 1, 0, 1])
+            correct_label_length = torch.tensor(4)
+
+            assert dataset[0][0].shape == torch.tensor([0.1] * 160000).shape
+            assert torch.equal(dataset[0][2], correct_label)
+            assert torch.equal(dataset[0][3], correct_label_length)
+
+            count = 0
+            for _ in dataset:
+                count += 1
+            assert count == 2

tests/collections/asr/test_ngram_lm.py

@@ -0,0 +1,209 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+from pathlib import Path
+
+import pytest
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from tqdm.auto import tqdm
+
+from nemo.collections.asr.parts.submodules.ngram_lm import KenLMBatchedWrapper, NGramGPULanguageModel
+from nemo.core.utils.optional_libs import KENLM_AVAILABLE, TRITON_AVAILABLE
+
+DEVICES = [torch.device("cpu")]
+
+if torch.cuda.is_available():
+    DEVICES.append('cuda')
+
+
+@pytest.fixture(scope="module")
+def n_gpu_lm(test_data_dir):
+    kenlm_model_path = Path(test_data_dir) / "asr/kenlm_ngram_lm/parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+    vocab_size = 1024
+    return NGramGPULanguageModel.from_arpa(kenlm_model_path, vocab_size=vocab_size, normalize_unk=False)
+
+
+@pytest.fixture(scope="module")
+def kenlm_wrapper(test_data_dir):
+    kenlm_model_path = Path(test_data_dir) / "asr/kenlm_ngram_lm/parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+    vocab_size = 1024
+    return KenLMBatchedWrapper.from_file(lm_path=kenlm_model_path, vocab_size=vocab_size)
+
+
+class TestNGramGPULanguageModel:
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    def test_load(self, test_data_dir):
+        kenlm_model_path = Path(test_data_dir) / "asr/kenlm_ngram_lm/parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+        _ = NGramGPULanguageModel.from_arpa(kenlm_model_path, vocab_size=1024)
+
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not KENLM_AVAILABLE, reason="KenLM is not available")
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("batch_size", [1, 3])
+    @pytest.mark.parametrize("bos", [True, False])
+    def test_initial_states(
+        self,
+        n_gpu_lm: NGramGPULanguageModel,
+        kenlm_wrapper: KenLMBatchedWrapper,
+        bos: bool,
+        batch_size: int,
+        device: torch.device,
+    ):
+        n_gpu_lm = n_gpu_lm.to(device)
+        init_states = n_gpu_lm.get_init_states(batch_size=batch_size, bos=bos)
+        init_states_kenlm = kenlm_wrapper.get_init_states(batch_size=batch_size, bos=bos)
+        scores_lm, _ = n_gpu_lm.advance(init_states)
+        scores_ref, _ = kenlm_wrapper.advance(init_states_kenlm)
+        assert torch.allclose(scores_lm, scores_ref.to(device))
+
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not TRITON_AVAILABLE, reason="Triton is not available")
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
+    def test_triton_vs_pytorch_random_states(self, n_gpu_lm: NGramGPULanguageModel, batch_size=2, num_iterations=100):
+        """Randomly initializes the states and compares the scores from Triton and PyTorch implementations."""
+        torch.manual_seed(777)
+        device = torch.device("cuda")
+        n_gpu_lm = n_gpu_lm.to(device)
+        for _ in tqdm(range(num_iterations)):
+            start_state = random.randint(0, n_gpu_lm.num_states - 1)
+            with torch.no_grad():
+                scores1, states1 = n_gpu_lm._advance_pytorch(
+                    states=torch.full([batch_size], fill_value=start_state, device=device, dtype=torch.int64)
+                )
+                scores2, states2 = n_gpu_lm._advance_triton(
+                    states=torch.full([batch_size], fill_value=start_state, device=device, dtype=torch.int64)
+                )
+            assert (states1 == states2).all()
+            assert torch.allclose(scores1, scores2)
+
+    @pytest.mark.unit
+    @pytest.mark.skipif(not KENLM_AVAILABLE, reason="KenLM is not available")
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("bos", [True, False])
+    def test_final(
+        self, n_gpu_lm: NGramGPULanguageModel, kenlm_wrapper: KenLMBatchedWrapper, bos: bool, device: torch.device
+    ):
+        """Test final (eos) scores"""
+        n_gpu_lm = n_gpu_lm.to(device)
+        sentences = [
+            [25, 70, 12],
+            [58, 41, 186, 293, 306, 999, 163, 264, 689, 683, 999],
+            [],  # empty sentence
+        ]
+        last_states = []
+        for sentence in sentences:
+            state = kenlm_wrapper.get_init_state(bos=bos)
+            for label in sentence:
+                _, state = kenlm_wrapper.advance_single(state=state, label=label)
+            last_states.append(state)
+        final_ref = kenlm_wrapper.get_final(states=last_states).to(device=device)
+
+        last_states = []
+        for sentence in sentences:
+            states = n_gpu_lm.get_init_states(batch_size=1, bos=bos)
+            for label in sentence:
+                _, states = n_gpu_lm.advance(states=states)
+                states = states[0, label].unsqueeze(0)
+            last_states.append(states)
+        final_lm = n_gpu_lm.get_final(states=torch.cat(last_states, dim=0))
+
+        assert torch.allclose(final_lm, final_ref), "Final scores do not match"
+
+    @pytest.mark.unit
+    @pytest.mark.skipif(not KENLM_AVAILABLE, reason="KenLM is not available")
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("bos", [True, False])
+    @pytest.mark.parametrize("eos", [True, False])
+    def test_sentences(
+        self,
+        n_gpu_lm: NGramGPULanguageModel,
+        kenlm_wrapper: KenLMBatchedWrapper,
+        bos: bool,
+        eos: bool,
+        device: torch.device,
+    ):
+        n_gpu_lm = n_gpu_lm.to(device)
+        sentences = [
+            [25, 70, 12],
+            [58, 41, 186, 293, 306, 999, 163, 264, 689, 683, 999],
+            [],  # empty sentence
+        ]
+        # non-batched
+        for sentence in sentences:
+            scores_ref = kenlm_wrapper.score_sentences([sentence], bos=bos, eos=eos).to(device)
+            scores_lm = n_gpu_lm(
+                labels=torch.LongTensor([sentence]).to(device),
+                bos=bos,
+                eos=eos,
+            )
+            assert torch.allclose(scores_ref, scores_lm), "Non-batched scores do not match"
+
+        # batched
+        scores_ref = kenlm_wrapper.score_sentences(sentences, bos=bos, eos=eos).to(device)
+        scores_lm = n_gpu_lm(
+            labels=pad_sequence([torch.LongTensor(sentence) for sentence in sentences], batch_first=True).to(device),
+            labels_lengths=torch.LongTensor([len(sentence) for sentence in sentences]).to(device),
+            bos=bos,
+            eos=eos,
+        )
+        assert torch.allclose(scores_lm, scores_ref), "Batched scores do not match"
+
+    @pytest.mark.unit
+    def test_save_load_nemo(self, tmp_path, test_data_dir):
+        vocab_size = 1024
+        kenlm_model_path = Path(test_data_dir) / "asr/kenlm_ngram_lm/parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+        n_gpu_lm = NGramGPULanguageModel.from_arpa(kenlm_model_path, vocab_size=vocab_size, normalize_unk=False)
+        nemo_path = tmp_path / "ngram_lm.nemo"
+        n_gpu_lm.save_to(f"{nemo_path}")
+        n_gpu_lm_loaded = NGramGPULanguageModel.from_nemo(f"{nemo_path}", vocab_size=vocab_size)
+
+        # arcs data
+        assert torch.allclose(n_gpu_lm_loaded.arcs_weights, n_gpu_lm.arcs_weights)
+        assert (n_gpu_lm_loaded.from_states == n_gpu_lm.from_states).all()
+        assert (n_gpu_lm_loaded.to_states == n_gpu_lm.to_states).all()
+        assert (n_gpu_lm_loaded.ilabels == n_gpu_lm.ilabels).all()
+
+        # states data
+        assert (n_gpu_lm_loaded.start_end_arcs == n_gpu_lm.start_end_arcs).all()
+        assert (n_gpu_lm_loaded.state_order == n_gpu_lm.state_order).all()
+        assert (n_gpu_lm_loaded.backoff_to_states == n_gpu_lm.backoff_to_states).all()
+        assert torch.allclose(n_gpu_lm_loaded.backoff_weights, n_gpu_lm.backoff_weights)
+        assert torch.allclose(n_gpu_lm_loaded.final_weights, n_gpu_lm.final_weights)
+
+    @pytest.mark.unit
+    def test_save_load_from_file(self, tmp_path, test_data_dir):
+        vocab_size = 1024
+        kenlm_model_path = Path(test_data_dir) / "asr/kenlm_ngram_lm/parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+        n_gpu_lm = NGramGPULanguageModel.from_file(kenlm_model_path, vocab_size=vocab_size, normalize_unk=False)
+        nemo_path = tmp_path / "ngram_lm.nemo"
+        n_gpu_lm.save_to(f"{nemo_path}")
+        n_gpu_lm_loaded = NGramGPULanguageModel.from_file(f"{nemo_path}", vocab_size=vocab_size)
+
+        # arcs data
+        assert torch.allclose(n_gpu_lm_loaded.arcs_weights, n_gpu_lm.arcs_weights)
+        assert (n_gpu_lm_loaded.from_states == n_gpu_lm.from_states).all()
+        assert (n_gpu_lm_loaded.to_states == n_gpu_lm.to_states).all()
+        assert (n_gpu_lm_loaded.ilabels == n_gpu_lm.ilabels).all()
+
+        # states data
+        assert (n_gpu_lm_loaded.start_end_arcs == n_gpu_lm.start_end_arcs).all()
+        assert (n_gpu_lm_loaded.state_order == n_gpu_lm.state_order).all()
+        assert (n_gpu_lm_loaded.backoff_to_states == n_gpu_lm.backoff_to_states).all()
+        assert torch.allclose(n_gpu_lm_loaded.backoff_weights, n_gpu_lm.backoff_weights)
+        assert torch.allclose(n_gpu_lm_loaded.final_weights, n_gpu_lm.final_weights)

tests/collections/asr/test_padding_and_batch_size_invariance.py

@@ -0,0 +1,145 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import pytest
+import torch.testing
+from lhotse.testing.random import deterministic_rng
+
+from nemo.collections.asr.modules import AudioToMelSpectrogramPreprocessor, ConformerEncoder
+from nemo.collections.asr.parts.preprocessing import FilterbankFeatures
+
+
+@pytest.mark.parametrize("length", list(range(15950, 16050, 3)))
+def test_preprocessor_invariant_to_padding(deterministic_rng, length):
+    # Settings corresponding to Canary-1B features
+    f = FilterbankFeatures(n_window_size=400, nfilt=128, pad_to=0).eval()
+
+    # Test data:
+    # * a1: 1s "audio"
+    # * a2: 1s "audio" + 1s padding, keep length tensor unchanged
+    a1 = torch.arange(0, length).unsqueeze(0) / 16000
+    a1l = torch.tensor([length])
+
+    a2 = torch.cat([a1, torch.zeros(1, 16000)], dim=1)
+    a2l = a1l.clone()
+
+    mels1, mels1l = f(a1, a1l)
+    mels2, mels2l = f(a2, a2l)
+
+    # Ideally, we'd have strictly identical results.
+    # However, we observed depending on PyTorch build and environment,
+    # Mel-spectrogram normalization tends to yield non-deterministic results;
+    # specifically, in the computation of numerator in
+    # nemo.collections.asr.parts.preprocessing.features.normalize_batch
+    # where identical inputs lead up to +/- 2e-3 numerical differences.
+    torch.testing.assert_close(mels1[..., :mels1l], mels2[..., :mels1l], atol=5e-2, rtol=0)
+
+
+@pytest.mark.parametrize("length", [16000])
+def test_canary_encoder_invariant_to_padding(deterministic_rng, length):
+    preprocessor = AudioToMelSpectrogramPreprocessor(
+        sample_rate=16000,
+        normalize="per_feature",
+        window_size=0.025,
+        window_stride=0.01,
+        window="hann",
+        features=128,
+        n_fft=512,
+        log=True,
+        frame_splicing=1,
+        dither=1e-5,
+        pad_to=0,
+        pad_value=0.0,
+    ).eval()
+    encoder = ConformerEncoder(
+        feat_in=128,
+        feat_out=-1,
+        n_layers=17,
+        d_model=512,
+        subsampling="dw_striding",
+        subsampling_factor=8,
+        subsampling_conv_channels=256,
+        causal_downsampling=True,
+        reduction=None,
+        reduction_factor=1,
+        ff_expansion_factor=4,
+        self_attention_model="rel_pos",
+        n_heads=8,
+        att_context_size=[-1, -1],
+        xscaling=False,
+        untie_biases=True,
+        pos_emb_max_len=5000,
+        conv_kernel_size=9,
+        conv_norm_type="batch_norm",
+        conv_context_size=None,
+        dropout=0.1,
+        dropout_pre_encoder=0.1,
+        dropout_emb=0.0,
+        dropout_att=0.1,
+    ).eval()
+
+    # Test data:
+    # * a1: 1s "audio"
+    # * a2: 1s "audio" + 1s padding, keep length tensor unchanged
+    a1 = torch.arange(0, length).unsqueeze(0) / 16000
+    a1l = torch.tensor([length])
+
+    a2 = torch.cat([a1, torch.zeros(1, 16000)], dim=1)
+    a2l = a1l.clone()
+
+    mels1, mels1l = preprocessor(input_signal=a1, length=a1l)
+    mels2, mels2l = preprocessor(input_signal=a2, length=a2l)
+
+    torch.testing.assert_close(mels1[..., :mels1l], mels2[..., :mels1l], atol=5e-4, rtol=0)
+
+    # SUBSAMPLING MODULE NOT MISMATCHING
+    activation = {}
+
+    def get_activation(name):
+        def hook(model, input, output):
+            activation[name] = torch.tensor(output.detach().tolist())
+
+        return hook
+
+    for i, layer in enumerate(encoder.pre_encode.conv):
+        if "ReLU" in str(layer):
+            continue
+        layer.register_forward_hook(get_activation(f"{i}:{layer}"))
+    h1, h1l = encoder.pre_encode(mels1.transpose(1, 2), mels1l)
+    inner1 = activation.copy()
+    h2, h2l = encoder.pre_encode(mels2.transpose(1, 2), mels2l)
+    inner2 = activation
+    for k in inner1:
+        torch.testing.assert_close(inner1[k], inner2[k][:, :, : inner1[k].shape[2]], atol=5e-5, rtol=0)
+
+    torch.testing.assert_close(h1[:, :h1l], h2[:, :h1l])
+
+    h1, h1l = encoder(audio_signal=mels1, length=mels1l)
+    h2, h2l = encoder(audio_signal=mels2, length=mels2l)
+
+    torch.testing.assert_close(h1[..., :h1l], h2[..., :h1l])
+
+
+def test_conformer_inference_invariant_to_batch_size(deterministic_rng):
+    model = ConformerEncoder(feat_in=128, n_layers=2, d_model=128, feat_out=128)
+    model = model.eval()
+
+    audio_signal_bs1, length_bs1 = model.input_example()
+    h_bs1, h_length_bs1 = model(audio_signal=audio_signal_bs1, length=length_bs1)
+
+    audio_signal_bs2 = audio_signal_bs1.repeat(2, 1, 1)
+    length_bs2 = length_bs1.repeat(2)
+    h_bs2, h_length_bs2 = model(audio_signal=audio_signal_bs2, length=length_bs2)
+
+    torch.testing.assert_close(h_bs1, h_bs2[:1])
+    torch.testing.assert_close(h_bs1, h_bs2[1:])

tests/collections/asr/test_parallel_expert_encoder.py

@@ -0,0 +1,639 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import io
+import tarfile
+
+import pytest
+import torch
+import torch.distributed as dist
+from omegaconf import DictConfig, OmegaConf
+from torch import nn
+
+from nemo.collections.asr.models import SortformerEncLabelModel
+from nemo.collections.asr.modules.conformer_encoder import ConformerEncoder
+from nemo.collections.asr.modules.parallel_expert_encoder import (
+    ParallelExpertEncoder,
+    ParallelExpertEncoderPT,
+    _clone_config,
+    _default_dtype,
+    _disable_dist_feature_sync,
+)
+
+# ``@experimental`` wraps the class in a wrapt proxy, so ``__new__`` (used to build
+# bare instances that skip the heavy real ``__init__``) must target the underlying
+# class. Attribute access / isinstance still go through the proxy name.
+_PEE = getattr(ParallelExpertEncoder, "__wrapped__", ParallelExpertEncoder)
+
+
+# ----------------------------------------------------------------------------- #
+# Module-level context managers / helpers
+# ----------------------------------------------------------------------------- #
+@pytest.mark.unit
+def test_clone_config_is_deep_and_handles_none():
+    cfg = OmegaConf.create({"a": {"b": 1}})
+    clone = _clone_config(cfg)
+    assert clone == cfg
+    clone.a.b = 2
+    assert cfg.a.b == 1  # original untouched
+    assert _clone_config(None) is None
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("target_dtype", [torch.float64, torch.float16])
+def test_default_dtype_sets_and_restores(target_dtype):
+    prev = torch.get_default_dtype()
+    with _default_dtype(target_dtype):
+        assert torch.get_default_dtype() == target_dtype
+    assert torch.get_default_dtype() == prev
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("noop_dtype", [torch.get_default_dtype(), torch.int32])
+def test_default_dtype_noop_paths(noop_dtype):
+    # Same-dtype and non-floating dtype are both no-ops.
+    prev = torch.get_default_dtype()
+    with _default_dtype(noop_dtype):
+        assert torch.get_default_dtype() == prev
+    assert torch.get_default_dtype() == prev
+
+
+@pytest.mark.unit
+def test_disable_dist_feature_sync_noop_when_uninitialized():
+    assert not dist.is_initialized()
+    orig = dist.is_initialized
+    with _disable_dist_feature_sync():
+        pass
+    assert dist.is_initialized is orig  # nothing patched when dist is down
+
+
+# ----------------------------------------------------------------------------- #
+# Static pure helpers on ParallelExpertEncoder
+# ----------------------------------------------------------------------------- #
+@pytest.mark.unit
+@pytest.mark.parametrize("max_pos, dim", [(4, 8), (1, 16), (10, 4)])
+def test_build_sinusoid_position_encoding(max_pos, dim):
+    pe = ParallelExpertEncoder._build_sinusoid_position_encoding(max_pos, dim)
+    assert pe.shape == (max_pos, dim)
+    # row 0: sin(0)=0 on even indices, cos(0)=1 on odd indices
+    assert torch.allclose(pe[0, 0::2], torch.zeros(dim // 2))
+    assert torch.allclose(pe[0, 1::2], torch.ones(dim // 2))
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "cur_len, target_len",
+    [(3, 6), (6, 3), (5, 5), (1, 4)],
+)
+def test_align_diar_frames_length_and_padding(cur_len, target_len):
+    n_spk = 3
+    diar = torch.arange(cur_len * n_spk, dtype=torch.float32).reshape(1, cur_len, n_spk)
+    out = ParallelExpertEncoder._align_diar_frames(diar, target_len)
+    assert out.shape == (1, target_len, n_spk)
+    if target_len <= cur_len:
+        # truncation keeps the leading frames unchanged
+        assert torch.equal(out, diar[:, :target_len, :])
+    else:
+        # padding repeats the last frame
+        assert torch.equal(out[:, :cur_len, :], diar)
+        for t in range(cur_len, target_len):
+            assert torch.equal(out[:, t, :], diar[:, -1, :])
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("param_dtype", [torch.float64, torch.float16])
+def test_match_module_io_casts_to_param_dtype(param_dtype):
+    module = nn.Linear(4, 4).to(param_dtype)
+    tensor = torch.zeros(2, 4, dtype=torch.float32)
+    out = ParallelExpertEncoder._match_module_io(tensor, module)
+    assert out.dtype == param_dtype
+
+
+@pytest.mark.unit
+def test_match_module_io_paramless_module_unchanged():
+    module = nn.Identity()  # no parameters
+    tensor = torch.zeros(2, 4, dtype=torch.float32)
+    out = ParallelExpertEncoder._match_module_io(tensor, module)
+    assert out.dtype == torch.float32
+    assert out is tensor
+
+
+# ----------------------------------------------------------------------------- #
+# forward() offline/online dispatch
+# ----------------------------------------------------------------------------- #
+def dispatch_stub(online_inference_length, chunk_feat_len, training):
+    """Build a bare ParallelExpertEncoder with stubbed branch methods."""
+    enc = _PEE.__new__(_PEE)
+    nn.Module.__init__(enc)
+    enc.online_inference_length = online_inference_length
+    enc.chunk_feat_len = chunk_feat_len
+    enc.training = training
+    enc._forward = lambda **kw: "offline"
+    enc._forward_online = lambda **kw: "online"
+    return enc
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "online_len, chunk_feat_len, training, n_frames, expected",
+    [
+        (500, 100, False, 200, "online"),  # eval + long enough -> online
+        (500, 100, False, 50, "offline"),  # eval but shorter than one window
+        (500, 100, True, 200, "offline"),  # training always offline
+        (0, 100, False, 200, "offline"),  # online disabled
+        (500, 100, False, 100, "offline"),  # exactly one window (not strictly greater)
+    ],
+)
+def test_forward_dispatch(online_len, chunk_feat_len, training, n_frames, expected):
+    enc = dispatch_stub(online_len, chunk_feat_len, training)
+    audio = torch.zeros(1, 8, n_frames)
+    length = torch.tensor([n_frames])
+    assert enc.forward(audio, length) == expected
+
+
+# ----------------------------------------------------------------------------- #
+# _forward_online orchestration (stubbed ASR encoder, provided spk_targets)
+# ----------------------------------------------------------------------------- #
+class _FakeASR(nn.Module):
+    """Minimal stand-in for the wrapped ConformerEncoder."""
+
+    def __init__(self, d_model: int, sf: int):
+        super().__init__()
+        self.subsampling_factor = sf
+        self.d_model = d_model
+        self._p = nn.Parameter(torch.zeros(1))
+
+    def forward(self, audio_signal, length):
+        b, _, t = audio_signal.shape
+        # generous frame count so the trim logic never clamps
+        t_out = (t + self.subsampling_factor - 1) // self.subsampling_factor + 8
+        out = torch.randn(b, self.d_model, t_out)
+        return out, length // self.subsampling_factor
+
+
+def online_stub(d_model, n_spk, sf, win, lc, rc):
+    enc = _PEE.__new__(_PEE)
+    nn.Module.__init__(enc)
+    enc.asr_encoder = _FakeASR(d_model, sf)
+    enc.asr_normalize_type = None
+    enc.online_inference_length = win
+    enc.chunk_left_context = lc
+    enc.chunk_right_context = rc
+    enc.chunk_feat_len = win * sf
+    enc.left_ctx_feat_len = lc * sf
+    enc.right_ctx_feat_len = rc * sf
+    enc.freeze_asr = True
+    enc.freeze_diar = False  # The stub has no `diarization_model`, so `freeze_diar` must be False to keep
+    enc.asr_norm = nn.LayerNorm(d_model)
+    enc.diar_norm = nn.LayerNorm(n_spk)
+    enc.register_buffer("diar_kernel", torch.randn(n_spk, d_model))
+    enc._suppress_online_pbar = True
+    enc.eval()
+    return enc
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "sf, win, lc, rc, n_frames",
+    [
+        (8, 10, 2, 2, 240),  # 3 full chunks
+        (8, 10, 0, 0, 200),  # partial last chunk, no context
+        (4, 5, 1, 1, 64),  # 4 chunks, small subsampling
+        (8, 50, 5, 5, 160),  # single chunk (n_frames < window)
+    ],
+)
+def test_forward_online_output_length_telescopes(sf, win, lc, rc, n_frames):
+    d_model, n_spk, b = 16, 4, 2
+    enc = online_stub(d_model, n_spk, sf, win, lc, rc)
+
+    mels = torch.randn(b, 80, n_frames)
+    length = torch.tensor([n_frames] * b)
+    spk_targets = torch.rand(b, 5, n_spk)  # arbitrary; aligned internally
+
+    outputs, encoded_len = enc._forward_online(audio_signal=mels, length=length, spk_targets=spk_targets)
+
+    expected_t = round(n_frames / sf)
+    assert outputs.shape == (b, d_model, expected_t)
+    assert encoded_len.tolist() == [expected_t] * b
+
+
+# ----------------------------------------------------------------------------- #
+# ParallelExpertEncoderPT.is_pe_nemo
+# ----------------------------------------------------------------------------- #
+def write_nemo(path, *, target=None, include_cfg=True):
+    with tarfile.open(path, "w") as tf:
+        if include_cfg:
+            data = (f"target: {target}\n" if target is not None else "foo: bar\n").encode()
+            info = tarfile.TarInfo(name="model_config.yaml")
+            info.size = len(data)
+            tf.addfile(info, io.BytesIO(data))
+        else:
+            data = b"not a config"
+            info = tarfile.TarInfo(name="weights.ckpt")
+            info.size = len(data)
+            tf.addfile(info, io.BytesIO(data))
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "target, expected",
+    [
+        ("nemo.collections.asr.modules.parallel_expert_encoder.ParallelExpertEncoderPT", True),
+        ("ParallelExpertEncoderPT", True),
+        ("nemo.collections.asr.models.SomethingElse", False),
+        (None, False),  # model_config.yaml present but no `target`
+    ],
+)
+def test_is_pe_nemo_by_target(tmp_path, target, expected):
+    nemo_path = str(tmp_path / "bundle.nemo")
+    write_nemo(nemo_path, target=target)
+    assert ParallelExpertEncoderPT.is_pe_nemo(nemo_path) is expected
+
+
+@pytest.mark.unit
+def test_is_pe_nemo_without_model_config(tmp_path):
+    nemo_path = str(tmp_path / "no_cfg.nemo")
+    write_nemo(nemo_path, include_cfg=False)
+    assert ParallelExpertEncoderPT.is_pe_nemo(nemo_path) is False
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "bad_path",
+    [None, 123, "missing.nemo", "not_a_nemo.txt"],
+)
+def test_is_pe_nemo_rejects_bad_paths(tmp_path, bad_path):
+    # a real-but-non-.nemo file to exercise the suffix check
+    if bad_path == "not_a_nemo.txt":
+        p = tmp_path / "not_a_nemo.txt"
+        p.write_text("hello")
+        bad_path = str(p)
+    assert ParallelExpertEncoderPT.is_pe_nemo(bad_path) is False
+
+
+# ----------------------------------------------------------------------------- #
+# ParallelExpertEncoderPT.save_to_nemo guard rails
+# ----------------------------------------------------------------------------- #
+@pytest.mark.unit
+def test_save_to_nemo_rejects_non_encoder(tmp_path):
+    with pytest.raises(TypeError):
+        ParallelExpertEncoderPT.save_to_nemo(
+            nn.Linear(2, 2), str(tmp_path / "out.nemo"), template_bundle_path=str(tmp_path / "tpl.nemo")
+        )
+
+
+@pytest.mark.unit
+def test_save_to_nemo_missing_template(tmp_path):
+    # __new__ produces a real ParallelExpertEncoder instance (passes isinstance)
+    # without running the heavy __init__, so we reach the template existence check.
+    fake_encoder = _PEE.__new__(_PEE)
+    with pytest.raises(FileNotFoundError):
+        ParallelExpertEncoderPT.save_to_nemo(
+            fake_encoder,
+            str(tmp_path / "out.nemo"),
+            template_bundle_path=str(tmp_path / "does_not_exist.nemo"),
+        )
+
+
+# ----------------------------------------------------------------------------- #
+# End-to-end fusion with real toy encoders
+#
+# ParallelExpertEncoder loads two real sub-encoders and fuses them:
+#   * an ASR ConformerEncoder (cf. tests/collections/asr/test_conformer_encoder.py)
+#   * a Sortformer diarizer    (cf. tests/collections/speaker_tasks/test_diar_sortformer_models.py)
+# These tests build tiny-but-real instances of both and run the wrapper end to end.
+# ----------------------------------------------------------------------------- #
+_MEL_FEATURES = 128
+_ASR_D_MODEL = 32
+_DIAR_FC_D_MODEL = 32
+_DIAR_TF_D_MODEL = 16
+_N_SPK = 4
+_SUBSAMPLING_FACTOR = 8
+
+
+def toy_asr_encoder_cfg() -> DictConfig:
+    """Tiny ConformerEncoder config the PE encoder mounts as its ASR branch."""
+    return DictConfig(
+        {
+            '_target_': 'nemo.collections.asr.modules.ConformerEncoder',
+            'feat_in': _MEL_FEATURES,
+            'feat_out': -1,
+            'n_layers': 1,
+            'd_model': _ASR_D_MODEL,
+            'subsampling': 'dw_striding',
+            'subsampling_factor': _SUBSAMPLING_FACTOR,
+            'subsampling_conv_channels': 16,
+            'ff_expansion_factor': 4,
+            'self_attention_model': 'rel_pos',
+            'n_heads': 4,
+            'att_context_size': [-1, -1],
+            'conv_kernel_size': 9,
+            'dropout': 0.0,
+            'dropout_pre_encoder': 0.0,
+            'dropout_emb': 0.0,
+            'dropout_att': 0.0,
+        }
+    )
+
+
+def toy_diarization_model_cfg() -> DictConfig:
+    """Tiny SortformerEncLabelModel config the PE encoder mounts as its diar branch."""
+    model_defaults = {'fc_d_model': _DIAR_FC_D_MODEL, 'tf_d_model': _DIAR_TF_D_MODEL}
+    return DictConfig(
+        {
+            'target': 'nemo.collections.asr.models.sortformer_diar_models.SortformerEncLabelModel',
+            'sample_rate': 16000,
+            'pil_weight': 0.5,
+            'ats_weight': 0.5,
+            'max_num_of_spks': _N_SPK,
+            'streaming_mode': False,
+            'async_streaming': False,
+            'model_defaults': DictConfig(model_defaults),
+            'preprocessor': DictConfig(
+                {
+                    '_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor',
+                    'normalize': 'per_feature',
+                    'window_size': 0.025,
+                    'sample_rate': 16000,
+                    'window_stride': 0.01,
+                    'window': 'hann',
+                    'features': _MEL_FEATURES,
+                    'n_fft': 512,
+                    'frame_splicing': 1,
+                    'dither': 0.00001,
+                }
+            ),
+            'encoder': DictConfig(
+                {
+                    '_target_': 'nemo.collections.asr.modules.ConformerEncoder',
+                    'feat_in': _MEL_FEATURES,
+                    'feat_out': -1,
+                    'n_layers': 1,
+                    'd_model': _DIAR_FC_D_MODEL,
+                    'subsampling': 'dw_striding',
+                    'subsampling_factor': _SUBSAMPLING_FACTOR,
+                    'subsampling_conv_channels': 16,
+                    'causal_downsampling': False,
+                    'ff_expansion_factor': 4,
+                    'self_attention_model': 'rel_pos',
+                    'n_heads': 4,
+                    'att_context_size': [-1, -1],
+                    'conv_kernel_size': 9,
+                    'conv_norm_type': 'batch_norm',
+                    'dropout': 0.0,
+                    'dropout_pre_encoder': 0.0,
+                    'dropout_emb': 0.0,
+                    'dropout_att': 0.0,
+                }
+            ),
+            'transformer_encoder': DictConfig(
+                {
+                    '_target_': 'nemo.collections.asr.modules.transformer.transformer_encoders.TransformerEncoder',
+                    'num_layers': 1,
+                    'hidden_size': _DIAR_TF_D_MODEL,
+                    'inner_size': 32,
+                    'num_attention_heads': 4,
+                    'attn_score_dropout': 0.0,
+                    'attn_layer_dropout': 0.0,
+                    'ffn_dropout': 0.0,
+                    'hidden_act': 'relu',
+                    'pre_ln': False,
+                    'pre_ln_final_layer_norm': True,
+                }
+            ),
+            'sortformer_modules': DictConfig(
+                {
+                    '_target_': 'nemo.collections.asr.modules.sortformer_modules.SortformerModules',
+                    'num_spks': _N_SPK,
+                    'dropout_rate': 0.0,
+                    'fc_d_model': _DIAR_FC_D_MODEL,
+                    'tf_d_model': _DIAR_TF_D_MODEL,
+                }
+            ),
+            'loss': DictConfig(
+                {
+                    '_target_': 'nemo.collections.asr.losses.bce_loss.BCELoss',
+                    'weight': None,
+                    'reduction': 'mean',
+                }
+            ),
+        }
+    )
+
+
+def build_toy_pe_encoder(**overrides) -> ParallelExpertEncoder:
+    """Construct a real ParallelExpertEncoder from the tiny ASR + diar configs."""
+    kwargs = dict(
+        asr_encoder_cfg=toy_asr_encoder_cfg(),
+        diarization_model_cfg=toy_diarization_model_cfg(),
+        asr_normalize_type='per_feature',
+        # Keep the input far below one window so forward() stays on the offline path.
+        online_inference_length=500,
+    )
+    kwargs.update(overrides)
+    return ParallelExpertEncoder(**kwargs)
+
+
+@pytest.mark.unit
+def test_pe_encoder_builds_and_wires_both_real_encoders():
+    enc = build_toy_pe_encoder()
+    # The two fused sub-encoders are the real classes, not stubs.
+    assert isinstance(enc.asr_encoder, ConformerEncoder)
+    assert isinstance(enc.diarization_model, SortformerEncLabelModel)
+    # ConformerEncoder-compatible drop-in properties come from the ASR branch.
+    assert enc.d_model == _ASR_D_MODEL
+    assert enc.subsampling_factor == _SUBSAMPLING_FACTOR
+    # Speaker count + fusion kernel come from the diar branch.
+    assert enc.n_spk == _N_SPK
+    assert enc.diar_kernel.shape == (_N_SPK, _ASR_D_MODEL)
+    # freeze_diar defaults to True -> diar params are frozen, ASR params remain trainable.
+    assert all(not p.requires_grad for p in enc.diarization_model.parameters())
+    assert any(p.requires_grad for p in enc.asr_encoder.parameters())
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("batch_size, n_frames", [(1, 160), (2, 200)])
+def test_pe_encoder_offline_forward_runs_internal_diarizer(batch_size, n_frames):
+    enc = build_toy_pe_encoder().eval()
+    mels = torch.randn(batch_size, _MEL_FEATURES, n_frames)
+    length = torch.full((batch_size,), n_frames, dtype=torch.long)
+
+    with torch.no_grad():
+        outputs, encoded_len = enc(mels, length)  # spk_targets=None -> Sortformer runs internally
+
+    expected_t = int(encoded_len[0].item())
+    assert outputs.shape == (batch_size, _ASR_D_MODEL, expected_t)
+    assert expected_t > 0
+    assert torch.isfinite(outputs).all()
+    assert encoded_len.tolist() == [expected_t] * batch_size
+
+
+@pytest.mark.unit
+def test_pe_encoder_offline_forward_accepts_diar_override_and_fuses_it():
+    enc = build_toy_pe_encoder().eval()
+    batch_size, n_frames = 2, 160
+    mels = torch.randn(batch_size, _MEL_FEATURES, n_frames)
+    length = torch.full((batch_size,), n_frames, dtype=torch.long)
+
+    # Arbitrary diar frame count: PE aligns it to the ASR frame count internally.
+    dp1 = torch.rand(batch_size, 7, _N_SPK)
+    dp2 = torch.rand(batch_size, 7, _N_SPK)
+
+    with torch.no_grad():
+        out1, len1 = enc(mels, length, spk_targets=dp1)
+        out2, len2 = enc(mels, length, spk_targets=dp2)
+
+    expected_t = int(len1[0].item())
+    assert out1.shape == (batch_size, _ASR_D_MODEL, expected_t)
+    assert torch.equal(len1, len2)
+    assert torch.isfinite(out1).all()
+    # Same audio + same (dropout-free, eval) ASR branch, but different speaker
+    # predictions must change the fused output -> proves the diar branch is fused in.
+    assert not torch.allclose(out1, out2)
+
+
+@pytest.mark.unit
+def test_pe_encoder_online_forward_matches_conformer_io_with_real_encoders():
+    # Small window so a modest input crosses onto the long-form online path.
+    enc = build_toy_pe_encoder(
+        online_inference_length=10,
+        chunk_left_context=2,
+        chunk_right_context=2,
+        diar_fifo_len=10,
+        diar_spkcache_update_period=20,
+        diar_spkcache_len=20,
+    ).eval()
+    enc._suppress_online_pbar = True
+
+    batch_size, n_frames = 1, 320  # > online_inference_length * subsampling_factor (=80)
+    mels = torch.randn(batch_size, _MEL_FEATURES, n_frames)
+    length = torch.full((batch_size,), n_frames, dtype=torch.long)
+
+    with torch.no_grad():
+        outputs, encoded_len = enc(mels, length)
+
+    expected_t = int(encoded_len[0].item())
+    assert outputs.shape == (batch_size, _ASR_D_MODEL, expected_t)
+    assert expected_t > 0
+    assert torch.isfinite(outputs).all()
+
+
+# ----------------------------------------------------------------------------- #
+# GPU end-to-end fusion with real toy encoders
+#
+# These mirror the CPU end-to-end tests but run on CUDA. They additionally
+# exercise the device/dtype-bridging machinery the wrapper exists for: fp32 mels
+# fed into (optionally) bf16 experts on the GPU, handled by `_match_module_io`
+# (offline) and `_default_dtype` / `_disable_dist_feature_sync` (online).
+# ----------------------------------------------------------------------------- #
+@pytest.mark.unit
+@pytest.mark.run_only_on('GPU')
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="PEE GPU test requires CUDA")
+@pytest.mark.parametrize("batch_size, n_frames", [(1, 160), (2, 200)])
+def test_pe_encoder_offline_forward_on_gpu(batch_size, n_frames):
+    enc = build_toy_pe_encoder().eval().cuda()
+    # Mels arrive un-normalised in fp32 (the SALM perception contract).
+    mels = torch.randn(batch_size, _MEL_FEATURES, n_frames, device="cuda", dtype=torch.float32)
+    length = torch.full((batch_size,), n_frames, dtype=torch.long, device="cuda")
+
+    with torch.no_grad():
+        outputs, encoded_len = enc(mels, length)  # spk_targets=None -> Sortformer runs internally
+
+    expected_t = int(encoded_len[0].item())
+    assert outputs.is_cuda
+    assert outputs.shape == (batch_size, _ASR_D_MODEL, expected_t)
+    assert expected_t > 0
+    assert torch.isfinite(outputs).all()
+    assert encoded_len.tolist() == [expected_t] * batch_size
+
+
+@pytest.mark.unit
+@pytest.mark.run_only_on('GPU')
+@pytest.mark.skipif(
+    not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
+    reason="PEE bf16 GPU test requires CUDA with bf16 support",
+)
+def test_pe_encoder_offline_forward_bf16_experts_on_gpu():
+    # Experts run in bf16 while mels stay fp32 -> exercises `_match_module_io`
+    # device/dtype bridging on both branches before their conv subsampling.
+    enc = build_toy_pe_encoder().eval().cuda().to(torch.bfloat16)
+    batch_size, n_frames = 2, 200
+    mels = torch.randn(batch_size, _MEL_FEATURES, n_frames, device="cuda", dtype=torch.float32)
+    length = torch.full((batch_size,), n_frames, dtype=torch.long, device="cuda")
+
+    with torch.no_grad():
+        outputs, encoded_len = enc(mels, length)
+
+    expected_t = int(encoded_len[0].item())
+    assert outputs.is_cuda
+    assert outputs.dtype == torch.bfloat16
+    assert outputs.shape == (batch_size, _ASR_D_MODEL, expected_t)
+    assert torch.isfinite(outputs).all()
+
+
+@pytest.mark.unit
+@pytest.mark.run_only_on('GPU')
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="PEE GPU test requires CUDA")
+def test_pe_encoder_offline_forward_accepts_diar_override_on_gpu():
+    enc = build_toy_pe_encoder().eval().cuda()
+    batch_size, n_frames = 2, 160
+    mels = torch.randn(batch_size, _MEL_FEATURES, n_frames, device="cuda", dtype=torch.float32)
+    length = torch.full((batch_size,), n_frames, dtype=torch.long, device="cuda")
+
+    dp1 = torch.rand(batch_size, 7, _N_SPK, device="cuda")
+    dp2 = torch.rand(batch_size, 7, _N_SPK, device="cuda")
+
+    with torch.no_grad():
+        out1, len1 = enc(mels, length, spk_targets=dp1)
+        out2, len2 = enc(mels, length, spk_targets=dp2)
+
+    expected_t = int(len1[0].item())
+    assert out1.is_cuda
+    assert out1.shape == (batch_size, _ASR_D_MODEL, expected_t)
+    assert torch.equal(len1, len2)
+    assert torch.isfinite(out1).all()
+    # Different speaker predictions must change the fused output.
+    assert not torch.allclose(out1, out2)
+
+
+@pytest.mark.unit
+@pytest.mark.run_only_on('GPU')
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="PEE GPU test requires CUDA")
+def test_pe_encoder_online_forward_on_gpu():
+    enc = (
+        build_toy_pe_encoder(
+            online_inference_length=10,
+            chunk_left_context=2,
+            chunk_right_context=2,
+            diar_fifo_len=10,
+            diar_spkcache_update_period=20,
+            diar_spkcache_len=20,
+        )
+        .eval()
+        .cuda()
+    )
+    enc._suppress_online_pbar = True
+
+    batch_size, n_frames = 1, 320  # > online_inference_length * subsampling_factor (=80)
+    mels = torch.randn(batch_size, _MEL_FEATURES, n_frames, device="cuda", dtype=torch.float32)
+    length = torch.full((batch_size,), n_frames, dtype=torch.long, device="cuda")
+
+    with torch.no_grad():
+        outputs, encoded_len = enc(mels, length)
+
+    expected_t = int(encoded_len[0].item())
+    assert outputs.is_cuda
+    assert outputs.shape == (batch_size, _ASR_D_MODEL, expected_t)
+    assert expected_t > 0
+    assert torch.isfinite(outputs).all()

tests/collections/asr/test_preprocessing_segment.py

@@ -0,0 +1,532 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+import tempfile
+from collections import namedtuple
+from typing import List, Type, Union
+
+import numpy as np
+import pytest
+import soundfile as sf
+
+from nemo.collections.asr.parts.preprocessing.perturb import NoisePerturbation, ShiftPerturbation, SilencePerturbation
+from nemo.collections.asr.parts.preprocessing.segment import AudioSegment, select_channels
+
+
+class TestSelectChannels:
+    num_samples = 1000
+    max_diff_tol = 1e-9
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("channel_selector", [None, 'average', 0, 1, [0, 1]])
+    def test_single_channel_input(self, channel_selector: Type[Union[str, int, List[int]]]):
+        """Cover the case with single-channel input signal.
+        Channel selector should not do anything in this case.
+        """
+        golden_out = signal_in = np.random.rand(self.num_samples)
+
+        if channel_selector not in [None, 0, 'average']:
+            # Expect a failure if looking for a different channel when input is 1D
+            with pytest.raises(ValueError):
+                # UUT
+                select_channels(signal_in, channel_selector)
+        else:
+            # UUT
+            signal_out = select_channels(signal_in, channel_selector)
+
+            # Check difference
+            max_diff = np.max(np.abs(signal_out - golden_out))
+            assert max_diff < self.max_diff_tol
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("num_channels", [2, 4])
+    @pytest.mark.parametrize("channel_selector", [None, 'average', 0, [1], [0, 1]])
+    def test_multi_channel_input(self, num_channels: int, channel_selector: Type[Union[str, int, List[int]]]):
+        """Cover the case with multi-channel input signal and single-
+        or multi-channel output.
+        """
+        signal_in = np.random.rand(self.num_samples, num_channels)
+
+        # calculate golden output
+        if channel_selector is None:
+            golden_out = signal_in
+        elif channel_selector == 'average':
+            golden_out = np.mean(signal_in, axis=1)
+        else:
+            golden_out = signal_in[:, channel_selector].squeeze()
+
+        # UUT
+        signal_out = select_channels(signal_in, channel_selector)
+
+        # Check difference
+        max_diff = np.max(np.abs(signal_out - golden_out))
+        assert max_diff < self.max_diff_tol
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("num_channels", [1, 2])
+    @pytest.mark.parametrize("channel_selector", [2, [1, 2]])
+    def test_select_more_channels_than_available(
+        self, num_channels: int, channel_selector: Type[Union[str, int, List[int]]]
+    ):
+        """This test is expecting the UUT to fail because we ask for more channels
+        than available in the input signal.
+        """
+        signal_in = np.random.rand(self.num_samples, num_channels)
+
+        # expect failure since we ask for more channels than available
+        with pytest.raises(ValueError):
+            # UUT
+            select_channels(signal_in, channel_selector)
+
+
+class TestAudioSegment:
+
+    sample_rate = 16000
+    signal_duration_sec = 2
+    max_diff_tol = 1e-9
+
+    @property
+    def num_samples(self):
+        return self.sample_rate * self.signal_duration_sec
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("num_channels", [1, 4])
+    @pytest.mark.parametrize("channel_selector", [None, 'average', 0, 1, [0, 1]])
+    def test_init_single_channel(self, num_channels: int, channel_selector: Type[Union[str, int, List[int]]]):
+        """Test the constructor directly."""
+        if num_channels == 1:
+            # samples is a one-dimensional vector for single-channel signal
+            samples = np.random.rand(self.num_samples)
+        else:
+            samples = np.random.rand(self.num_samples, num_channels)
+
+        if (isinstance(channel_selector, int) and channel_selector >= num_channels) or (
+            isinstance(channel_selector, list) and max(channel_selector) >= num_channels
+        ):
+            # Expect a failure if looking for a different channel when input is 1D
+            with pytest.raises(ValueError):
+                # Construct UUT
+                uut = AudioSegment(samples=samples, sample_rate=self.sample_rate, channel_selector=channel_selector)
+        else:
+            # Construct UUT
+            uut = AudioSegment(samples=samples, sample_rate=self.sample_rate, channel_selector=channel_selector)
+
+            # Create golden reference
+            # Note: AudioSegment converts input samples to float32
+            golden_samples = select_channels(samples.astype('float32'), channel_selector)
+            expected_num_channels = 1 if golden_samples.ndim == 1 else golden_samples.shape[1]
+
+            # Test UUT
+            assert uut.num_channels == expected_num_channels
+            assert uut.num_samples == self.num_samples
+            assert uut.sample_rate == self.sample_rate
+            assert uut.duration == self.signal_duration_sec
+            max_diff = np.max(np.abs(uut.samples - golden_samples))
+            assert max_diff < self.max_diff_tol
+
+            # Test zero padding
+            pad_length = 42
+            uut.pad(pad_length, symmetric=False)
+            # compare to golden references
+            assert uut.num_samples == self.num_samples + pad_length
+            assert np.all(uut.samples[-pad_length:] == 0.0)
+            max_diff = np.max(np.abs(uut.samples[:-pad_length] - golden_samples))
+            assert max_diff < self.max_diff_tol
+
+            # Test subsegment
+            start_time = 0.2 * self.signal_duration_sec
+            end_time = 0.5 * self.signal_duration_sec
+            uut.subsegment(start_time=start_time, end_time=end_time)
+            # compare to golden references
+            start_sample = int(round(start_time * self.sample_rate))
+            end_sample = int(round(end_time * self.sample_rate))
+            max_diff = np.max(np.abs(uut.samples - golden_samples[start_sample:end_sample]))
+            assert max_diff < self.max_diff_tol
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("num_channels", [1, 4])
+    @pytest.mark.parametrize("channel_selector", [None, 'average', 0])
+    def test_from_file(self, num_channels, channel_selector):
+        """Test loading a signal from a file."""
+        with tempfile.TemporaryDirectory() as test_dir:
+            # Prepare a wav file
+            audio_file = os.path.join(test_dir, 'audio.wav')
+            if num_channels == 1:
+                # samples is a one-dimensional vector for single-channel signal
+                samples = np.random.rand(self.num_samples)
+            else:
+                samples = np.random.rand(self.num_samples, num_channels)
+            sf.write(audio_file, samples, self.sample_rate, 'float')
+
+            # Create UUT
+            uut = AudioSegment.from_file(audio_file, channel_selector=channel_selector)
+
+            # Create golden reference
+            # Note: AudioSegment converts input samples to float32
+            golden_samples = select_channels(samples.astype('float32'), channel_selector)
+            expected_num_channels = 1 if golden_samples.ndim == 1 else golden_samples.shape[1]
+
+            # Test UUT
+            assert uut.num_channels == expected_num_channels
+            assert uut.num_samples == self.num_samples
+            assert uut.sample_rate == self.sample_rate
+            assert uut.duration == self.signal_duration_sec
+            max_diff = np.max(np.abs(uut.samples - golden_samples))
+            assert max_diff < self.max_diff_tol
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("data_channels", [1, 4])
+    @pytest.mark.parametrize("noise_channels", [1, 4])
+    def test_noise_perturb_channels(self, data_channels, noise_channels):
+        """Test loading a signal from a file."""
+        with tempfile.TemporaryDirectory() as test_dir:
+            # Prepare a wav file
+            audio_file = os.path.join(test_dir, 'audio.wav')
+            if data_channels == 1:
+                # samples is a one-dimensional vector for single-channel signal
+                samples = np.random.rand(self.num_samples)
+            else:
+                samples = np.random.rand(self.num_samples, data_channels)
+            sf.write(audio_file, samples, self.sample_rate, 'float')
+
+            noise_file = os.path.join(test_dir, 'noise.wav')
+            if noise_channels == 1:
+                # samples is a one-dimensional vector for single-channel signal
+                samples = np.random.rand(self.num_samples)
+            else:
+                samples = np.random.rand(self.num_samples, noise_channels)
+            sf.write(noise_file, samples, self.sample_rate, 'float')
+
+            manifest_file = os.path.join(test_dir, 'noise_manifest.json')
+            with open(manifest_file, 'w') as fout:
+                item = {'audio_filepath': os.path.abspath(noise_file), 'label': '-', 'duration': 0.1, 'offset': 0.0}
+                fout.write(f'{json.dumps(item)}\n')
+
+            perturber = NoisePerturbation(manifest_file)
+            audio = AudioSegment.from_file(audio_file)
+            noise = AudioSegment.from_file(noise_file)
+
+            if data_channels == noise_channels:
+                try:
+                    _ = perturber.perturb_with_input_noise(audio, noise, ref_mic=0)
+                except ValueError as e:
+                    assert False, "perturb_with_input_noise failed with ref_mic=0"
+
+                with pytest.raises(ValueError):
+                    _ = perturber.perturb_with_input_noise(audio, noise, ref_mic=data_channels)
+
+                try:
+                    _ = perturber.perturb_with_foreground_noise(audio, noise, ref_mic=0)
+                except ValueError as e:
+                    assert False, "perturb_with_foreground_noise failed with ref_mic=0"
+
+                with pytest.raises(ValueError):
+                    _ = perturber.perturb_with_foreground_noise(audio, noise, ref_mic=data_channels)
+            else:
+                with pytest.raises(ValueError):
+                    _ = perturber.perturb_with_input_noise(audio, noise)
+                with pytest.raises(ValueError):
+                    _ = perturber.perturb_with_foreground_noise(audio, noise)
+
+    def test_silence_perturb(self):
+        """Test loading a signal from a file and apply silence perturbation"""
+        with tempfile.TemporaryDirectory() as test_dir:
+            # Prepare a wav file
+            audio_file = os.path.join(test_dir, 'audio.wav')
+            # samples is a one-dimensional vector for single-channel signal
+            samples = np.random.rand(self.num_samples)
+            sf.write(audio_file, samples, self.sample_rate, 'float')
+
+            dur = 2
+            perturber = SilencePerturbation(
+                min_start_silence_secs=dur,
+                max_start_silence_secs=dur,
+                min_end_silence_secs=dur,
+                max_end_silence_secs=dur,
+            )
+
+            audio = AudioSegment.from_file(audio_file)
+            ori_audio_len = len(audio._samples)
+            _ = perturber.perturb(audio)
+
+            assert len(audio._samples) == ori_audio_len + 2 * dur * self.sample_rate
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "num_channels, channel_selectors",
+        [
+            (1, [None, 'average', 0]),
+            (3, [None, 'average', 0, 1, [0, 1]]),
+        ],
+    )
+    @pytest.mark.parametrize("sample_rate", [8000, 16000, 22500])
+    def test_audio_segment_from_file(self, tmpdir, num_channels, channel_selectors, sample_rate):
+        """Test loading and audio signal from a file."""
+        signal_len_sec = 4
+        num_samples = signal_len_sec * sample_rate
+        num_examples = 10
+        rtol, atol = 1e-5, 1e-6
+
+        for n in range(num_examples):
+            # Create a test vector
+            audio_file = os.path.join(tmpdir, f'test_audio_{n:02}.wav')
+            samples = np.random.randn(num_samples, num_channels)
+            sf.write(audio_file, samples, sample_rate, 'float')
+
+            for channel_selector in channel_selectors:
+                if channel_selector is None:
+                    ref_samples = samples
+                elif isinstance(channel_selector, int) or isinstance(channel_selector, list):
+                    ref_samples = samples[:, channel_selector]
+                elif channel_selector == 'average':
+                    ref_samples = np.mean(samples, axis=1)
+                else:
+                    raise ValueError(f'Unexpected value of channel_selector {channel_selector}')
+
+                # 1) Load complete audio
+                # Reference
+                ref_samples = ref_samples.squeeze()
+                ref_channels = 1 if ref_samples.ndim == 1 else ref_samples.shape[1]
+
+                # UUT
+                audio_segment = AudioSegment.from_file(audio_file, channel_selector=channel_selector)
+
+                # Test
+                assert (
+                    audio_segment.sample_rate == sample_rate
+                ), f'channel_selector {channel_selector}, sample rate not matching: {audio_segment.sample_rate} != {sample_rate}'
+                assert (
+                    audio_segment.num_channels == ref_channels
+                ), f'channel_selector {channel_selector}, num channels not matching: {audio_segment.num_channels} != {ref_channels}'
+                assert audio_segment.num_samples == len(
+                    ref_samples
+                ), f'channel_selector {channel_selector}, num samples not matching: {audio_segment.num_samples} != {len(ref_samples)}'
+                assert np.allclose(
+                    audio_segment.samples, ref_samples, rtol=rtol, atol=atol
+                ), f'channel_selector {channel_selector}, samples not matching'
+
+                # 2) Load a with duration=None and offset=None, should load the whole audio
+
+                # UUT
+                audio_segment = AudioSegment.from_file(
+                    audio_file, offset=None, duration=None, channel_selector=channel_selector
+                )
+
+                # Test
+                assert (
+                    audio_segment.sample_rate == sample_rate
+                ), f'channel_selector {channel_selector}, offset {offset}, duration {duration}, sample rate not matching: {audio_segment.sample_rate} != {sample_rate}'
+                assert (
+                    audio_segment.num_channels == ref_channels
+                ), f'channel_selector {channel_selector}, offset {offset}, duration {duration}, num channels not matching: {audio_segment.num_channels} != {ref_channels}'
+                assert audio_segment.num_samples == len(
+                    ref_samples
+                ), f'channel_selector {channel_selector}, offset {offset}, duration {duration}, num samples not matching: {audio_segment.num_samples} != {len(ref_samples)}'
+                assert np.allclose(
+                    audio_segment.samples, ref_samples, rtol=rtol, atol=atol
+                ), f'channel_selector {channel_selector}, offset {offset}, duration {duration}, samples not matching'
+
+                # 3) Load a random segment
+                offset = 0.45 * np.random.rand() * signal_len_sec
+                duration = 0.45 * np.random.rand() * signal_len_sec
+
+                # Reference
+                start = int(offset * sample_rate)
+                end = start + int(duration * sample_rate)
+                ref_samples = ref_samples[start:end, ...]
+
+                # UUT
+                audio_segment = AudioSegment.from_file(
+                    audio_file, offset=offset, duration=duration, channel_selector=channel_selector
+                )
+
+                # Test
+                assert (
+                    audio_segment.sample_rate == sample_rate
+                ), f'channel_selector {channel_selector}, offset {offset}, duration {duration}, sample rate not matching: {audio_segment.sample_rate} != {sample_rate}'
+                assert (
+                    audio_segment.num_channels == ref_channels
+                ), f'channel_selector {channel_selector}, offset {offset}, duration {duration}, num channels not matching: {audio_segment.num_channels} != {ref_channels}'
+                assert audio_segment.num_samples == len(
+                    ref_samples
+                ), f'channel_selector {channel_selector}, offset {offset}, duration {duration}, num samples not matching: {audio_segment.num_samples} != {len(ref_samples)}'
+                assert np.allclose(
+                    audio_segment.samples, ref_samples, rtol=rtol, atol=atol
+                ), f'channel_selector {channel_selector}, offset {offset}, duration {duration}, samples not matching'
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "num_channels, channel_selectors",
+        [
+            (1, [None, 'average', 0]),
+            (3, [None, 'average', 0, 1, [0, 1]]),
+        ],
+    )
+    @pytest.mark.parametrize("offset", [0, 1.5])
+    @pytest.mark.parametrize("duration", [1, 2])
+    def test_audio_segment_multichannel_with_list(self, tmpdir, num_channels, channel_selectors, offset, duration):
+        """Test loading an audio signal from a list of single-channel files."""
+        sample_rate = 16000
+        signal_len_sec = 5
+        num_samples = signal_len_sec * sample_rate
+        rtol, atol = 1e-5, 1e-6
+
+        # Random samples
+        samples = np.random.rand(num_samples, num_channels)
+
+        # Save audio
+        audio_files = []
+        for m in range(num_channels):
+            a_file = os.path.join(tmpdir, f'ch_{m}.wav')
+            sf.write(a_file, samples[:, m], sample_rate)
+            audio_files.append(a_file)
+        mc_file = os.path.join(tmpdir, f'mc.wav')
+        sf.write(mc_file, samples, sample_rate)
+
+        for channel_selector in channel_selectors:
+
+            # UUT: loading audio from a list of files
+            uut_segment = AudioSegment.from_file(
+                audio_file=audio_files, offset=offset, duration=duration, channel_selector=channel_selector
+            )
+
+            # Reference: load from the original file
+            ref_segment = AudioSegment.from_file(
+                audio_file=mc_file, offset=offset, duration=duration, channel_selector=channel_selector
+            )
+
+            # Check
+            assert (
+                uut_segment.sample_rate == ref_segment.sample_rate
+            ), f'channel_selector {channel_selector}: expecting {ref_segment.sample_rate}, but UUT segment has {uut_segment.sample_rate}'
+            assert (
+                uut_segment.num_samples == ref_segment.num_samples
+            ), f'channel_selector {channel_selector}: expecting {ref_segment.num_samples}, but UUT segment has {uut_segment.num_samples}'
+            assert np.allclose(
+                uut_segment.samples, ref_segment.samples, rtol=rtol, atol=atol
+            ), f'channel_selector {channel_selector}: samples not matching'
+
+        # Try to get a channel that is out of range.
+        with pytest.raises(RuntimeError, match="Channel cannot be selected"):
+            AudioSegment.from_file(audio_file=audio_files, channel_selector=num_channels)
+
+        if num_channels > 1:
+            # Try to load a list of multichannel files
+            # This is expected to fail since we only support loading a single-channel signal
+            # from each file when audio_file is a list
+            with pytest.raises(RuntimeError, match="Expecting a single-channel audio signal"):
+                AudioSegment.from_file(audio_file=[mc_file, mc_file])
+
+            with pytest.raises(RuntimeError, match="Expecting a single-channel audio signal"):
+                AudioSegment.from_file(audio_file=[mc_file, mc_file], channel_selector=0)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("target_sr", [8000, 16000])
+    def test_audio_segment_trim_match(self, tmpdir, target_sr):
+        """Test loading and audio signal from a file matches when using a path and a list
+        for different target_sr, int_values and trim setups.
+        """
+        sample_rate = 24000
+        signal_len_sec = 2
+        num_samples = signal_len_sec * sample_rate
+        num_examples = 10
+
+        TrimSetup = namedtuple("TrimSetup", "ref top_db frame_length hop_length")
+        trim_setups = []
+        trim_setups.append(TrimSetup(np.max, 10, 2048, 1024))
+        trim_setups.append(TrimSetup(1.0, 35, 2048, 1024))
+        trim_setups.append(TrimSetup(0.8, 45, 2048, 1024))
+
+        for n in range(num_examples):
+            # Create a test vector
+            audio_file = os.path.join(tmpdir, f'test_audio_{n:02}.wav')
+            samples = np.random.randn(num_samples)
+            # normalize
+            samples = samples / np.max(samples)
+            # apply random scaling and window to have some samples cut by trim
+            samples = np.random.rand() * np.hanning(num_samples) * samples
+            sf.write(audio_file, samples, sample_rate, 'float')
+
+            for trim_setup in trim_setups:
+                # UUT 1: load from a path
+                audio_segment_1 = AudioSegment.from_file(
+                    audio_file,
+                    target_sr=target_sr,
+                    trim=True,
+                    trim_ref=trim_setup.ref,
+                    trim_top_db=trim_setup.top_db,
+                    trim_frame_length=trim_setup.frame_length,
+                    trim_hop_length=trim_setup.hop_length,
+                )
+
+                # UUT 2: load from a list
+                audio_segment_2 = AudioSegment.from_file(
+                    [audio_file],
+                    target_sr=target_sr,
+                    trim=True,
+                    trim_ref=trim_setup.ref,
+                    trim_top_db=trim_setup.top_db,
+                    trim_frame_length=trim_setup.frame_length,
+                    trim_hop_length=trim_setup.hop_length,
+                )
+
+                # Test
+                assert audio_segment_1 == audio_segment_2, f'trim setup {trim_setup}, loaded segments not matching'
+
+
+class TestShiftPerturbation:
+    sample_rate = 16000
+
+    def _make_audio_segment(self, duration_sec=1.0):
+        """Create a simple AudioSegment with a sine wave for testing."""
+        num_samples = int(duration_sec * self.sample_rate)
+        t = np.linspace(0, duration_sec, num_samples, dtype=np.float32)
+        samples = np.sin(2 * np.pi * 440 * t)
+        return AudioSegment(samples=samples, sample_rate=self.sample_rate)
+
+    def test_shift_perturbation_normal(self):
+        """Shift perturbation modifies audio when shift is within duration."""
+        perturb = ShiftPerturbation(min_shift_ms=-5.0, max_shift_ms=5.0)
+        segment = self._make_audio_segment(duration_sec=1.0)
+        original = segment.samples.copy()
+        perturb.perturb(segment)
+        assert segment.samples.shape == original.shape, "Audio length should not change"
+
+    def test_shift_perturbation_short_audio_not_skipped(self):
+        """Shift perturbation should clamp and apply shift for short audio, not silently skip."""
+        perturb = ShiftPerturbation(min_shift_ms=10.0, max_shift_ms=20.0)
+        duration_sec = 0.005  # 5ms — shorter than min_shift_ms
+        segment = self._make_audio_segment(duration_sec=duration_sec)
+        original = segment.samples.copy()
+        perturb.perturb(segment)
+        assert segment.samples.shape == original.shape, "Audio length should not change"
+
+    @pytest.mark.parametrize("duration_sec", [0.001, 0.01, 0.1, 1.0])
+    def test_shift_perturbation_preserves_length(self, duration_sec):
+        """Audio length must be preserved regardless of duration."""
+        perturb = ShiftPerturbation(min_shift_ms=-50.0, max_shift_ms=50.0)
+        segment = self._make_audio_segment(duration_sec=duration_sec)
+        original_len = len(segment.samples)
+        perturb.perturb(segment)
+        assert len(segment.samples) == original_len, "Shift perturbation must preserve audio length"
+
+    def test_shift_perturbation_zero_shift(self):
+        """When min and max shift are both 0, audio should be unchanged."""
+        perturb = ShiftPerturbation(min_shift_ms=0.0, max_shift_ms=0.0)
+        segment = self._make_audio_segment(duration_sec=0.5)
+        original = segment.samples.copy()
+        perturb.perturb(segment)
+        np.testing.assert_array_equal(segment.samples, original, "Zero shift should not modify audio")

tests/collections/asr/test_ssl_models.py

@@ -0,0 +1,446 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+from omegaconf import DictConfig
+
+from nemo.collections.asr.losses import ContrastiveLoss
+from nemo.collections.asr.models import EncDecDenoiseMaskedTokenPredModel, SpeechEncDecSelfSupervisedModel
+from nemo.core.classes.common import typecheck
+
+
+@pytest.fixture()
+def ssl_model():
+    preprocessor = {
+        'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor',
+        'params': dict({'pad_to': 16, 'dither': 0}),
+    }
+
+    model_defaults = {'enc_hidden': 32, 'dec_out': 128}
+
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': model_defaults['enc_hidden'],
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                },
+                {
+                    'filters': model_defaults['enc_hidden'],
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                },
+                {
+                    'filters': model_defaults['enc_hidden'],
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                },
+            ],
+        },
+    }
+
+    spec_augment = {
+        '_target_': 'nemo.collections.asr.modules.MaskedPatchAugmentation',
+        'freq_masks': 3,
+        'freq_width': 20,
+        'patch_size': 16,
+        'mask_patches': 0.5,
+    }
+
+    loss_list_contr_mlm = {
+        'contr': {
+            'decoder': {
+                '_target_': 'nemo.collections.asr.modules.ConvASRDecoderReconstruction',
+                'feat_in': model_defaults['enc_hidden'],
+                'feat_hidden': 128,
+                'feat_out': model_defaults['dec_out'],
+                'stride_layers': 0,
+                'non_stride_layers': 0,
+                'stride_transpose': False,
+            },
+            'loss': {
+                '_target_': 'nemo.collections.asr.losses.ContrastiveLoss',
+                'in_dim': 64,
+                'proj_dim': model_defaults['dec_out'],
+                'combine_time_steps': 1,
+                'quantized_targets': True,
+                'codebook_size': 64,
+                'sample_from_same_utterance_only': True,
+                'sample_from_non_masked': False,
+                'num_negatives': 3,
+            },
+        },
+        'mlm': {
+            'decoder': {
+                '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+                'feat_in': model_defaults['enc_hidden'],
+                'num_classes': 4096,
+            },
+            'loss': {'_target_': 'nemo.collections.asr.losses.MLMLoss', 'combine_time_steps': 1},
+            'targets_from_loss': "contr",
+        },
+    }
+
+    modelConfig_contr_mlm = DictConfig(
+        {
+            'preprocessor': DictConfig(preprocessor),
+            'spec_augment': DictConfig(spec_augment),
+            'model_defaults': DictConfig(model_defaults),
+            'encoder': DictConfig(encoder),
+            'loss_list': DictConfig(loss_list_contr_mlm),
+        }
+    )
+    ssl_model = SpeechEncDecSelfSupervisedModel(cfg=modelConfig_contr_mlm)
+    return ssl_model
+
+
+@pytest.fixture()
+def denoise_mlm_ssl_model():
+
+    model_defaults = {
+        "subsampling_factor": 1,
+        'enc_hidden': 32,
+        'dec_out': 128,
+        "sample_rate": 16000,
+        "num_classes": 32,
+        "num_books": 1,
+        "code_dim": 16,
+        "squeeze_single": False,
+        "mask_position": "pre_conv",  # position to apply masking, before or after conv subsampling, choices in ['pre_conv', 'post_conv']
+    }
+
+    preprocessor = {
+        "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+        "sample_rate": model_defaults["sample_rate"],
+        "normalize": "per_feature",
+        "window_size": 0.025,
+        "window_stride": 0.01,
+        "window": "hann",
+        "features": 80,
+        "n_fft": 512,
+        "log": True,
+        "frame_splicing": 1,
+        "dither": 0.00001,
+        "pad_to": 16,
+        "pad_value": 0.0,
+    }
+
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': preprocessor["features"],
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': model_defaults['enc_hidden'],
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                },
+                {
+                    'filters': model_defaults['enc_hidden'],
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                },
+                {
+                    'filters': model_defaults['enc_hidden'],
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                },
+            ],
+        },
+    }
+
+    spec_augment = {
+        '_target_': 'nemo.collections.asr.modules.SpectrogramAugmentation',
+        'freq_masks': 0,
+        'time_masks': 0,
+        'freq_width': 16,
+        'time_width': 0.05,
+    }
+
+    masking = {
+        "_target_": "nemo.collections.asr.modules.RandomBlockMasking",
+        "block_size": 40,  # for pre_conv masking, 10ms per frame, 400ms per block with block_size=40
+        "mask_prob": 0.01,  # for allow_overlap=True, this means the mask prob for each frame; otherwise it means the overall masked proportion
+        "feat_in": preprocessor["features"],
+        "freeze": True,
+        "allow_overlap": True,
+    }
+
+    quantizer = {
+        "_target_": "nemo.collections.asr.modules.RandomProjectionVectorQuantizer",
+        "feat_in": preprocessor["features"],
+        "code_dim": model_defaults["code_dim"],
+        "num_books": model_defaults["num_books"],
+        "num_classes": model_defaults["num_classes"],
+        "dist_fn": "l2",  # choices=["l2", "cosine"]
+        "freeze": True,
+        "squeeze_single": model_defaults["squeeze_single"],
+        "combine_time_steps": model_defaults["subsampling_factor"],  # conformer sub-sampling ratio
+    }
+
+    decoder = {
+        "_target_": "nemo.collections.asr.modules.MultiSoftmaxDecoder",
+        "feat_in": model_defaults["enc_hidden"],
+        "num_classes": model_defaults["num_classes"],
+        "num_decoders": model_defaults["num_books"],
+        "squeeze_single": model_defaults["squeeze_single"],
+        "use_bias": True,
+    }
+
+    loss = {
+        "_target_": "nemo.collections.asr.losses.MultiMLMLoss",
+        "combine_time_steps": model_defaults[
+            "subsampling_factor"
+        ],  # conformer sub-sampling ratio for 'pre_conv', 1 for 'post_conv'
+        "mask_threshold": 0.8,
+        "num_decoders": model_defaults["num_books"],
+        "squeeze_single": model_defaults["squeeze_single"],
+    }
+
+    optim = {
+        "name": "adamw",
+        "lr": 5.0,
+        # optimizer arguments
+        "betas": [0.9, 0.98],
+        "weight_decay": 1e-3,
+    }
+
+    model_config = DictConfig(
+        {
+            "preprocessor": DictConfig(preprocessor),
+            "spec_augment": DictConfig(spec_augment),
+            'model_defaults': DictConfig(model_defaults),
+            "masking": DictConfig(masking),
+            "quantizer": DictConfig(quantizer),
+            "encoder": DictConfig(encoder),
+            "decoder": DictConfig(decoder),
+            "loss": DictConfig(loss),
+            "optim": DictConfig(optim),
+        }
+    )
+    ssl_model = EncDecDenoiseMaskedTokenPredModel(cfg=model_config)
+    return ssl_model
+
+
+class TestSSLModel:
+    @pytest.mark.unit
+    def test_constructor(self, ssl_model):
+        confdict = ssl_model.to_config_dict()
+        instance2 = SpeechEncDecSelfSupervisedModel.from_config_dict(confdict)
+        assert isinstance(instance2, SpeechEncDecSelfSupervisedModel)
+
+    @pytest.mark.unit
+    def test_contr_nonquant(self, ssl_model):
+        modelConfig_contr_nonquant = ssl_model.to_config_dict()
+
+        loss_list_contr_nonquant = dict(modelConfig_contr_nonquant['loss_list'])
+        del loss_list_contr_nonquant['mlm']
+
+        loss_list_contr_nonquant['contr']['loss']['quantized_targets'] = False
+
+        modelConfig_contr_nonquant['loss_list'] = DictConfig(loss_list_contr_nonquant)
+
+        ssl_model = SpeechEncDecSelfSupervisedModel(cfg=modelConfig_contr_nonquant)
+
+        input_signal = torch.randn(size=(4, 64000))
+        length = torch.randint(low=48000, high=64000, size=[4])
+
+        with torch.no_grad():
+            spectrograms, spec_masks, encoded, encoded_len = ssl_model.forward(
+                input_signal=input_signal, input_signal_length=length
+            )
+
+            loss_value, loss_val_dict = ssl_model.decoder_loss_step(spectrograms, spec_masks, encoded, encoded_len)
+
+        assert len(loss_val_dict) == 1
+
+    @pytest.mark.unit
+    def test_contr_mlm(self, ssl_model):
+
+        input_signal = torch.randn(size=(4, 64000))
+        length = torch.randint(low=48000, high=64000, size=[4])
+
+        with torch.no_grad():
+            spectrograms, spec_masks, encoded, encoded_len = ssl_model.forward(
+                input_signal=input_signal, input_signal_length=length
+            )
+
+        loss_value, loss_val_dict = ssl_model.decoder_loss_step(spectrograms, spec_masks, encoded, encoded_len)
+
+        assert len(loss_val_dict) == 2
+
+    @pytest.mark.unit
+    def test_contr_mlm_multi(self, ssl_model):
+        modelConfig_contr_mlm_multi = ssl_model.to_config_dict()
+
+        model_defaults = modelConfig_contr_mlm_multi['model_defaults']
+
+        loss_list_contr_mlm_multi = dict(modelConfig_contr_mlm_multi['loss_list'])
+        loss_list_contr_mlm_multi['mlm_2'] = {
+            'decoder': {
+                '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+                'feat_in': model_defaults['enc_hidden'],
+                'num_classes': 4096,
+            },
+            'loss': {'_target_': 'nemo.collections.asr.losses.MLMLoss', 'combine_time_steps': 1},
+            'output_from_layer': "encoder.0",
+            'targets_from_loss': "contr",
+        }
+        loss_list_contr_mlm_multi['mlm_3'] = {
+            'decoder': {
+                '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
+                'feat_in': model_defaults['enc_hidden'],
+                'num_classes': 4096,
+            },
+            'loss': {'_target_': 'nemo.collections.asr.losses.MLMLoss', 'combine_time_steps': 1},
+            'output_from_layer': "encoder.1",
+            'targets_from_loss': "contr",
+        }
+        modelConfig_contr_mlm_multi['loss_list'] = DictConfig(loss_list_contr_mlm_multi)
+
+        ssl_model = SpeechEncDecSelfSupervisedModel(cfg=modelConfig_contr_mlm_multi)
+
+        input_signal = torch.randn(size=(4, 64000))
+        length = torch.randint(low=48000, high=64000, size=[4])
+
+        with torch.no_grad():
+            spectrograms, spec_masks, encoded, encoded_len = ssl_model.forward(
+                input_signal=input_signal, input_signal_length=length
+            )
+
+            loss_value, loss_val_dict = ssl_model.decoder_loss_step(spectrograms, spec_masks, encoded, encoded_len)
+
+        assert len(loss_val_dict) == 4
+
+
+class TestContrastiveLoss:
+    @pytest.mark.unit
+    def test_sample_negatives_fewer_frames_than_num_negatives(self):
+        num_negatives = 40
+        num_frames = 5
+        num = num_frames
+        feat_dim = 128
+
+        loss = ContrastiveLoss(in_dim=64, proj_dim=feat_dim, num_negatives=num_negatives, quantized_targets=False)
+        y = torch.randn(num_frames, feat_dim)
+
+        negs, neg_idxs = loss.sample_negatives(y, num)
+
+        assert neg_idxs.shape == (num, num_negatives)
+        assert negs.shape == (num_negatives, num, feat_dim)
+
+
+class TestDenoiseMLMSSLModel:
+    @pytest.mark.unit
+    def test_forward(self, denoise_mlm_ssl_model):
+        input_signal = torch.randn(size=(4, 64000))
+        input_length = torch.randint(low=48000, high=64000, size=[4])
+        noise = 0.1 * torch.ones_like(input_signal)
+        noisy_input_signal = input_signal + noise
+        noisy_input_length = input_length
+        with torch.no_grad():
+            with typecheck.disable_checks():
+                log_probs, encoded_len, masks, tokens = denoise_mlm_ssl_model.forward(
+                    input_signal=input_signal,
+                    input_signal_length=input_length,
+                    noisy_input_signal=noisy_input_signal,
+                    noisy_input_signal_length=noisy_input_length,
+                )
+
+        assert log_probs.size(0) == 4
+        assert log_probs.size(2) == denoise_mlm_ssl_model.cfg.model_defaults.num_classes
+        assert encoded_len.size(0) == 4
+        assert masks.size(0) == 4
+        assert tokens.size(0) == 4
+        assert masks.sum() == 0.0  # no mask should be applied to the input by default
+
+    @pytest.mark.unit
+    def test_forward_masked(self, denoise_mlm_ssl_model: EncDecDenoiseMaskedTokenPredModel):
+        input_signal = torch.randn(size=(4, 64000))
+        input_length = torch.randint(low=48000, high=64000, size=[4])
+        noise = 0.1 * torch.ones_like(input_signal)
+        noisy_input_signal = input_signal + noise
+        noisy_input_length = input_length
+
+        with torch.no_grad():
+            with typecheck.disable_checks():
+                log_probs, encoded_len, masks, tokens = denoise_mlm_ssl_model.forward(
+                    input_signal=input_signal,
+                    input_signal_length=input_length,
+                    noisy_input_signal=noisy_input_signal,
+                    noisy_input_signal_length=noisy_input_length,
+                    apply_mask=True,
+                )
+
+        loss_value = denoise_mlm_ssl_model.loss(
+            masks=masks, decoder_outputs=log_probs, targets=tokens, decoder_lengths=encoded_len
+        )
+
+        assert log_probs.size(0) == 4
+        assert log_probs.size(2) == denoise_mlm_ssl_model.cfg.model_defaults.num_classes
+        assert encoded_len.size(0) == 4
+        assert masks.size(0) == 4
+        assert tokens.size(0) == 4
+        assert masks.sum() > 0.0  # mask should be applied to the input
+        assert not torch.isnan(loss_value)

tests/collections/asr/test_text_to_text_dataset.py

@@ -0,0 +1,213 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import multiprocessing
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import pytest
+from omegaconf import OmegaConf
+
+from nemo.core.classes.common import safe_instantiate
+
+nemo_text_processing = pytest.importorskip("nemo_text_processing", reason="Requires nemo_text_processing to run")
+
+try:
+    from nemo_text_processing.text_normalization.normalize import Normalizer
+except (ImportError, ModuleNotFoundError):
+    raise ModuleNotFoundError(
+        "The package `nemo_text_processing` was not installed in this environment. Please refer to"
+        " https://github.com/NVIDIA/NeMo-text-processing and install this package before using "
+        "this script"
+    )
+
+from nemo.collections.asr.data.text_to_text import TextToTextDataset, TextToTextItem, TextToTextIterableDataset
+from nemo.collections.common import tokenizers
+
+BASE_DIR = Path(__file__).parent.parent.parent.parent
+
+
+@pytest.fixture(scope="module")
+def set_multiprocessing_method():
+    """
+    Try to set 'fork' multiprocessing method to avoid problems with multiprocessing in PyTest on MacOS
+    """
+    if multiprocessing.get_start_method(allow_none=True) != "fork":
+        multiprocessing.set_start_method("fork", force=True)
+
+
+@pytest.fixture(scope="module")
+def speakers_path(tmp_path_factory):
+    path = tmp_path_factory.mktemp("textonly") / "speakers.txt"
+    with open(path, "w", encoding="utf-8") as f:
+        for speaker in [1, 2, 3]:
+            print(f"{speaker}", file=f)
+    return path
+
+
+@pytest.fixture(scope="module")
+def textonly_manifest_path(tmp_path_factory):
+    path = tmp_path_factory.mktemp("textonly") / "manifest.json"
+    texts = [
+        "lorem ipsum dolor sit amet consectetur adipiscing elit",
+        "nullam rhoncus sapien eros eu mollis sem euismod non",
+    ]
+    with open(path, "w", encoding="utf-8") as f:
+        for text in texts:
+            print(json.dumps(dict(text=text, tts_text_normalized=text)), file=f)
+    return path
+
+
+@pytest.fixture(scope="module")
+def textonly_unnormalized_manifest_path(tmp_path_factory):
+    path = tmp_path_factory.mktemp("textonly") / "manifest_nonorm.json"
+    texts = [
+        (
+            "lorem ipsum dolor sit amet consectetur adipiscing elit",
+            "Lorem ipsum dolor sit amet, consectetur adipiscing elit.",
+        ),
+        (
+            "nullam rhoncus sapien eros eu mollis sem euismod non nineteen",
+            "Nullam rhoncus sapien eros, eu mollis sem euismod non 19.",
+        ),
+    ]
+    with open(path, "w", encoding="utf-8") as f:
+        for asr_text, tts_text in texts:
+            print(json.dumps(dict(text=asr_text, tts_text=tts_text)), file=f)
+    return path
+
+
+@pytest.fixture(scope="module")
+def tts_normalizer():
+    normalizer = Normalizer(
+        lang="en",
+        input_case="cased",
+        overwrite_cache=True,
+        cache_dir=None,
+    )
+    return normalizer
+
+
+@pytest.fixture(scope="module")
+def asr_tokenizer(test_data_dir):
+    tokenizer_path = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
+    tokenizer = tokenizers.AutoTokenizer(pretrained_model_name='bert-base-cased', vocab_file=tokenizer_path)
+    return tokenizer
+
+
+@pytest.fixture(scope="module")
+def tts_tokenizer():
+    @dataclass
+    class G2PConfig:
+        _target_: str = "nemo.collections.tts.g2p.models.en_us_arpabet.EnglishG2p"
+        phoneme_dict: str = str(BASE_DIR / "scripts/tts_dataset_files/cmudict-0.7b_nv22.10")
+        heteronyms: str = str(BASE_DIR / "scripts/tts_dataset_files/heteronyms-052722")
+        phoneme_probability: float = 0.5
+
+    @dataclass
+    class TextTokenizerCfg:
+        _target_: str = "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.EnglishPhonemesTokenizer"
+        punct: bool = True
+        stresses: bool = True
+        chars: bool = True
+        apostrophe: bool = True
+        pad_with_space: bool = True
+        add_blank_at: bool = True
+        g2p: G2PConfig = field(default_factory=lambda: G2PConfig())
+
+    config = OmegaConf.create(OmegaConf.to_yaml(TextTokenizerCfg()))
+    return safe_instantiate(config)
+
+
+class TestTextToTextDataset:
+    @pytest.mark.unit
+    @pytest.mark.parametrize("tokenizer_workers", [1, 2])
+    def test_text_to_text_dataset(
+        self,
+        textonly_manifest_path,
+        tokenizer_workers,
+        speakers_path,
+        asr_tokenizer,
+        tts_tokenizer,
+        tts_normalizer,
+        set_multiprocessing_method,
+    ):
+        """
+        Test map-style text-to-text dataset with ASR and TTS tokenizers with normalized text
+        """
+        dataset = TextToTextDataset(
+            manifest_filepath=textonly_manifest_path,
+            speakers_filepath=speakers_path,
+            asr_tokenizer=asr_tokenizer,
+            asr_use_start_end_token=False,
+            tts_parser=tts_tokenizer,
+            tts_text_pad_id=0,
+            tts_text_normalizer=tts_normalizer,
+            tts_text_normalizer_call_kwargs=dict(),
+            tokenizer_workers=tokenizer_workers,
+        )
+        assert len(dataset) == 2
+        item = dataset[0]
+        assert isinstance(item, TextToTextItem)
+
+    @pytest.mark.unit
+    def test_text_to_text_dataset_unnormalized(
+        self, textonly_unnormalized_manifest_path, speakers_path, asr_tokenizer, tts_tokenizer, tts_normalizer
+    ):
+        """
+        Test TextToTextDataset with ASR and TTS tokenizers with non-normalized text
+        """
+        dataset = TextToTextDataset(
+            manifest_filepath=textonly_unnormalized_manifest_path,
+            speakers_filepath=speakers_path,
+            asr_tokenizer=asr_tokenizer,
+            asr_use_start_end_token=False,
+            tts_parser=tts_tokenizer,
+            tts_text_pad_id=0,
+            tts_text_normalizer=tts_normalizer,
+            tts_text_normalizer_call_kwargs=dict(),
+        )
+        assert len(dataset) == 2
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("tokenizer_workers", [1, 2])
+    def test_text_to_text_iterable_dataset(
+        self,
+        textonly_manifest_path,
+        tokenizer_workers,
+        speakers_path,
+        asr_tokenizer,
+        tts_tokenizer,
+        tts_normalizer,
+        set_multiprocessing_method,
+    ):
+        """
+        Test iterable text-to-text dataset with ASR and TTS tokenizers with normalized text
+        """
+        dataset = TextToTextIterableDataset(
+            manifest_filepath=textonly_manifest_path,
+            speakers_filepath=speakers_path,
+            asr_tokenizer=asr_tokenizer,
+            asr_use_start_end_token=False,
+            tts_parser=tts_tokenizer,
+            tts_text_pad_id=0,
+            tts_text_normalizer=tts_normalizer,
+            tts_text_normalizer_call_kwargs=dict(),
+            tokenizer_workers=tokenizer_workers,
+        )
+        assert len(dataset) == 2
+        item = next(iter(dataset))
+        assert isinstance(item, TextToTextItem)

tests/collections/asr/test_transformer_encoder.py

@@ -0,0 +1,647 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pytest
+import torch
+
+from nemo.collections.asr.modules.transformer_encoder import (
+    FeatureStacking,
+    TransformerEncoder,
+    TransformerEncoderConfig,
+)
+from nemo.collections.asr.parts.submodules.multi_head_attention import RotaryPositionalEncoding
+
+
+class TestTransformerEncoderConfig:
+    @pytest.mark.unit
+    def test_default_config(self):
+        cfg = TransformerEncoderConfig()
+        assert cfg.feat_in == 128
+        assert cfg.d_model == 512
+        assert cfg.n_heads == 8
+        assert cfg.n_layers == 17
+        assert cfg.drop_rate == 0.1
+        assert cfg.qkv_bias is False
+        assert cfg.qk_norm is False
+        assert cfg.ff_expansion == 4.0
+        assert cfg.pre_block_norm is True
+        assert cfg.subsampling_factor == 4
+        assert cfg.attn_mode == "full"
+        assert cfg.self_attention_model == "rel_pos"
+        assert cfg.rope_base == 10000.0
+        assert cfg.rotary_fraction == 1.0
+
+    @pytest.mark.unit
+    def test_custom_config(self):
+        cfg = TransformerEncoderConfig(
+            feat_in=128, d_model=1280, n_heads=16, n_layers=32, qk_norm=True, self_attention_model="abs_pos"
+        )
+        assert cfg.feat_in == 128
+        assert cfg.d_model == 1280
+        assert cfg.n_heads == 16
+        assert cfg.n_layers == 32
+        assert cfg.qk_norm is True
+        assert cfg.self_attention_model == "abs_pos"
+
+
+class TestFeatureStacking:
+    @pytest.mark.unit
+    @pytest.mark.parametrize("subsampling_factor", [2, 4, 8])
+    def test_output_shape(self, subsampling_factor):
+        B, C, T = 2, 80, 400
+        stacking = FeatureStacking(subsampling_factor=subsampling_factor, feat_in=C, feat_out=256)
+        x = torch.randn(B, C, T)
+        lengths = torch.tensor([400, 300])
+
+        out, out_lengths = stacking(x, lengths)
+        expected_t = stacking.compute_num_out_frames(T)
+        assert out.shape == (B, expected_t, 256)
+        assert out_lengths[0].item() == expected_t
+
+    @pytest.mark.unit
+    def test_padding_when_not_divisible(self):
+        B, C, T = 1, 80, 401
+        subsampling_factor = 4
+        stacking = FeatureStacking(subsampling_factor=subsampling_factor, feat_in=C, feat_out=256)
+        x = torch.randn(B, C, T)
+        lengths = torch.tensor([401])
+
+        out, out_lengths = stacking(x, lengths)
+        expected_t = stacking.compute_num_out_frames(T)
+        assert out.shape == (B, expected_t, 256)
+        assert out_lengths[0].item() == expected_t
+
+    @pytest.mark.unit
+    def test_length_shorter_than_batch(self):
+        """Output length must be ceil(sample_length / factor), not dependent on batch T."""
+        B, C, T = 2, 80, 403
+        subsampling_factor = 4
+        stacking = FeatureStacking(subsampling_factor=subsampling_factor, feat_in=C, feat_out=256)
+        x = torch.randn(B, C, T)
+        lengths = torch.tensor([401, 397])
+
+        _, out_lengths = stacking(x, lengths)
+        assert out_lengths[0].item() == stacking.compute_num_out_frames(401)
+        assert out_lengths[1].item() == stacking.compute_num_out_frames(397)
+
+    @pytest.mark.unit
+    def test_no_padding_when_divisible(self):
+        B, C, T = 1, 80, 400
+        stacking = FeatureStacking(subsampling_factor=4, feat_in=C, feat_out=256)
+        x = torch.randn(B, C, T)
+        lengths = torch.tensor([400])
+
+        out, out_lengths = stacking(x, lengths)
+        assert out.shape == (B, stacking.compute_num_out_frames(T), 256)
+        assert out_lengths[0].item() == stacking.compute_num_out_frames(T)
+
+
+class TestBypassPreEncode:
+    """Testing bypass pre-encode functionality."""
+
+    def test_bypass_pre_encode_forward(self):
+        """Testing that forward works with "bypass pre-encode" mode.
+
+        Forwards are wrapped in ``torch.no_grad()`` so the test runs on CPU as well as GPU:
+        FlexAttention's CPU path refuses to run when any input requires gradients (parameters
+        of an ``nn.Module`` do by default), and we are only checking output shapes here, never
+        calling ``.backward()``.
+        """
+        # For pre-encoded embeddings, the shape is (batch_size, n_frames, emb_dim)
+        batch_size = 2
+        n_frames, emb_dim, feat_out = 17, 64, 8  # emb_dim=64 with n_heads=4 -> head_dim=16 (>= 16)
+        random_input = torch.rand((batch_size, n_frames, emb_dim))
+        random_length = torch.tensor([n_frames] * batch_size, dtype=torch.int64)
+
+        model = TransformerEncoder(
+            feat_in=10,
+            n_layers=3,
+            d_model=emb_dim,
+            n_heads=4,
+            feat_out=feat_out,
+            drop_rate=0.0,
+            dropout_pre_encoder=0.0,
+            dropout_emb=0.0,
+        )
+        model.train()
+        with torch.no_grad():
+            fwd_outputs = model(audio_signal=random_input, length=random_length, bypass_pre_encode=True)[0]
+        assert fwd_outputs.shape == (batch_size, feat_out, n_frames)
+
+        model.eval()
+        with torch.no_grad():
+            fwd_outputs = model(audio_signal=random_input, length=random_length, bypass_pre_encode=True)[0]
+        assert fwd_outputs.shape == (batch_size, feat_out, n_frames)
+
+    def test_error_shape_invalid_bypass_pre_encode_forward(self):
+        """
+        Testing that error messages are correctly triggered regarding "bypass pre-encode" mode.
+        Both correct samples and wrongs samples are tested.
+
+        (1) bypass_pre_encode = False (default):
+            `audio_signal` must be a tensor containing audio features.
+            Shape: (batch, self._feat_in, n_frames)
+        (2) bypass_pre_encode = True:
+            `audio_signal` must be a tensor containing pre-encoded embeddings.
+            Shape: (batch, n_frame, self.d_model)
+        """
+        batch_size = 2
+        n_frames, emb_dim, feat_in, feat_out = 17, 64, 10, 8  # emb_dim=64 with n_heads=4 -> head_dim=16 (>= 16)
+
+        pre_encode_input = torch.rand((batch_size, n_frames, emb_dim))
+        feat_input = torch.rand((batch_size, feat_in, n_frames))
+        input_length = torch.tensor([n_frames] * batch_size, dtype=torch.int64)
+
+        model = TransformerEncoder(
+            feat_in=feat_in,
+            n_layers=3,
+            d_model=emb_dim,
+            n_heads=4,
+            feat_out=feat_out,
+            drop_rate=0.0,
+            dropout_pre_encoder=0.0,
+            dropout_emb=0.0,
+        )
+        sub_sampled_n_frames = np.ceil(n_frames / model.subsampling_factor)
+
+        # Test with bypass_pre_encode = True, should be pre_encode_input but given feat_input.
+        model.train()
+        with pytest.raises(ValueError):
+            model(audio_signal=feat_input, length=input_length, bypass_pre_encode=True)
+
+        model.eval()
+        with pytest.raises(ValueError):
+            model(audio_signal=feat_input, length=input_length, bypass_pre_encode=True)
+
+        # Test with bypass_pre_encode = True, given the correct input pre_encode_input.
+        # NB: forwards that actually reach FlexAttention are wrapped in ``torch.no_grad()`` so
+        # the test passes on CPU (FlexAttention's CPU path refuses inputs that require grad).
+        # The ``pytest.raises(ValueError)`` blocks above/below intentionally do *not* need this
+        # wrapper because the shape check in ``TransformerEncoder.forward()`` raises before any
+        # attention computation.
+        model.train()
+        with torch.no_grad():
+            fwd_outputs = model(audio_signal=pre_encode_input, length=input_length, bypass_pre_encode=True)[0]
+        assert fwd_outputs.shape == (batch_size, feat_out, n_frames)
+
+        model.eval()
+        with torch.no_grad():
+            fwd_outputs = model(audio_signal=pre_encode_input, length=input_length, bypass_pre_encode=True)[0]
+        assert fwd_outputs.shape == (batch_size, feat_out, n_frames)
+
+        # Test with bypass_pre_encode = False, should be feat_input but given pre_encode_input.
+        model.train()
+        with pytest.raises(ValueError):
+            model(audio_signal=pre_encode_input, length=input_length, bypass_pre_encode=False)
+
+        model.eval()
+        with pytest.raises(ValueError):
+            model(audio_signal=pre_encode_input, length=input_length, bypass_pre_encode=False)
+
+        # Test with bypass_pre_encode = False, given the correct input feat_input.
+        model.train()
+        with torch.no_grad():
+            fwd_outputs = model(audio_signal=feat_input, length=input_length, bypass_pre_encode=False)[0]
+        assert fwd_outputs.shape == (batch_size, feat_out, sub_sampled_n_frames)
+
+        model.eval()
+        with torch.no_grad():
+            fwd_outputs = model(audio_signal=feat_input, length=input_length, bypass_pre_encode=False)[0]
+        assert fwd_outputs.shape == (batch_size, feat_out, sub_sampled_n_frames)
+
+    @pytest.mark.unit
+    def test_bypass_pre_encode_matches_manual_pre_encode(self):
+        """``bypass_pre_encode=True`` must skip *only* the pre-encoder.
+
+        Running the pre-encoder by hand and feeding its output back in with
+        ``bypass_pre_encode=True`` should reproduce the full forward
+        (``bypass_pre_encode=False``) exactly, because the positional-encoding, norm and
+        Transformer-block stack downstream of the pre-encoder is identical on both paths.
+        """
+        B, feat_in, T, d_model, feat_out = 2, 32, 64, 64, 8  # d_model=64 with n_heads=4 -> head_dim=16 (>= 16)
+        model = TransformerEncoder(
+            feat_in=feat_in,
+            d_model=d_model,
+            n_heads=4,
+            n_layers=2,
+            feat_out=feat_out,
+            subsampling_factor=4,
+            drop_rate=0.0,
+            dropout_pre_encoder=0.0,
+            dropout_emb=0.0,
+        )
+        model.eval()
+
+        mel = torch.randn(B, feat_in, T)
+        lengths = torch.tensor([T, T - 8], dtype=torch.int64)
+
+        with torch.no_grad():
+            out_full, len_full = model(audio_signal=mel, length=lengths, bypass_pre_encode=False)
+
+            # Reproduce just the pre-encoder, then bypass it on the next call.
+            pre_x, pre_len = model.pre_encode(mel, lengths)
+            out_bypass, len_bypass = model(audio_signal=pre_x, length=pre_len, bypass_pre_encode=True)
+
+        assert out_full.shape == out_bypass.shape == (B, feat_out, pre_x.shape[1])
+        assert torch.equal(len_full, len_bypass)
+        assert torch.allclose(out_full, out_bypass, atol=1e-5)
+
+
+class TestTransformerEncoder:
+    @pytest.mark.unit
+    def test_model_creation(self):
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2)
+        total_params = sum(p.numel() for p in model.parameters())
+        assert total_params > 0
+        assert len(model.layers) == 2
+
+    @pytest.mark.unit
+    def test_model_creation_with_qk_norm(self):
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, qk_norm=True)
+        attn = model.layers[0].attn
+        assert hasattr(attn, 'q_norm')
+        assert hasattr(attn, 'k_norm')
+
+    @pytest.mark.unit
+    def test_model_creation_without_qk_norm(self):
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, qk_norm=False)
+        attn = model.layers[0].attn
+        assert not hasattr(attn, 'q_norm')
+        assert not hasattr(attn, 'k_norm')
+
+    @pytest.mark.unit
+    def test_invalid_attn_mode(self):
+        with pytest.raises(ValueError, match="not yet supported"):
+            TransformerEncoder(feat_in=80, d_model=64, n_heads=4, n_layers=2, attn_mode="sliding_window")
+
+    @pytest.mark.unit
+    def test_head_dim_below_16_raises(self):
+        """head_dim = d_model // n_heads must be >= 16 (PyTorch FlexAttention CUDA requirement).
+
+        The check happens at construction time, so an unsupported (d_model, n_heads) pair raises
+        before any forward pass.
+        """
+        # d_model=32, n_heads=4 -> head_dim=8 (< 16).
+        with pytest.raises(ValueError, match="per-head embedding dimension >= 16"):
+            TransformerEncoder(feat_in=128, d_model=32, n_heads=4, n_layers=2)
+
+    @pytest.mark.unit
+    def test_causal_forward_cpu(self):
+        model = TransformerEncoder(feat_in=80, d_model=64, n_heads=4, n_layers=2, drop_rate=0.0, attn_mode="causal")
+        model.eval()
+
+        x = torch.randn(2, 80, 400)
+        lengths = torch.tensor([400, 300])
+
+        with torch.no_grad():
+            out, out_lengths = model(x, lengths)
+
+        assert out.shape == (2, 64, 100)
+        assert out_lengths.tolist() == [100, 75]
+        assert not torch.isnan(out).any()
+
+    @pytest.mark.unit
+    def test_causal_future_does_not_affect_past(self):
+        """Output at position t must be invariant to changes at positions > t."""
+        model = TransformerEncoder(feat_in=80, d_model=64, n_heads=4, n_layers=2, drop_rate=0.0, attn_mode="causal")
+        model.eval()
+
+        B, C, T = 1, 80, 400
+        x_a = torch.randn(B, C, T)
+        x_b = x_a.clone()
+        # Perturb only the second half of frames.
+        x_b[:, :, T // 2 :] = torch.randn(B, C, T - T // 2)
+        lengths = torch.tensor([T])
+
+        with torch.no_grad():
+            out_a, _ = model(x_a, lengths)
+            out_b, _ = model(x_b, lengths)
+
+        # Output frames covering only past + present should be identical.
+        # First half of *output* frames corresponds to first half of input frames after subsampling.
+        safe_t = (T // 2) // model.pre_encode.subsampling_factor
+        assert torch.allclose(out_a[:, :, :safe_t], out_b[:, :, :safe_t], atol=1e-5)
+
+    @pytest.mark.unit
+    def test_freeze_unfreeze_partial_restores_prior_state(self):
+        model = TransformerEncoder(feat_in=80, d_model=64, n_heads=4, n_layers=2)
+        for p in model.final_norm.parameters():
+            p.requires_grad = False
+        prior = {n: p.requires_grad for n, p in model.named_parameters()}
+
+        model.freeze()
+        assert all(not p.requires_grad for p in model.parameters())
+        assert not model.training
+
+        model.unfreeze(partial=True)
+        assert {n: p.requires_grad for n, p in model.named_parameters()} == prior
+        assert model.training
+
+    @pytest.mark.unit
+    def test_forward_cpu(self):
+        """Forward pass on CPU uses unfused FlexAttention fallback."""
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, drop_rate=0.0, subsampling_factor=4)
+        model.eval()
+
+        B, C, T = 2, 128, 400
+        x = torch.randn(B, C, T)
+        lengths = torch.tensor([400, 300])
+
+        with torch.no_grad():
+            out, out_lengths = model(audio_signal=x, length=lengths)
+
+        assert out.shape == (B, 64, T // 4)
+        assert out_lengths[0].item() == T // 4
+        assert out_lengths[1].item() == 300 // 4
+        assert not torch.isnan(out).any()
+
+    @pytest.mark.unit
+    def test_forward_cpu_with_qk_norm(self):
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, drop_rate=0.0, qk_norm=True)
+        model.eval()
+
+        x = torch.randn(1, 128, 200)
+        lengths = torch.tensor([200])
+
+        with torch.no_grad():
+            out, _ = model(audio_signal=x, length=lengths)
+
+        assert out.shape == (1, 64, 50)
+        assert not torch.isnan(out).any()
+
+    @pytest.mark.run_only_on('GPU')
+    def test_forward_basic(self):
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, drop_rate=0.0, subsampling_factor=4)
+        model = model.cuda().to(torch.bfloat16)
+
+        B, C, T = 2, 128, 400
+        x = torch.randn(B, C, T, device='cuda', dtype=torch.bfloat16)
+        lengths = torch.tensor([400, 300], device='cuda')
+
+        model.eval()
+        with torch.no_grad():
+            out, out_lengths = model(audio_signal=x, length=lengths)
+
+        assert out.shape == (B, 64, T // 4)
+        assert out_lengths[0].item() == T // 4
+        assert out_lengths[1].item() == 300 // 4
+        assert not torch.isnan(out).any()
+
+    @pytest.mark.run_only_on('GPU')
+    def test_forward_with_qk_norm(self):
+        model = TransformerEncoder(
+            feat_in=128, d_model=128, n_heads=8, n_layers=2, drop_rate=0.0, qk_norm=True, subsampling_factor=8
+        )
+        model = model.cuda().to(torch.bfloat16)
+
+        B, C, T = 2, 128, 800
+        x = torch.randn(B, C, T, device='cuda', dtype=torch.bfloat16)
+        lengths = torch.tensor([800, 640], device='cuda')
+
+        model.eval()
+        with torch.no_grad():
+            out, out_lengths = model(audio_signal=x, length=lengths)
+
+        assert out.shape == (B, 128, T // 8)
+        assert out_lengths[1].item() == 640 // 8
+        assert not torch.isnan(out).any()
+
+    @pytest.mark.run_only_on('GPU')
+    def test_forward_output_channels_first(self):
+        """Verify output is (B, D, T) channels-first as expected by downstream decoders."""
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=1, drop_rate=0.0)
+        model = model.cuda().to(torch.bfloat16)
+
+        x = torch.randn(1, 128, 200, device='cuda', dtype=torch.bfloat16)
+        lengths = torch.tensor([200], device='cuda')
+
+        model.eval()
+        with torch.no_grad():
+            out, _ = model(audio_signal=x, length=lengths)
+
+        assert out.shape[1] == 64  # D dimension
+        assert out.shape[2] == 200 // 4  # T dimension
+
+    @pytest.mark.run_only_on('GPU')
+    def test_eval_deterministic(self):
+        """In eval mode with no dropout, repeated forward passes should produce identical output."""
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, drop_rate=0.0)
+        model = model.cuda().to(torch.bfloat16).eval()
+
+        x = torch.randn(1, 128, 200, device='cuda', dtype=torch.bfloat16)
+        lengths = torch.tensor([200], device='cuda')
+
+        with torch.no_grad():
+            out1, _ = model(audio_signal=x, length=lengths)
+            out2, _ = model(audio_signal=x, length=lengths)
+
+        assert torch.allclose(out1, out2, atol=1e-6)
+
+    @pytest.mark.run_only_on('GPU')
+    def test_padding_does_not_affect_valid_output(self):
+        """Padding frames should not change the encoded output at valid positions."""
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, drop_rate=0.0)
+        model = model.cuda().to(torch.bfloat16).eval()
+
+        T_valid = 200
+        x_short = torch.randn(1, 128, T_valid, device='cuda', dtype=torch.bfloat16)
+        lengths_short = torch.tensor([T_valid], device='cuda')
+
+        T_padded = 400
+        x_long = torch.zeros(1, 128, T_padded, device='cuda', dtype=torch.bfloat16)
+        x_long[:, :, :T_valid] = x_short
+        lengths_long = torch.tensor([T_valid], device='cuda')
+
+        with torch.no_grad():
+            out_short, len_short = model(audio_signal=x_short, length=lengths_short)
+            out_long, len_long = model(audio_signal=x_long, length=lengths_long)
+
+        assert len_short[0].item() == len_long[0].item()
+        valid_t = len_short[0].item()
+        # bf16 + different block mask shapes cause small numerical differences in Triton kernels
+        assert torch.allclose(out_short[:, :, :valid_t], out_long[:, :, :valid_t], atol=5e-2)
+
+    @pytest.mark.run_only_on('GPU')
+    def test_backward_pass(self):
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, drop_rate=0.0)
+        model = model.cuda().to(torch.bfloat16).train()
+
+        x = torch.randn(2, 128, 200, device='cuda', dtype=torch.bfloat16)
+        lengths = torch.tensor([200, 160], device='cuda')
+
+        out, _ = model(audio_signal=x, length=lengths)
+        loss = out.sum()
+        loss.backward()
+
+        for name, param in model.named_parameters():
+            assert param.grad is not None, f"No gradient for {name}"
+            assert not torch.isnan(param.grad).any(), f"NaN gradient for {name}"
+
+
+class TestSelfAttentionModel:
+    """Tests for the ``self_attention_model`` positional encoding option."""
+
+    @pytest.mark.unit
+    def test_default_is_rel_pos(self):
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2)
+        assert model.self_attention_model == "rel_pos"
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("mode", ["abs_pos", "rel_pos", "no_pos", "rope"])
+    def test_valid_modes_are_accepted(self, mode):
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, self_attention_model=mode)
+        assert model.self_attention_model == mode
+
+    @pytest.mark.unit
+    def test_none_aliases_no_pos(self):
+        """Passing ``self_attention_model=None`` must be equivalent to ``"no_pos"``."""
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, self_attention_model=None)
+        assert model.self_attention_model == "no_pos"
+        assert model.pos_enc is None
+
+    @pytest.mark.unit
+    def test_invalid_mode_raises(self):
+        with pytest.raises(ValueError, match="not supported"):
+            TransformerEncoder(
+                feat_in=128, d_model=64, n_heads=4, n_layers=2, self_attention_model="rel_pos_local_attn"
+            )
+
+    @pytest.mark.unit
+    def test_rel_pos_attention_params_allocated(self):
+        """rel_pos mode allocates the Transformer-XL bias parameters per attention layer."""
+        d_model, n_heads, n_layers = 64, 4, 2
+        model = TransformerEncoder(
+            feat_in=128, d_model=d_model, n_heads=n_heads, n_layers=n_layers, self_attention_model="rel_pos"
+        )
+        head_dim = d_model // n_heads
+        assert model.pos_enc is not None
+        for layer in model.layers:
+            attn = layer.attn
+            assert attn.linear_pos is not None
+            assert attn.pos_bias_u is not None
+            assert attn.pos_bias_v is not None
+            assert attn.pos_bias_u.shape == (n_heads, head_dim)
+            assert attn.pos_bias_v.shape == (n_heads, head_dim)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("mode", ["abs_pos", "no_pos", "rope"])
+    def test_non_rel_pos_modes_have_no_rel_params(self, mode):
+        """abs_pos, no_pos and rope modes must not allocate the rel-pos parameters."""
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, self_attention_model=mode)
+        for layer in model.layers:
+            attn = layer.attn
+            assert attn.linear_pos is None
+            assert attn.pos_bias_u is None
+            assert attn.pos_bias_v is None
+
+    @pytest.mark.unit
+    def test_no_pos_has_no_positional_encoding_module(self):
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, self_attention_model="no_pos")
+        assert model.pos_enc is None
+        # set_max_audio_length is invoked in __init__; it must not crash for no_pos and must
+        # still record the requested max length so update_max_seq_length works normally.
+        assert model.max_audio_length == model.pos_emb_max_len
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("mode", ["abs_pos", "rel_pos", "no_pos", "rope", None])
+    def test_forward_each_mode_cpu(self, mode):
+        """Each ``self_attention_model`` choice (including ``None``) must produce a valid forward."""
+        model = TransformerEncoder(
+            feat_in=128,
+            d_model=64,
+            n_heads=4,
+            n_layers=2,
+            drop_rate=0.0,
+            subsampling_factor=4,
+            self_attention_model=mode,
+        )
+        model.eval()
+
+        B, C, T = 2, 128, 200
+        x = torch.randn(B, C, T)
+        lengths = torch.tensor([T, 160])
+
+        with torch.no_grad():
+            out, out_lengths = model(audio_signal=x, length=lengths)
+
+        assert out.shape == (B, 64, T // 4)
+        assert out_lengths[0].item() == T // 4
+        assert out_lengths[1].item() == 160 // 4
+        assert not torch.isnan(out).any()
+
+    @pytest.mark.unit
+    def test_rel_pos_broadcasts_when_T_differs_from_n_heads(self):
+        """Regression test for the Transformer-XL bias broadcasting.
+
+        ``pos_bias_{u,v}`` has shape ``(H, D)`` and must broadcast against the head axis of
+        ``q`` which has shape ``(B, H, T, D)``. A naive add would right-align ``H`` against
+        ``T`` and either crash (``T != H``) or silently apply the bias on the wrong axis
+        (``T == H``). This test exercises a configuration where ``T_attn != n_heads`` so the
+        broken broadcast would surface as an error.
+        """
+        # 200 input frames / subsampling_factor=4 -> 50 attention frames; n_heads=4 -> T != H.
+        model = TransformerEncoder(
+            feat_in=128, d_model=64, n_heads=4, n_layers=2, drop_rate=0.0, self_attention_model="rel_pos"
+        )
+        model.eval()
+
+        B, C, T = 2, 128, 200
+        x = torch.randn(B, C, T)
+        lengths = torch.tensor([T, 160])
+
+        with torch.no_grad():
+            out, _ = model(audio_signal=x, length=lengths)
+
+        assert out.shape == (B, 64, T // 4)
+        assert not torch.isnan(out).any()
+
+    @pytest.mark.unit
+    def test_rope_uses_shared_rotary_pos_enc(self):
+        """rope mode builds a single ``RotaryPositionalEncoding`` reused by every attention layer.
+
+        The cos/sin buffers are computed once on the shared module (see ``TransformerEncoder``),
+        so each layer's ``attn.rope`` must be the *same* object as ``model.pos_enc``.
+        """
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=3, self_attention_model="rope")
+        assert isinstance(model.pos_enc, RotaryPositionalEncoding)
+        for layer in model.layers:
+            attn = layer.attn
+            assert attn._uses_rope is True
+            assert attn.rope is model.pos_enc
+
+    @pytest.mark.unit
+    def test_rope_partial_rotation_forward_cpu(self):
+        """``rotary_fraction`` < 1.0 rotates only part of each head dim (exercises the pass-through split)."""
+        model = TransformerEncoder(
+            feat_in=128,
+            d_model=64,
+            n_heads=4,
+            n_layers=2,
+            drop_rate=0.0,
+            subsampling_factor=4,
+            self_attention_model="rope",
+            rotary_fraction=0.5,
+        )
+        model.eval()
+
+        B, C, T = 2, 128, 200
+        x = torch.randn(B, C, T)
+        lengths = torch.tensor([T, 160])
+
+        with torch.no_grad():
+            out, _ = model(audio_signal=x, length=lengths)
+
+        assert out.shape == (B, 64, T // 4)
+        assert not torch.isnan(out).any()

tests/collections/asr/utils/test_asr_multispeaker_utils.py

@@ -0,0 +1,129 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.collections.asr.parts.utils.asr_multispeaker_utils import (
+    find_first_nonzero,
+    get_hidden_length_from_sample_length,
+    read_rttm_supervisions_lenient,
+)
+
+
+def _write_rttm(path, lines):
+    path.write_text("\n".join(lines) + "\n")
+    return path
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "line",
+    [
+        "SPEAKER rec1 1 1.25 2.50 <NA> <NA> speaker_A <NA>",
+        "SPEAKER rec1 1 1.25 2.50 <NA> <NA> speaker_A <NA> <NA>",
+    ],
+)
+def test_read_rttm_supervisions_lenient_accepts_9_and_10_column_lines(tmp_path, line):
+    rttm_path = _write_rttm(tmp_path / "valid.rttm", [line])
+
+    supervisions = read_rttm_supervisions_lenient(rttm_path)
+
+    assert len(supervisions) == 1
+    segment = supervisions[0]
+    assert segment.id == "rec1-000000"
+    assert segment.recording_id == "rec1"
+    assert segment.channel == 1
+    assert segment.start == pytest.approx(1.25)
+    assert segment.duration == pytest.approx(2.50)
+    assert segment.speaker == "speaker_A"
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "line",
+    [
+        "SPEAKER rec1 1 0.0 1.0 <NA> <NA>",
+        "SPEAKER rec1 1 0.0 1.0",
+        "too short",
+    ],
+)
+def test_read_rttm_supervisions_lenient_rejects_short_lines(tmp_path, line):
+    rttm_path = _write_rttm(tmp_path / "invalid.rttm", [line])
+
+    with pytest.raises(ValueError, match="Invalid RTTM line"):
+        read_rttm_supervisions_lenient(rttm_path)
+
+
+@pytest.mark.unit
+def test_read_rttm_supervisions_lenient_skips_blank_and_zero_duration_lines(tmp_path):
+    rttm_path = _write_rttm(
+        tmp_path / "skip.rttm",
+        [
+            "",
+            "SPEAKER rec1 1 0.00 0.00 <NA> <NA> speaker_A <NA>",
+            "SPEAKER rec1 1 3.00 1.25 <NA> <NA> speaker_B <NA>",
+        ],
+    )
+
+    supervisions = read_rttm_supervisions_lenient(rttm_path)
+
+    assert len(supervisions) == 1
+    segment = supervisions[0]
+    assert segment.id == "rec1-000002"
+    assert segment.start == pytest.approx(3.00)
+    assert segment.duration == pytest.approx(1.25)
+    assert segment.speaker == "speaker_B"
+
+
+@pytest.mark.unit
+def test_read_rttm_supervisions_lenient_accepts_multiple_files(tmp_path):
+    rttm_path_a = _write_rttm(tmp_path / "a.rttm", ["SPEAKER rec_a 1 0.00 1.00 <NA> <NA> speaker_A <NA>"])
+    rttm_path_b = _write_rttm(tmp_path / "b.rttm", ["SPEAKER rec_b 2 1.00 2.00 <NA> <NA> speaker_B <NA>"])
+
+    supervisions = read_rttm_supervisions_lenient([rttm_path_a, rttm_path_b])
+
+    assert len(supervisions) == 2
+    assert [segment.recording_id for segment in supervisions] == ["rec_a", "rec_b"]
+    assert [segment.channel for segment in supervisions] == [1, 2]
+    assert [segment.speaker for segment in supervisions] == ["speaker_A", "speaker_B"]
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    ("num_samples", "expected_hidden_length"),
+    [
+        (0, 0),
+        (1, 1),
+        (1280, 1),
+        (1281, 2),
+    ],
+)
+def test_get_hidden_length_rounds_up_to_encoder_frames(num_samples, expected_hidden_length):
+    assert get_hidden_length_from_sample_length(num_samples) == expected_hidden_length
+
+
+@pytest.mark.unit
+def test_find_first_nonzero_returns_first_threshold_crossing_or_cap():
+    mat = torch.tensor(
+        [
+            [0.0, 0.2, 0.6],
+            [0.0, 0.0, 0.0],
+            [0.7, 0.8, 0.0],
+        ]
+    )
+
+    result = find_first_nonzero(mat, max_cap_val=99, thres=0.5)
+
+    assert torch.equal(result, torch.tensor([2, 99, 0]))

tests/collections/asr/utils/test_chunking_utils.py

@@ -0,0 +1,248 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.collections.asr.parts.utils.chunking_utils import (
+    join_char_level_timestamps,
+    merge_all_hypotheses,
+    merge_hypotheses_of_same_audio,
+)
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+
+
+def _make_char(char, token_id, start_off, end_off, token=None):
+    return {
+        "char": char,
+        "token": token if token is not None else char,
+        "token_id": token_id,
+        "start_offset": start_off,
+        "end_offset": end_off,
+    }
+
+
+@pytest.mark.unit
+def test_join_char_level_timestamps_without_filter():
+    # Merging char level timestamps within same audio segment.
+    subsampling_factor = 8
+    window_stride = 0.01
+    chunk_offsets = [0, 32]
+
+    h0 = Hypothesis(
+        score=0.0,
+        y_sequence=torch.tensor([]),
+        timestamp={
+            "char": [
+                _make_char("a", 10, 0, 1),
+                _make_char("b", 11, 2, 3),
+            ]
+        },
+    )
+    h1 = Hypothesis(
+        score=0.0,
+        y_sequence=torch.tensor([]),
+        timestamp={
+            "char": [
+                _make_char("b", 12, 0, 1),
+                _make_char("c", 13, 2, 3),
+            ]
+        },
+    )
+
+    out = join_char_level_timestamps(
+        hypotheses=[h0, h1],
+        chunk_offsets=chunk_offsets,
+        subsampling_factor=subsampling_factor,
+        window_stride=window_stride,
+        merged_tokens=None,
+    )
+
+    assert len(out) == 4
+    shift = chunk_offsets[1] // subsampling_factor
+
+    assert out[0]["start_offset"] == 0 and out[0]["end_offset"] == 1
+    assert out[1]["start_offset"] == 2 and out[1]["end_offset"] == 3
+
+    assert out[2]["start_offset"] == 0 + shift and out[2]["end_offset"] == 1 + shift
+    assert out[3]["start_offset"] == 2 + shift and out[3]["end_offset"] == 3 + shift
+
+    sec_per_subsample = window_stride * subsampling_factor
+    assert out[0]["start"] == pytest.approx(out[0]["start_offset"] * sec_per_subsample)
+    assert out[3]["end"] == pytest.approx(out[3]["end_offset"] * sec_per_subsample)
+
+
+@pytest.mark.unit
+def test_join_char_level_timestamps_with_filter():
+    # Merging char level timestamps within same audio segment.
+    subsampling_factor = 8
+    window_stride = 0.01
+    chunk_offsets = [0, 200]
+
+    # Chunk0: tokens 1..4
+    h0 = Hypothesis(
+        score=0.0,
+        y_sequence=torch.tensor([]),
+        timestamp={
+            "char": [
+                _make_char("a", 1, 0, 0),
+                _make_char("b", 2, 1, 1),
+                _make_char("c", 3, 2, 2),
+                _make_char("d", 4, 3, 3),
+            ]
+        },
+    )
+    # Chunk1: overlaps and -1 offsets as provided
+    h1 = Hypothesis(
+        score=0.0,
+        y_sequence=torch.tensor([]),
+        timestamp={
+            "char": [
+                _make_char("a", 1, 0, 0),
+                _make_char("c", 3, 1, 1),
+                _make_char("d", 4, 2, 2),
+                _make_char("e", 5, -1, 3),
+                _make_char("f", 6, 4, 4),
+                _make_char("g", 7, -1, -1),
+            ]
+        },
+    )
+
+    merged_tokens = [1, 2, 3, 4, 5, 6, 7]
+
+    out = join_char_level_timestamps(
+        hypotheses=[h0, h1],
+        chunk_offsets=chunk_offsets,
+        subsampling_factor=subsampling_factor,
+        window_stride=window_stride,
+        merged_tokens=merged_tokens,
+    )
+
+    # Token IDs in order
+    assert [d["token_id"] for d in out] == merged_tokens
+    # Expected global offsets (from your provided output)
+    expected_start_offsets = [0, 1, 2, 3, -1, 29, -1]
+    expected_end_offsets = [0, 1, 2, 3, 28, 29, -1]
+    assert [d["start_offset"] for d in out] == expected_start_offsets
+    assert [d["end_offset"] for d in out] == expected_end_offsets
+
+    # Expected times
+    expected_starts = [0.0, 0.08, 0.16, 0.24, -1, 2.32, -1]
+    expected_ends = [0.0, 0.08, 0.16, 0.24, 2.24, 2.32, -1]
+
+    assert [d["start"] for d in out] == pytest.approx(expected_starts)
+    assert [d["end"] for d in out] == pytest.approx(expected_ends)
+
+
+@pytest.mark.unit
+def test_merge_hypotheses_of_same_audio():
+    # Different segments of the same audio file are correctly combined
+    subsampling_factor = 8
+    chunk_duration_seconds = 10
+    frame_offset = int(chunk_duration_seconds * 1000 / subsampling_factor)
+
+    h0 = Hypothesis(
+        score=0.0,
+        y_sequence=torch.tensor([1]),
+        timestamp={
+            "word": [{"word": "a", "start": 0.0, "end": 0.1, "start_offset": 0, "end_offset": 2}],
+            "segment": [{"segment": "a", "start": 0.0, "end": 0.1, "start_offset": 0, "end_offset": 2}],
+        },
+    )
+    h1 = Hypothesis(
+        score=0.0,
+        y_sequence=torch.tensor([2]),
+        timestamp={
+            "word": [{"word": "b", "start": 0.2, "end": 0.3, "start_offset": 0, "end_offset": 3}],
+            "segment": [{"segment": "b", "start": 0.2, "end": 0.3, "start_offset": 0, "end_offset": 3}],
+        },
+    )
+    h2 = Hypothesis(
+        score=0.0,
+        y_sequence=torch.tensor([3]),
+        timestamp={
+            "word": [],
+            "segment": [],
+        },
+    )
+
+    merged = merge_hypotheses_of_same_audio(
+        hypotheses_list=[h0, h1, h2],
+        timestamps=True,
+        subsampling_factor=subsampling_factor,
+        chunk_duration_seconds=chunk_duration_seconds,
+    )
+
+    words = merged.timestamp["word"]
+    segs = merged.timestamp["segment"]
+
+    assert [w["word"] for w in words] == ["a", "b"]
+    assert words[0]["start"] == pytest.approx(0.0)
+    assert words[0]["start_offset"] == 0
+    assert words[1]["start"] == pytest.approx(0.2 + chunk_duration_seconds)
+    assert words[1]["start_offset"] == frame_offset
+
+    assert [s["segment"] for s in segs] == ["a", "b"]
+    assert segs[1]["end"] == pytest.approx(0.3 + chunk_duration_seconds)
+    assert segs[1]["end_offset"] == 3 + frame_offset
+
+
+@pytest.mark.unit
+def test_merge_all_hypotheses():
+    # Testing if merging by id works
+    def H(text, id_):
+        h = Hypothesis(score=0.0, y_sequence=torch.tensor([1]), timestamp={"word": [], "segment": []})
+        h.text = text
+        h.id = id_
+        return h
+
+    hyps = [H("a", 1), H("b", 1), H("c", 2), H("d", 2)]
+
+    merged_list = merge_all_hypotheses(
+        hypotheses_list=hyps,
+        timestamps=False,
+        subsampling_factor=2,
+        chunk_duration_seconds=3600,
+    )
+
+    assert len(merged_list) == 2
+    texts = {m.text for m in merged_list}
+    assert texts == {"a b", "c d"}
+
+
+@pytest.mark.unit
+def test_merge_all_hypotheses_with_cut_segmented_suffix():
+    def H(text, id_):
+        h = Hypothesis(score=0.0, y_sequence=torch.tensor([1]), timestamp={"word": [], "segment": []})
+        h.text = text
+        h.id = id_
+        return h
+
+    hyps = [
+        H("root", "11-0"),
+        H("cont1", "11-1_cut_segmented"),
+        H("cont2", "11-2_cut_segmented"),
+        H("other", "12-0"),
+    ]
+
+    merged_list = merge_all_hypotheses(
+        hypotheses_list=hyps,
+        timestamps=False,
+        subsampling_factor=8,
+        chunk_duration_seconds=3600,
+    )
+
+    assert len(merged_list) == 2
+    texts = sorted(m.text for m in merged_list)
+    assert texts == ["other", "root cont1 cont2"]