BDF solver for stiff ODEs

[skrsna]

This PR addresses #3686 .

Implements BDF solver with JAX control flow and JIT support using code from tensforflow_probability. The current implementation is tested against SciPy's VODE wrapper with stiff chemical kinetics.

To-Do:

• Adjoint sensitivities using custom_vjp

• Lazy jacobian evaluation

• Automatically calculate jacobian using JAX's autodiff if user doesn't provide jacobian function

• Automatically choose times if user only provides end_time

Comments and suggestions are welcome. @shoyer

[googlebot]

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it! and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[skrsna]

@googlebot I signed it!

[googlebot]

CLAs look good, thanks!

ℹ️ Googlers: Go here for more info.

[skrsna]

I'm converting the PR to ready for review cause I need some guidance and suggestions with custom_vjp implementation. Right now the results from integration are stored in named tuple and I'd want some input on how to structure/store adjoint gradient method results etc. @shoyer @mattjj

[skrsna]

Hi @shoyer and @mattjj ,

I have been benchmarking the current solver against SciPy's VODE with large chemical systems (e.g. systems with IVPs of ~600-1000) and the solver seems to be really really slow even on an accelerator (NVIDIA Tesla V100). I searched through issues and couldn't find anything solid to track recompilation. I assumed JIT+GPU should at least be comparable to CPU times for large systems. I profiled the code using NVProf but I can only see kernel calls and I'm in the process of understanding and setting up profiling using TensorboardX so I can track python function calls. Any help would be greatly appreciated regarding this as the potential use case for this solver (along with adjoint sensitivities) is to solve chemical systems in the order of ~10000-40000 IVPs. Thanks!!!! 🙏

[mattjj]

Hey @skrsna , awesome work! Sorry we haven't had a chance to follow up; we're really slammed with work right now.

About the slowness, that's something we'd certainly like to fix! Can you try profiling using the TensorBoard profiler? That can show a lot more information, including both host and device activity in detail and often with source code provenance information.

[skrsna]

Sure! I had trouble setting up Tensorboard profiler on our campus' cluster due to port/network issues and our on premise GPU workstation is really old so I kept running into CUDA OOM for large systems. I'll work with our sys admins to figure out Tensorboard. The project I'm working on is in really early stages of development so it's not that urgent anyway. Thanks!!

[jekbradbury]

If you're concerned about recompilation, try setting JAX_LOG_COMPILES=1 and see if any recompilation still happens after the program reaches "steady state"?

[skrsna]

Thanks @jekbradbury, I tried your suggestion and I'm seeing lots of Compiling eval_jaxpr for args but with different length of args. Is this normal? Here's the full output. I'm sorry if a PR is not a right place to talk about these issues. We can move this conversation to the related issue at #3686.

Full compile log.

WARNING:absl:Compiling _where for args (AbstractUnit(), ShapedArray(float64[], weak_type=True), ShapedArray(int64[], weak_type=True)).
WARNING:absl:Compiling calculate_rate_constants for args (AbstractUnit(), ShapedArray(float64[], weak_type=True), ShapedArray(float64[325,3])).
WARNING:absl:Compiling _where for args (AbstractUnit(), ShapedArray(float64[], weak_type=True), ShapedArray(int64[], weak_type=True)).
WARNING:absl:Compiling calculate_rate_constants for args (AbstractUnit(), ShapedArray(float64[], weak_type=True), ShapedArray(float64[325,3]), AbstractUnit()).
WARNING:absl:Compiling interpolate_backward_differences for args (AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling _where for args (AbstractUnit(), ShapedArray(float64[], weak_type=True), ShapedArray(int64[], weak_type=True)).
WARNING:absl:Compiling calculate_rate_constants for args (AbstractUnit(), ShapedArray(float64[], weak_type=True), ShapedArray(float64[325,3])).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling eval_jaxpr for args (AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit(), AbstractUnit()).
WARNING:absl:Compiling solve_to_end_time for args (ShapedArray(float64[54]), ShapedArray(float64[], weak_type=True), ShapedArray(float64[], weak_type=True), ShapedArray(float64[], weak_type=True), ShapedArray(float64[], weak_type=True)).