src.model.deeplearn.layer.torf_output_mix_layer

Classes

TORFOutputMixLayer(*args, **kwargs)

class src.model.deeplearn.layer.torf_output_mix_layer.TORFOutputMixLayer(*args, **kwargs)
Author:

Alberto M. Esmoris Pena

Fused pre-pool per-point mixer for the TORF output head (intervention B). Replaces the 14-layer Concat+Pool+MLP chain in TransfOctoRFPwiseClassif.build_output_mix() with a single Keras Layer whose internal Dense sub-layers are owned as attributes.

The fused layout sidesteps a glibc heap-corruption SIGABRT triggered during process cleanup when many small per-instance layers from the out_mix_* namespace are destructed under high model-state pressure (e.g., DALES at leaf_voxel_length=1.0). The mathematical behavior is identical to the unfused version; only the layer topology changes.

Given \(\mathbf{H} \in \mathbb{R}^{B \times K \times n_h}\) and a per-neighbor float mask \(\mathbf{m} \in [0, 1]^{B \times K}\), the layer computes:

  1. Pool: a global summary \(\mathbf{g} \in \mathbb{R}^{B \times n_h}\) via either masked-mean (pool='mean') or masked-max (pool='max').

  2. Broadcast \(\mathbf{g}\) along K to give \(\mathbf{G} \in \mathbb{R}^{B \times K \times n_h}\).

  3. Concatenate \(\mathbf{C} = [\mathbf{H} ; \mathbf{G}] \in \mathbb{R}^{B \times K \times 2 n_h}\).

  4. Apply an MLP whose widths are configured via units and whose activation is activation. Optional BatchNorm between each Dense and the activation when batch_normalization is true.

  5. If the last MLP width does not equal \(n_h\), project to \(n_h\).

  6. Residual: \(\mathbf{Z} = \mathbf{H} + \mathbf{M}\) where \(\mathbf{M}\) is the MLP output.

Variables:

  • n_h (int) – Output hidden dimensionality.

  • K (int) – Number of K-neighbors (used for broadcast).

  • units (list[int]) – List of hidden MLP widths.

  • activation (str) – Nonlinear activation between Dense layers.

  • pool (str) – Pool strategy "mean" or "max".

  • use_bn (bool) – Whether to apply BatchNorm between Dense layers.

  • bn_momentum (float) – Momentum for BatchNorm (if used).

__init__(n_h, K, units, activation='relu', pool='mean', batch_normalization=False, bn_momentum=0.99, kernel_initializer='glorot_uniform', **kwargs)

See Layer and Layer.__init__().

build(input_shape)

Build all owned sublayers so that keras.models.load_model can restore their weights without the parent’s call() having been invoked first.

Parameters:

input_shape – List of two shapes [(B, K, n_h), (B, K)].

call(inputs, training=False, mask=False)

Apply the fused output-mix.

Parameters:

inputs – List [hidden, mask_float] of two tensors with shapes \((B, K, n_h)\) and \((B, K)\).

Returns:

Mixed tensor of shape \((B, K, n_h)\).

Return type:

keras tensor

compute_output_shape(input_shape)

Return the static output shape (B, K, n_h).

Parameters:

input_shape – List of two input shapes [(B, K, n_h), (B, K)].

Returns:

(B, K, n_h).

Return type:

tuple

get_config()

Return necessary data to deserialize the layer.

classmethod from_config(config)

Use given config data to deserialize the layer.