generalize the jax pytree interface

[jcmgray]

No description provided.

[codecov]

Codecov Report

Merging #151 (3d4777f) into develop (2dbbcb8) will increase coverage by 34.92%.
The diff coverage is 50.66%.

@@             Coverage Diff              @@
##           develop     #151       +/-   ##
============================================
+ Coverage    35.51%   70.44%   +34.92%     
============================================
  Files           43       44        +1     
  Lines        17314    17369       +55     
============================================
+ Hits          6149    12235     +6086     
+ Misses       11165     5134     -6031     

Impacted Files	Coverage Δ
quimb/tensor/interface.py	36.58% <36.58%> (ø)
quimb/tensor/tensor_core.py	69.62% <59.09%> (+40.31%)	⬆️
quimb/utils.py	83.93% <71.42%> (+20.52%)	⬆️
quimb/tensor/__init__.py	100.00% <100.00%> (ø)
quimb/tensor/optimize.py	31.46% <100.00%> (+12.03%)	⬆️
quimb/experimental/tnvmc.py	0.00% <0.00%> (ø)
quimb/gen/rand.py	96.92% <0.00%> (+0.02%)	⬆️
quimb/linalg/approx_spectral.py	84.83% <0.00%> (+0.36%)	⬆️
quimb/gen/operators.py	98.42% <0.00%> (+0.52%)	⬆️
quimb/calc.py	97.86% <0.00%> (+0.97%)	⬆️
... and 25 more

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[jcmgray]

@mofeing, this was my attempt to generalize #150 a bit. I.e. provide a standard interface for packing and unpacking Tensor and TensorNetwork objects. Let me know if for some reason it doesn't capture the same behavior.

[mofeing]

I'm now running into an error within opt_einsum with these changes. More specifically inside the contract_path function at https://github.com/dgasmith/opt_einsum/blob/7c8f193f90b6771a6b3065bb5cf6ec2747af8209/opt_einsum/contract.py#L284-L287

Looks like the size of some label changes.

[jcmgray]

Hmm thanks, do you have a small snippet to reproduce? It's not come from loading/saving a TN across commits e.g.?

[jcmgray]

I'm gonna merge this as I think it fixes a problem with the previous jax registration (trying to register Tensor multiple times), and I can't see any obvious way it should break contractions.. but of course let me know if you can come up with a reproducible error.

[mofeing]

I don't know what changed but the error is gone, so great!

Could you do me a favour and publish this version as 1.4.1? This way we can properly set the dependency in our code.