successor library for autograd?

[CamDavidsonPilon]

HIPS/autograd entered into unmaintained status, so we should plan to move away from it.

1. Fork it and maintain it ourselves.
2. Choose a different library like JAX, tinygrad, any other options?
3. vendor it and keep only the parts of autograd we need.

1. Fork it and maintain it ourselves

This is probably the least attractive option. I'm not an expert in autodiff, or the implementation in autograd, and I don't want to support or be responsible for other libraries using autograd. Pass.

The only benefit is that users don't need to update their code from import autograd.numpy as anp to import x.numpy as xnp.

2. Choose a different library like JAX, tinygrad etc.

• One important part of any migration is being able to use a numpy/scipy submodule that mimics the numpy/scipy namespace. This is easiest for me and users to reason about.

  • autograd: yes, but not numpy 2.0 compatible.
  • jax: yes, and numpy 2.0 compatible.
  • tinygrad: no

• I currently maintain the autograd-gamma library, which is a numerical approximation to gamma and beta functions (which didn't have nice derivatives for core autograd). How to other libraries handle gamma & beta functions?

  • autograd: yes, with my support.
  • jax: yes? https://github.com/google/jax/issues/479, https://github.com/google/jax/pull/18397, and they are probably better versions that my implementation
  • tinygrad: no

• @Zethson brings up a good point about size and bloat. Autograd is a small 70kB, and jax for example is 1700kB (OTOH matplotlib is 36000 kB). Do users who have autograd/lifelines also often have JAX installed anyways.

  • Importing jax (for the first time) on my M1 MBP takes like 5s =S

3. vendor it and keep only the parts of autograd we need.

This is close to 1., but we don't need to support a community, just our own need. This is probably strictly better than 1. However, some notes:

• We leave the upgrade train =(
• We can incorporate autograd-gamma natively =)
• Users still need to update imports =(
• Future maintainers need to know and be experts in autodiff (and numpy/scipy) =(

[imessien]

Tinygrad is great!

[Zethson]

These are quite big dependencies. If you only need a small part and can vendor it, I'd advocate for it

[Zethson]

Different option: Make Jax an optional dependency and require it when accessing certain lifelines functionality where it would be required.

[fkiraly]

Different option: Make Jax an optional dependency and require it when accessing certain lifelines functionality where it would be required.

In sktime and skpro we manage dependencies on the level of estimators, so adopting the patterns therein would allow solving the issue along the lines of @Zethson. Happy to help if we go down that route, as I am familiar with dependency isolation per estimator.

There seems to be a substantial footprint - I would suggest making a battle plan:

• deciding on a target design
• making a list of estimators that need to be refactored
• possibly isolating the dependencies already early, so numpy 2.0 compatibility can be done more quickly

[fkiraly]

My preference would be option 2, jax.

I also wonder why there is a discontinuity in the package history, the two packages share some authors. Any gossip about jax/autograd relation?

[CamDavidsonPilon]

So things are happening with autograd, and some new maintainers have stepped up. I will keep any eye on their progress.

[CamDavidsonPilon]

Well that was a lucky break: https://github.com/HIPS/autograd/releases/tag/v1.7.0

I'll test and update this shortly.

[AGCaesar]

I had the same problem with numpy 2.0, my python-autograd was just updated to 1.7 and I rebuild lifelines and everything works great again :)