brainstate.interop.to_nnx

Contents

brainstate.interop.to_nnx#

brainstate.interop.to_nnx(model, *, rngs=None)[source]#

Convert a brainstate.nn model into an equivalent flax.nnx model.

Parameters:

  • model (Any) – The source model.

  • rngs (Any) – RNG container used to construct the nnx layers. Construction weights are overwritten immediately, so this only matters if you reuse the rngs afterwards.

Returns:

The converted model.

Return type:

Any

Contents