Unverified Commit c64861d6 authored by Chih-Chieh Yang's avatar Chih-Chieh Yang Committed by GitHub
Browse files

[Bugfix] Mamba2 remove bugged initial state condition in chunk scan (#22034)


Signed-off-by: default avatarChih-Chieh-Yang <7364402+cyang49@users.noreply.github.com>
parent 8564dc94
...@@ -476,15 +476,8 @@ def _chunk_scan_fwd( ...@@ -476,15 +476,8 @@ def _chunk_scan_fwd(
# with initial states, we need to take care of how # with initial states, we need to take care of how
# seq_idx crosses the boundaries # seq_idx crosses the boundaries
assert batch == 1, "chunk scan only supports initial states with batch 1" assert batch == 1, "chunk scan only supports initial states with batch 1"
assert chunk_indices is not None and chunk_offsets is not None, \
if initial_states.shape[0] == 1: "chunk_indices and chunk_offsets should have been set"
# no in this case no point to use initial states
initial_states = None
else:
assert chunk_indices is not None and chunk_offsets is not None, \
(
"chunk_indices and chunk_offsets should have been set"
)
else: else:
chunk_indices, chunk_offsets = None, None chunk_indices, chunk_offsets = None, None
else: else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment