Skip to content

Validate tensor ranks before padded batching - #4149

Open
x-Spartacus wants to merge 1 commit into
tensorflow:masterfrom
x-Spartacus:padded-batching-rank-validation
Open

x-Spartacus wants to merge 1 commit into
tensorflow:masterfrom
x-Spartacus:padded-batching-rank-validation

Conversation

@x-Spartacus

Copy link
Copy Markdown

This change validates corresponding tensor ranks and schemas before calculating maximum dimensions or adding padding. The checks cover both the classic and TFRT batching paths, with an additional guard in AddPadding.

Previously, maximum dimensions were sized using the first task and then indexed using the ranks of later tasks. A mixed-rank batch could therefore index beyond the dimension vector.

Testing:

  • bazel test //tensorflow_serving/batching:batching_util_test
  • bazel test //tensorflow_serving/batching:batching_session_test

Regression coverage was also added for the TFRT path. That target could not be executed locally because existing external TFRT/XLA build rules fail first on undeclared includes.

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