Commit b4ec483f authored by Gaetan Lepage's avatar Gaetan Lepage
Browse files

python312Packages.dm-haiku: adapt to jax 0.6.0 API changes 

https://docs.jax.dev/en/latest/jax.extend.html
https://github.com/google-deepmind/dm-haiku/commit/cfe8480d253a93100bf5e2d24c40435a95399c96
parent a3ef8a9f

pkgs/development/python-modules/dm-haiku/default.nix

+11 −0
Original line number Diff line number Diff line
@@ -58,6 +58,17 @@ let 
      })

    ];



    # AttributeError: jax.core.Var was removed in JAX v0.6.0. Use jax.extend.core.Var instead, and

    # see https://docs.jax.dev/en/latest/jax.extend.html for details.

    # Alrady on master: https://github.com/google-deepmind/dm-haiku/commit/cfe8480d253a93100bf5e2d24c40435a95399c96

    # TODO: remove at the next release

    postPatch = ''

      substituteInPlace haiku/_src/jaxpr_info.py \

        --replace-fail "jax.core.JaxprEqn" "jax.extend.core.JaxprEqn" \

        --replace-fail "jax.core.Var" "jax.extend.core.Var" \

        --replace-fail "jax.core.Jaxpr" "jax.extend.core.Jaxpr"

    '';



    build-system = [ setuptools ];



    dependencies = [