Metadata & State Export

Logic parameter extraction utilities for model portability and symbolic reasoning.

This module provides functions to traverse the Flax NNX module hierarchy and extract learned logical parameters (weights, biases, thresholds) into standard Python dictionaries. This allows for saving model metadata as JSON or exporting the neural logic back into symbolic rules.

jlnn.export.data.export_module_recursive(module: Any, name: str = 'root') → Dict[str, Any][source]

Recursively extracts logical parameters while preserving the module hierarchy.

This function traverses the Flax NNX module tree, identifying logical gates and capturing their specific parameters. It maintains the tree structure, allowing for the reconstruction of nested logical expressions (e.g., reduction trees in n-ary gates).

Parameters:

  • module (Any) – An instance of an NNX Module or a sub-gate to be inspected.

  • name (str) – The attribute name of the module within its parent. Defaults to “root”.

Returns:

A dictionary containing the gate’s name, type,

extracted parameters (weights, beta, etc.), and optional nested sub-gates.

Return type:

Dict[str, Any]

jlnn.export.data.extract_logic_parameters(model: Module) → Dict[str, Any][source]

Primary API to extract the full logical structure and weights of a JLNN model.

This function produces a portable, structured representation of the entire logical network. The resulting dictionary is suitable for JSON serialization and can be used by external symbolic reasoners to verify or visualize learned rules.

Parameters:

model (nnx.Module) – The top-level Flax NNX logical model instance.

Returns:

A nested dictionary containing framework metadata

and the recursive gate structure of the model.

Return type:

Dict[str, Any]

This module handles the serialization of model state (parameters) and its topology. It allows saving a trained model so that it can be later reconstructed without recompiling from formulas.

  • Weight Serialization: Exporting nnx.Param (weights and betas) to Msgpack or JSON format.

  • Topology Export: Saving the graph structure for visualization or static analysis purposes.