probabilistic_model.probabilistic_circuit.jax.inner_layer

probabilistic_model.probabilistic_circuit.jax.inner_layer#

Classes#

`Layer`	Abstract class for Layers of a layered circuit.
`InnerLayer`	Abstract Base Class for inner layers
`InputLayer`	Abstract base class for univariate input units.
`SumLayer`	Abstract Base Class for inner layers
`SparseSumLayer`	Abstract Base Class for inner layers
`DenseSumLayer`	Abstract Base Class for inner layers
`ProductLayer`	A layer that represents the product of multiple other units.
`NXConverterLayer`	Class used for conversion from a probabilistic circuit in networkx to a layered circuit in jax.

Functions#

inverse_class_of(→ typing_extensions.Type[Layer])

Module Contents#

probabilistic_model.probabilistic_circuit.jax.inner_layer.inverse_class_of(clazz: typing_extensions.Type[probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.Unit]) → typing_extensions.Type[Layer]#

class probabilistic_model.probabilistic_circuit.jax.inner_layer.Layer#

Bases: equinox.Module, random_events.utils.SubclassJSONSerializer, abc.ABC

Abstract class for Layers of a layered circuit.

Layers have the same scope (set of variables) for every node in them.

_variables: typing_extensions.Optional[jax.numpy.array]#: The variable indices of the layer.

property variables: jax.Array#

Abstractmethod:

abstract set_variables(value: jax.Array)#

abstract log_likelihood_of_nodes_single(x: jax.numpy.array) → jax.numpy.array#

Calculate the log-likelihood of the distribution.

Parameters:: x – The input vector.
Returns:: The log-likelihood of every node in the layer for x.

log_likelihood_of_nodes(x: jax.numpy.array) → jax.numpy.array#: Vectorized version of log_likelihood_of_nodes_single()

abstract validate()#: Validate the parameters and their layouts.

property number_of_nodes: int#

Abstractmethod:
Returns:: The number of nodes in the layer.

all_layers() → typing_extensions.List[Layer]#

Returns:: A list of all layers in the circuit.

all_layers_with_depth(depth: int = 0) → typing_extensions.List[typing_extensions.Tuple[int, Layer]]#

Returns:: A list of tuples of all layers in the circuit with their depth.

abstract __deepcopy__(memo=None) → Layer#

Create a deep copy of the layer.

Parameters:: memo – A dictionary that is used to keep track of objects that have already been copied.

classmethod nx_classes() → typing_extensions.Tuple[typing_extensions.Type, Ellipsis]#

Returns:: The tuple of matching classes of the layer in the probabilistic_model.probabilistic_circuit.rx package.

abstract to_nx(variables: sortedcontainers.SortedSet[random_events.variable.Variable], result: probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.ProbabilisticCircuit, progress_bar: typing_extensions.Optional[tqdm.tqdm] = None) → typing_extensions.List[probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.Unit]#

Convert the layer to a networkx circuit. For every node in this circuit, a corresponding node in the networkx circuit is created. The nodes all belong to the same circuit.

Parameters:

variables – The variables of the circuit.
result – The resulting circuit to write into
progress_bar – A progress bar to show the progress.

Returns:

The nodes of the networkx circuit.

static create_layers_from_nodes(nodes: typing_extensions.List[probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.Unit], child_layers: typing_extensions.List[NXConverterLayer], progress_bar: bool = True) → typing_extensions.List[NXConverterLayer]#: Create a layer from a list of nodes.

classmethod create_layer_from_nodes_with_same_type_and_scope(nodes: typing_extensions.List[probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.Unit], child_layers: typing_extensions.List[NXConverterLayer], progress_bar: bool = True) → NXConverterLayer#

Abstractmethod:

Create a layer from a list of nodes with the same type and scope.

partition() → typing_extensions.Tuple[typing_extensions.Any, typing_extensions.Any]#

Partition the layer into the parameters and the static structure.

Returns:: A tuple containing the parameters and the static structure as pytrees.

property number_of_trainable_parameters#

Returns:: The trainable parameters of the layer and all child layers.

property number_of_components: int#

Returns:: The number of components (leaves + edges) of the entire circuit

class probabilistic_model.probabilistic_circuit.jax.inner_layer.InnerLayer(child_layers: typing_extensions.List[Layer])#

Bases: Layer, abc.ABC

Abstract Base Class for inner layers

child_layers: typing_extensions.List[Layer]#: The child layers of this layer.

set_variables(value: jax.numpy.array)#

reset_variables()#

all_layers() → typing_extensions.List[Layer]#

Returns:: A list of all layers in the circuit.

all_layers_with_depth(depth: int = 0) → typing_extensions.List[typing_extensions.Tuple[int, Layer]]#

Returns:: A list of tuples of all layers in the circuit with their depth.

to_json() → typing_extensions.Dict[str, typing_extensions.Any]#

class probabilistic_model.probabilistic_circuit.jax.inner_layer.InputLayer(variable: int)#

