search

jax-ml / coix

Inference Combinators in JAX
https://coix.readthedocs.io/en/latest/
Apache License 2.0
43 stars 2 forks source link

Do not use cond in params update #40

Closed fehiepsi closed 4 months ago

fehiepsi commented 4 months ago

This is to keep memory friendly for training large models.