FrEIA.framework package#

Module contents#

The framework module contains the logic used in building the graph and inferring the order that the nodes have to be executed in forward and backward direction.

class FrEIA.framework.ConditionNode(*dims: int, name=None)[source]#

Bases: AbstractNode

Special type of node that represents conditional input to the internal networks inside coupling layers.

__init__(*dims: int, name=None)[source]#
build_module(condition_shapes: List[Tuple[int]], input_shapes: List[Tuple[int]]) → Tuple[None, List[Tuple[int]]][source]#

Instantiates the module and determines the output dimension.

consume_output(out_idx: int, by_node: AbstractNode, in_idx: int)[source]#
rev_input(out_idx)[source]#

Compute which value should be read in reverse mode.

class FrEIA.framework.FeedForwardNode(conditions: AbstractNode | Tuple[AbstractNode, int] | Iterable[Tuple[AbstractNode, int]], output_dims: Tuple[int], module_type: Type | Module, module_args: dict | None = None, name=None)[source]#

Bases: AbstractNode

Special type of node that computes output without Jacobian and inverse.

__init__(conditions: AbstractNode | Tuple[AbstractNode, int] | Iterable[Tuple[AbstractNode, int]], output_dims: Tuple[int], module_type: Type | Module, module_args: dict | None = None, name=None)[source]#
build_module(condition_shapes: List[Tuple[int]], input_shapes: List[Tuple[int]]) → Tuple[Module, List[Tuple[int]]][source]#

Instantiates the module and determines the output dimension by calling InvertibleModule#output_dims.

rev_input(out_idx)[source]#

Compute which value should be read in reverse mode.

class FrEIA.framework.GraphINN(node_list: Iterable[AbstractNode], force_tuple_output=False, verbose=False)[source]#

Bases: InvertibleModule

This class represents the invertible net itself. It is a subclass of InvertibleModule and supports the same methods.

The forward method has an additional option ‘rev’, with which the net can be computed in reverse. Passing jac to the forward method additionally computes the log determinant of the (inverse) Jacobian of the forward (backward) pass.

__init__(node_list: Iterable[AbstractNode], force_tuple_output=False, verbose=False)[source]#
Parameters:

  • dims_in – list of tuples specifying the shape of the inputs to this operator: dims_in = [shape_x_0, shape_x_1, ...]

  • dims_c – list of tuples specifying the shape of the conditions to this operator.

get_module_by_name(name) → Module | None[source]#

Return module of the first node in the graph with the provided name.

get_node_by_name(name) → Node | None[source]#

Return the first node in the graph with the provided name.

log_jacobian_numerical(x, c=None, rev=False, h=0.0001)[source]#

Approximate log Jacobian determinant via finite differences.

property node_list#
class FrEIA.framework.InputNode(*dims: int, name=None)[source]#

Bases: AbstractNode

Special type of node that represents the input data of the whole net (or the output when running reverse)

__init__(*dims: int, name=None)[source]#
build_module(condition_shapes: List[Tuple[int]], input_shapes: List[Tuple[int]]) → Tuple[None, List[Tuple[int]]][source]#

Instantiates the module and determines the output dimension.

rev_input(out_idx)[source]#

Compute which value should be read in reverse mode.

class FrEIA.framework.Node(inputs: AbstractNode | Tuple[AbstractNode, int] | Iterable[Tuple[AbstractNode, int]], module_type: Type | InvertibleModule, module_args: dict | None = None, conditions=None, name=None)[source]#

Bases: AbstractNode

This Node class represents one transformation in the graph, with an arbitrary number of in- and outputs. These in and outputs count as being invertible, so they must be consumed.

The user specifies the input, and the underlying module computes the number of outputs.

build_module(condition_shapes: List[Tuple[int]], input_shapes: List[Tuple[int]]) → Tuple[InvertibleModule, List[Tuple[int]]][source]#

Instantiates the module and determines the output dimension by calling InvertibleModule#output_dims.

rev_input(out_idx)[source]#

Compute which value should be read in reverse mode.

class FrEIA.framework.OutputNode(in_node: AbstractNode | Tuple[AbstractNode, int], name=None)[source]#

Bases: AbstractNode

Special type of node that represents the output of the whole net (or the input when running in reverse).

__init__(in_node: AbstractNode | Tuple[AbstractNode, int], name=None)[source]#
build_module(condition_shapes, input_shapes) → Tuple[None, List[Tuple[int]]][source]#

Instantiates the module and determines the output dimension.

consume_output(out_idx: int, by_node: AbstractNode, in_idx: int)[source]#
rev_input(out_idx)[source]#

Compute which value should be read in reverse mode.

class FrEIA.framework.ReversibleGraphNet(node_list, ind_in=None, ind_out=None, verbose=True, force_tuple_output=False)[source]#

Bases: GraphINN

__init__(node_list, ind_in=None, ind_out=None, verbose=True, force_tuple_output=False)[source]#
Parameters:

  • dims_in – list of tuples specifying the shape of the inputs to this operator: dims_in = [shape_x_0, shape_x_1, ...]

  • dims_c – list of tuples specifying the shape of the conditions to this operator.

class FrEIA.framework.ReversibleSequential(*dims: int)[source]#

Bases: SequenceINN

__init__(*dims: int)[source]#
Parameters:

  • dims_in – list of tuples specifying the shape of the inputs to this operator: dims_in = [shape_x_0, shape_x_1, ...]

  • dims_c – list of tuples specifying the shape of the conditions to this operator.

class FrEIA.framework.SequenceINN(*dims: int, force_tuple_output=False)[source]#

Bases: InvertibleModule

Simpler than FrEIA.framework.GraphINN: Only supports a sequential series of modules (no splitting, merging, branching off). Has an append() method, to add new blocks in a more simple way than the computation-graph based approach of GraphINN. For example:

inn = SequenceINN(channels, dims_H, dims_W)

for i in range(n_blocks):
    inn.append(FrEIA.modules.AllInOneBlock, clamp=2.0, permute_soft=True)
inn.append(FrEIA.modules.HaarDownsampling)
# and so on
__init__(*dims: int, force_tuple_output=False)[source]#
Parameters:

  • dims_in – list of tuples specifying the shape of the inputs to this operator: dims_in = [shape_x_0, shape_x_1, ...]

  • dims_c – list of tuples specifying the shape of the conditions to this operator.

append(module_class, cond=None, cond_shape=None, **kwargs)[source]#

Append a reversible block from FrEIA.modules to the network.

Parameters:

  • module_class – Class from FrEIA.modules.

  • cond (int) – index of which condition to use (conditions will be passed as list to forward()). Conditioning nodes are not needed for SequenceINN.

  • cond_shape (tuple[int]) – the shape of the condition tensor.

  • **kwargs – Further keyword arguments that are passed to the constructor of module_class (see example).

FrEIA.framework.collect_nodes(*input_nodes: AbstractNode)[source]#