Batching rule not implemented for aten::rsub.Scalar, aten::diag, aten::where.Scalar, aten::allclose.

[pytorch]

Batching rule not implemented for aten::rsub.Scalar, aten::diag, aten::where.Scalar, aten::allclose.

[Chillee]

Thanks for providing lists of batching rules where we still need to fill in the gaps :) We'll work towards adding them.

Other than allclose (where we couldn't generate a fallback), does the API suffice for your needs? We're looking for feedback on APIs and such :)

Ps: thank you for using the same vmap syntax as jax that saved a lot of time converting the code.

Yeah, all credit to the Jax team for coming up with such a great API :)

[Chillee]

Looking closer at the usage of aten::allclose, I'm not sure that you'll be able to vmap through it without some modifications to the library.

Unfortunately, the code looks something like this:

is_hermitian: bool = torch.allclose(a, a.t())

The problem is that a is now batched, so what should torch.allclose return? A list of booleans? Then what happens if you have

is_hermitian: bool = torch.allclose(a, a.t())
if is_hermitian:
    do foo
else:
   do blah

My suspicion is that the code will need to be restructured in order to vmap over it.

[JonathanSchmidt1]

Thank you for taking a look at the xitorch issue, than I will rewrite the code as this check is not 100% required.
Concerning the api if you just use the jax api I have no suggestions for improvements, it already works great.

[Chillee]

I suspect that in order to resolve this, we may need to have control flow operators (like Jax's lax.cond).