LayerMapping

Contents

LayerMapping#

class brainstate.interop.LayerMapping(bst_type, framework, foreign_type, to_bst, to_foreign)[source]#

A bidirectional conversion mapping for one layer type.

Parameters:

  • bst_type (type) – The brainstate Module subclass.

  • framework (str) – One of "nnx", "linen", "equinox".

  • foreign_type (type) – The foreign layer class.

  • to_bst (Callable) – (foreign_module, ctx) -> bst.nn.Module. Builds and weight-fills a brainstate layer.

  • to_foreign (Callable) – (bst_module, ctx) -> foreign. Builds and weight-fills the foreign layer. For the functional linen framework the return value is whatever the linen adapter expects (a (module, params) pair handled by the adapter).

Contents