Skip to content

Fix SSMState.update_states slice_update end index missing trailing dim - #283

Draft
stikves wants to merge 2 commits into
apple:mainfrom
stikves:sukru/fix-ssmstate-slice-update-index
Draft

stikves wants to merge 2 commits into
apple:mainfrom
stikves:sukru/fix-ssmstate-slice-update-index

Conversation

@stikves

@stikves stikves commented Sep 22, 2026

Copy link
Copy Markdown
Contributor

The end bounds tensor was built with range(1, cache.dim() - 1), producing cache.dim() - 2 size entries and cache.dim() - 1 elements total, one short of begin. Export/AIProgram.optimize rejects the slice_update with mismatched begin/end ranks; eager only passed because zip uses strict=False and treats the missing trailing dim as a full slice. Iterate over range(1, cache.dim()) so begin and end both span every cache dimension.

stikves and others added 2 commits September 22, 2026 06:55
apple#275)

The `end` bounds tensor was built with range(1, cache.dim() - 1), producing
cache.dim() - 2 size entries and cache.dim() - 1 elements total, one short of
`begin`. Export/AIProgram.optimize rejects the slice_update with mismatched
begin/end ranks; eager only passed because zip uses strict=False and treats the
missing trailing dim as a full slice. Iterate over range(1, cache.dim()) so
begin and end both span every cache dimension.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant