Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add rewrite for matmul when only one of the inputs has batched dimensions #558

Merged
merged 1 commit into from
Dec 20, 2023
Merged

Add rewrite for matmul when only one of the inputs has batched dimensions #558

merged 1 commit into from
Dec 20, 2023

Commits on Dec 18, 2023

  1. Add rewrite for matmul when only one of the inputs has batched dimens… 

    …ions

This rewrites replaces a batched matmul by a core matmul by raveling the batched dimensions of the batched input, doing a 2D matmul and then unravelling the batched dimensions of the output.

This sidesteps the Blockwise Dot operation and allows specialization into BLAS routines.

The idea was taken from these two discussions:
numpy/numpy#7569
numpy/numpy#8957
    @ricardoV94
    ricardoV94 committed Dec 18, 2023
    Configuration menu
    Copy the full SHA
    b19c6ab View commit details
    Browse the repository at this point in the history