RFC-0034-batch-kronecker-product #61
Conversation
|
Hi @dannygoodacre! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at [email protected]. Thanks! |
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
|
|
||
|
|
||
| ## **Drawbacks** | ||
| This is not a breaking change. |
There was a problem hiding this comment.
Adding this new kwarg to torch.kron wouldn't be bc-breaking either but would make it diverge from numpy's kron.
cc @rgommers any plans on adding kron to the arrray API that we should follow here?
|
|
||
|
|
||
| ## **Proposed Implementation** | ||
| In the following, `a` and `b` are tensors of the specified sizes, and `batch_kron(a, b, batch_dim)` the proposed feature, where `batch_dim = 0` refers to the outermost dimension. |
There was a problem hiding this comment.
Not sure what you mean by outermost here. But 0 should refer to the first dimension indeed.
Note that to match existing function, we can also consider doing the same as mm/bmm and only have the 0th dimension always be batch for a bkron op.
There was a problem hiding this comment.
Keeping consistent with bmm makes sense. Then the only condition on the tensors would be that the first dimensions agree. Plus, batching in any other dimension is rather niche and potentially overkill to include.
| ``` | ||
| a.shape = (2,5,4) | ||
| b.shape = (2,5,7) | ||
| batch_kron(a, b, 0).shape = (2,25,28) |
There was a problem hiding this comment.
cc @zou3519 do you think that using vmap here would be the best way to expose this kind of feature rather than adding a new function? Can we implement a fast batched kron without having a native op for it?
No description provided.