use BCSR in sparse format of jaxoperators instead of BCOO

[PhilipVinc]

BCOO should just be used to construct it, not to use it for matrix vector multiplications...

[gcarleo]

any speedup?

[wdphy16]

Yes I agree. Although in MCState we don't need to multiply a NetKet operator on a dense vector, when working with FullSumState it indeed speeds up in my use case.

By the way, in FullSumState we can put the sparse matrix - dense vector multiplication into JIT to further speed it up, at the cost of longer compiling time (specializing on the sparsity structure can take a long time) and higher memory usage, as in my commit

[PhilipVinc]

@wdphy16 the only questionable thing is that jax sparse matrices do not support multiplication by another jax CSR sparse matrix, so I had to ensure that when one tries to do jaxoperator @ sparsevector we convert it to dense first.

This breaks current behaviour where, as we use BCOO, we get a sparse vector out.

But I think that's not important? Hopefully one never does this kind of multiplication...

[wdphy16]

Ah it's a bit unexpected for me, but anyway JAX sparse matrices are experimental so we can make breaking changes

In real use cases, I guess we usually first do NetKet operator multiplications then convert it to a JAX sparse matrix

[PhilipVinc]

@wdphy16 how bad is compile time in that case?

[wdphy16]

@wdphy16 how bad is compile time in that case?

I don't remember exactly, but for large experiments it's much less than the actual computation time, and for small experiments I did not feel much more lag

[PhilipVinc]

@wdphy16 as the improvement that you propose makes sense, would you care making a PR at some point? I think it's reasonable .