search

pyro-ppl / numpyro

Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.
https://num.pyro.ai
Apache License 2.0
2.15k stars 235 forks source link

Fix ECS to work with Haiku #1750

Closed OlaRonning closed 7 months ago

OlaRonning commented 7 months ago

The Haiku module requires a set of params (net$params) with layer shapes and rng_state to run. ECS shouldn't try to map these to constraint space.