search

stanford-crfm / haliax

Named Tensors for Legible Deep Learning in JAX
Apache License 2.0
141 stars 9 forks source link

Upgraded to latest Equinox; removed Chex dependency. #7

Closed patrick-kidger closed 1 year ago

patrick-kidger commented 1 year ago

The primary purpose of this PR is to upgrade Haliax to be compatible with the latest Equinox (v0.10.6).

En passant, I also made use of the recently-added jaxtyping.PRNGKeyArray, which also made it possible to drop the dependency on Chex.

I've bumped the version number so that Levanter can use this.

I didn't run the pre-commit hooks as they tried to reformat a lot of things that I didn't touch.

dlwh commented 1 year ago

yeah commit hooks are a bit messed up from my splitting of the repo and bumping the pre-commit libs recently. I'll do a fix to this repo soon.