factor out lambda in coix.util.train to avoid recompiling lax.cond.

jax-ml / coix

Inference Combinators in JAX

https://coix.readthedocs.io/en/latest/

Apache License 2.0

43 stars 2 forks source link

Closed copybara-service[bot] closed 1 year ago

copybara-service[bot] commented 1 year ago

factor out lambda in coix.util.train to avoid recompiling lax.cond.

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/jax-ml/coix/pull/15 from jax-ml:prng 222781afd64b8f77482a0f4adee181830b291783