-
Notifications
You must be signed in to change notification settings - Fork 60
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.
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
Operator collections arraylias integration #291
Operator collections arraylias integration #291
Conversation
…o include jax_sparse
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It is definitely better code than before.
I think the addition of array_library and the use of arraylias has improved the maintainability of the code considerably.
Co-authored-by: Kento Ueda <[email protected]>
Co-authored-by: Kento Ueda <[email protected]>
Co-authored-by: Kento Ueda <[email protected]>
Co-authored-by: Kento Ueda <[email protected]>
Co-authored-by: Kento Ueda <[email protected]>
Co-authored-by: Kento Ueda <[email protected]>
Co-authored-by: Kento Ueda <[email protected]>
Co-authored-by: Kento Ueda <[email protected]>
Co-authored-by: Kento Ueda <[email protected]>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you, Dan. I've approved.
c0d6b95
into
qiskit-community:models_arraylias_integration
Co-authored-by: Kento Ueda <[email protected]>
Co-authored-by: Kento Ueda <[email protected]>
Co-authored-by: Kento Ueda <[email protected]>
Co-authored-by: Kento Ueda <[email protected]>
Summary
Updating operator_collections.py to work with arraylias.
General list of changes:
numpy
,jax
, andjax_sparse
are merged into a single class using arraylias dispatching (one each forOperatorCollection
,LindbladCollection
, andVectorizedLindbladCollection
. However, I've decided to keepscipy_sparse
as a separate special case in each instance, as the interface is just too different from the others. It would be nice to have it all handled in a single class, but I'm worried that this would require building too much scipy sparse into our numpy alias that wouldn't be used anywhere else."jax_sparse"
:conjugate
andtranspose
.linear_combo
function to the registration for"numpy"
,"jax"
, and"jax_sparse"
. This is justtensordot
for the first two but forjax_sparse
a somewhat-hacky implementation is necessary._preferred_lib
to prefer"jax_sparse"
if present.GeneratorModel
,HamiltonianModel
, andLindbladModel
) are currently broken. I've commented out the lines using the old versions of the operator collection classes to avoid errors. I'll address these in the next PR updating these classes.One more general change:
Solver
class, in which signals need to be added to, and removed from, models. This PR won't be the full change, but I'm going to mark this PR as closing Parameter dependent simulation for optimal bayesian experimentation - JAX @jit compatibility #243, in which the immutability of model classes would have prevented the weird behaviour from occurring in the first place. For this PR, that means getting rid of all of thesetter
methods for the different operator types. I will do this in the model classes in a subsequent PR as well.Details and comments
Current list of test files we need to make sure are passing before merge to feature branch:
We should keep track of this list in subsequent PRs.
One technical issue:
LindbladCollection
, it isn't possible to directly pass JAX sparse arrays at construction. This was always the case, but I had hoped to change that here. The handling of then_batch
argument atBCOO
construction needs to be set to a particular value to get it to work here, and this isn't something I want the user to need to do. For now, they will need to pass dense arrays and selectarray_library="jax_sparse"
.