Bases: Layer, abc.ABC

Abstract base class for univariate input units.

Input layers contain only one type of distribution such that the vectorization of the log likelihood calculation works without bottleneck statements like if/else or loops.

_variables#: The variable indices of the layer.

property variables: jax.Array#

set_variables(value: jax.Array)#

to_json() → typing_extensions.Dict[str, typing_extensions.Any]#

property variable#

class probabilistic_model.probabilistic_circuit.jax.inner_layer.SumLayer(child_layers: typing_extensions.List[Layer], log_weights: typing_extensions.List[typing_extensions.Union[jax.array, jax.experimental.sparse.BCOO]])#

Bases: InnerLayer, abc.ABC

Abstract Base Class for inner layers

log_weights: typing_extensions.List[typing_extensions.Union[jax.array, jax.experimental.sparse.BCOO]]#

child_layers: typing_extensions.Union[typing_extensions.List[[ProductLayer]], typing_extensions.List[InputLayer]]#: The child layers of this layer.

validate()#: Validate the parameters and their layouts.

property log_weighted_child_layers: typing_extensions.Iterator[typing_extensions.Tuple[jax.experimental.sparse.BCOO, Layer]]#

Returns:: Yields log log_weights and the child layers zipped together.

property variables: jax.Array#

property number_of_nodes: int#

Returns:: The number of nodes in the layer.

class probabilistic_model.probabilistic_circuit.jax.inner_layer.SparseSumLayer(child_layers: typing_extensions.List[Layer], log_weights: typing_extensions.List[typing_extensions.Union[jax.array, jax.experimental.sparse.BCOO]])#

Bases: SumLayer

Abstract Base Class for inner layers

log_weights: typing_extensions.List[jax.experimental.sparse.BCOO]#

classmethod nx_classes() → typing_extensions.Tuple[typing_extensions.Type, Ellipsis]#

Returns:: The tuple of matching classes of the layer in the probabilistic_model.probabilistic_circuit.rx package.

property number_of_components: int#

Returns:: The number of components (leaves + edges) of the entire circuit

property concatenated_log_weights: jax.experimental.sparse.BCOO#

Returns:: The concatenated log_weights of the child layers for each node.

property log_normalization_constants: jax.Array#

property normalized_weights#

Returns:: The normalized log_weights of the child layers for each node.

log_likelihood_of_nodes_single(x: jax.Array) → jax.Array#

Calculate the log-likelihood of the distribution.

Parameters:: x – The input vector.
Returns:: The log-likelihood of every node in the layer for x.

__deepcopy__(memo=None)#

Create a deep copy of the layer.

Parameters:: memo – A dictionary that is used to keep track of objects that have already been copied.

to_json() → typing_extensions.Dict[str, typing_extensions.Any]#

classmethod _from_json(data: typing_extensions.Dict[str, typing_extensions.Any]) → typing_extensions.Self#

Create a variable from a json dict. This method is called from the from_json method after the correct subclass is determined and should be overwritten by the respective subclass.

Parameters:: data – The json dict
Returns:: The deserialized object

classmethod create_layer_from_nodes_with_same_type_and_scope(nodes: typing_extensions.List[probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.SumUnit], child_layers: typing_extensions.List[NXConverterLayer], progress_bar: bool = True) → NXConverterLayer#: Create a layer from a list of nodes with the same type and scope.

to_nx(variables: sortedcontainers.SortedSet[random_events.variable.Variable], result: probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.ProbabilisticCircuit, progress_bar: typing_extensions.Optional[tqdm.tqdm] = None) → typing_extensions.List[probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.Unit]#

Convert the layer to a networkx circuit. For every node in this circuit, a corresponding node in the networkx circuit is created. The nodes all belong to the same circuit.

Parameters:

variables – The variables of the circuit.
result – The resulting circuit to write into
progress_bar – A progress bar to show the progress.

Returns:

The nodes of the networkx circuit.

class probabilistic_model.probabilistic_circuit.jax.inner_layer.DenseSumLayer(child_layers: typing_extensions.List[Layer], log_weights: typing_extensions.List[typing_extensions.Union[jax.array, jax.experimental.sparse.BCOO]])#

Bases: SumLayer

Abstract Base Class for inner layers

log_weights: typing_extensions.List[jax.numpy.array]#

child_layers: typing_extensions.Union[typing_extensions.List[[ProductLayer]], typing_extensions.List[InputLayer]]#: The child layers of this layer.

classmethod create_layer_from_nodes_with_same_type_and_scope(nodes: typing_extensions.List[probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.Unit], child_layers: typing_extensions.List[NXConverterLayer], progress_bar: bool = True) → NXConverterLayer#

Abstractmethod:

Create a layer from a list of nodes with the same type and scope.

property number_of_components: int#

Returns:: The number of components (leaves + edges) of the entire circuit

classmethod nx_classes() → typing_extensions.Tuple[typing_extensions.Type, Ellipsis]#

Returns:: The tuple of matching classes of the layer in the probabilistic_model.probabilistic_circuit.rx package.

property concatenated_log_weights: jax.numpy.array#

Returns:: The concatenated log_weights of the child layers for each node.

property log_normalization_constants: jax.Array#

property normalized_weights#

Returns:: The normalized log_weights of the child layers for each node.

log_likelihood_of_nodes_single(x: jax.Array) → jax.Array#

Calculate the log-likelihood of the distribution.

Parameters:: x – The input vector.
Returns:: The log-likelihood of every node in the layer for x.

__deepcopy__(memo=None)#

Create a deep copy of the layer.

Parameters:: memo – A dictionary that is used to keep track of objects that have already been copied.

to_json() → typing_extensions.Dict[str, typing_extensions.Any]#

classmethod _from_json(data: typing_extensions.Dict[str, typing_extensions.Any]) → typing_extensions.Self#

Create a variable from a json dict. This method is called from the from_json method after the correct subclass is determined and should be overwritten by the respective subclass.

Parameters:: data – The json dict
Returns:: The deserialized object

to_nx(variables: sortedcontainers.SortedSet[random_events.variable.Variable], result: probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.ProbabilisticCircuit, progress_bar: typing_extensions.Optional[tqdm.tqdm] = None) → typing_extensions.List[probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.Unit]#

Convert the layer to a networkx circuit. For every node in this circuit, a corresponding node in the networkx circuit is created. The nodes all belong to the same circuit.

Parameters:

variables – The variables of the circuit.
result – The resulting circuit to write into
progress_bar – A progress bar to show the progress.

Returns:

The nodes of the networkx circuit.

class probabilistic_model.probabilistic_circuit.jax.inner_layer.ProductLayer(child_layers: typing_extensions.List[Layer], edges: jax.experimental.sparse.BCOO)#

Bases: InnerLayer

A layer that represents the product of multiple other units.

child_layers: typing_extensions.List[typing_extensions.Union[SparseSumLayer, InputLayer]]#: The child of a product layer is a list that contains groups sum units with the same scope or groups of input units with the same scope.

edges: jaxtyping.Int[jax.experimental.sparse.BCOO, len(child_layers), number_of_nodes]#

The edges consist of a sparse matrix containing integers. The first dimension describes the edges for each child layer. The second dimension describes the edges for each node in the child layer. The integers are interpreted in such a way that n-th value represents a edge (n, edges[n]).

Nodes in the child layer can be mapped to by multiple nodes in this layer.

The shape is (#child_layers, #nodes).

validate()#: Validate the parameters and their layouts.

property number_of_nodes: int#

Returns:: The number of nodes in the layer.

classmethod nx_classes() → typing_extensions.Tuple[typing_extensions.Type, Ellipsis]#

Returns:: The tuple of matching classes of the layer in the probabilistic_model.probabilistic_circuit.rx package.

property number_of_components: int#

Returns:: The number of components (leaves + edges) of the entire circuit

variables() → jax.Array#

log_likelihood_of_nodes_single(x: jax.Array) → jax.Array#

Calculate the log-likelihood of the distribution.

Parameters:: x – The input vector.
Returns:: The log-likelihood of every node in the layer for x.

__deepcopy__(memo=None)#

Create a deep copy of the layer.

Parameters:: memo – A dictionary that is used to keep track of objects that have already been copied.

to_json() → typing_extensions.Dict[str, typing_extensions.Any]#

classmethod _from_json(data: typing_extensions.Dict[str, typing_extensions.Any]) → typing_extensions.Self#

Create a variable from a json dict. This method is called from the from_json method after the correct subclass is determined and should be overwritten by the respective subclass.

Parameters:: data – The json dict
Returns:: The deserialized object

classmethod create_layer_from_nodes_with_same_type_and_scope(nodes: typing_extensions.List[probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.Unit], child_layers: typing_extensions.List[NXConverterLayer], progress_bar: bool = True) → NXConverterLayer#: Create a layer from a list of nodes with the same type and scope.

to_nx(variables: sortedcontainers.SortedSet[random_events.variable.Variable], result: probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.ProbabilisticCircuit, progress_bar: typing_extensions.Optional[tqdm.tqdm] = None) → typing_extensions.List[probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.Unit]#

Convert the layer to a networkx circuit. For every node in this circuit, a corresponding node in the networkx circuit is created. The nodes all belong to the same circuit.

Parameters:

variables – The variables of the circuit.
result – The resulting circuit to write into
progress_bar – A progress bar to show the progress.

Returns:

The nodes of the networkx circuit.

class probabilistic_model.probabilistic_circuit.jax.inner_layer.NXConverterLayer#

Class used for conversion from a probabilistic circuit in networkx to a layered circuit in jax.

layer: Layer#

nodes: typing_extensions.List[probabilistic_model.probabilistic_circuit.rx.probabilistic_circuit.Unit]#

hash_remap: typing_extensions.Dict[int, int]#