Advanced numerical embedding layer for continuous features.
This layer embeds each continuous numerical feature into a higher-dimensional space by combining two- branches: 1. Continuous Branch: Each feature is processed via a small MLP (using TimeDistributed layers). 2. Discrete Branch: Each feature is discretized into bins using learnable min/max boundaries and then an embedding is looked up for its bin.
A learnable gate (of shape (num_features, embedding_dim)) combines the two branch outputs per feature and per embedding dimension. Additionally, the continuous branch uses a residual connection and optional batch normalization to improve training stability.
The layer supports inputs of shape (batch, num_features) for any number of features and returns outputs of shape (batch, num_features, embedding_dim).
embedding_dim (int): Output embedding dimension per feature.
mlp_hidden_units (int): Hidden units for the continuous branch MLP.
num_bins (int): Number of bins for discretization.
init_min (float or list): Initial minimum values for discretization boundaries. If a scalar is
provided, it is applied to all features.
init_max (float or list): Initial maximum values for discretization boundaries.
dropout_rate (float): Dropout rate applied to the continuous branch.
use_batch_norm (bool): Whether to apply batch normalization to the continuous branch.
__init__(self, embedding_dim: int = 8, mlp_hidden_units: int = 16, num_bins: int = 10, init_min: float | list[float] = -3.0, init_max: float | list[float] = 3.0, dropout_rate: float = 0.1, use_batch_norm: bool = True, **kwargs)Initialize the NumericalEmbedding layer.
- mlp_hidden_units: Number of hidden units in the MLP.
- num_bins: Number of bins for discretization.
- init_min: Minimum value(s) for initialization. Can be a single float or list of floats.
- init_max: Maximum value(s) for initialization. Can be a single float or list of floats.
- dropout_rate: Dropout rate for regularization.
- use_batch_norm: Whether to use batch normalization. **kwargs: Additional layer arguments.
add_loss(self, loss)Can be called inside of the call() method to add a scalar loss.
class- **MyLayer(Layer)**: ...
def call(self, x):
self.add_loss(ops.sum(x))
return xadd_variable(self, shape, initializer, dtype=None, trainable=True, autocast=True, regularizer=None, constraint=None, name=None)Add a weight variable to the layer.
Alias of add_weight().
add_weight(self, shape=None, initializer=None, dtype=None, trainable=True, autocast=True, regularizer=None, constraint=None, aggregation='none', overwrite_with_gradient=False, name=None)Add a weight variable to the layer.
(no `None` entries). Defaults to `()` (scalar) if unspecified.
- initializer: Initializer object to use to populate the initial
variable value, or string name of a built-in initializer
(e.g.
"random_normal"). If unspecified, defaults to"glorot_uniform"for floating-point variables and to"zeros"for all other types (e.g. int, bool). - dtype: Dtype of the variable to create, e.g.
"float32". If unspecified, defaults to the layer's variable dtype (which itself defaults to"float32"if unspecified). - trainable: Boolean, whether the variable should be trainable via
backprop or whether its updates are managed manually. Defaults
to
True. - autocast: Boolean, whether to autocast layers variables when
accessing them. Defaults to
True. - regularizer: Regularizer object to call to apply penalty on the
weight. These penalties are summed into the loss function
during optimization. Defaults to
None. - constraint: Contrainst object to call on the variable after any
optimizer update, or string name of a built-in constraint.
Defaults to
None. - aggregation: Optional string, one of
None,"none","mean","sum"or"only_first_replica". Annotates the variable with the type of multi-replica aggregation to be used for this variable when writing custom data parallel training loops. Defaults to"none". - overwrite_with_gradient: Boolean, whether to overwrite the variable
with the computed gradient. This is useful for float8 training.
Defaults to
False. - name: String name of the variable. Useful for debugging purposes.
build(self, input_shape) -> NoneBuild the layer's weights for a given input shape.
build_from_config(self, config)Builds the layer's states with the supplied config dict.
By default, this method calls the build(config["input_shape"]) method,
which creates weights based on the layer's input shape in the supplied
config. If your config contains other information needed to load the
layer's state, you should override this method.
call(self, inputs: tensorflow.python.framework.tensor.Tensor, training: bool = False) -> tensorflow.python.framework.tensor.TensorApply the layer to its inputs.
- training: Whether the layer is being called in training mode.
The transformed tensor.
count_params(self)Count the total number of scalars composing the weights.
An integer count.
get_build_config(self)Returns a dictionary with the layer's input shape.
This method returns a config dict that can be used by
build_from_config(config) to create all states (e.g. Variables and
Lookup tables) needed by the layer.
By default, the config only contains the input shape that the layer was built with. If you're writing a custom layer that creates state in an unusual way, you should override this method to make sure this state is already created when Keras attempts to load its value upon model loading.
A dict containing the input shape associated with the layer.
get_config(self) -> dictReturn the configuration needed to re-create this layer.
The layer configuration.
get_weights(self)Return the values of layer.weights as a list of NumPy arrays.
load_own_variables(self, store)Loads the state of the layer.
You can override this method to take full control of how the state of
the layer is loaded upon calling keras.models.load_model().
rematerialized_call(self, layer_call, *args, **kwargs)Enable rematerialization dynamically for layer's call method.
Rematerialized layer's `call` method.
save_own_variables(self, store)Saves the state of the layer.
You can override this method to take full control of how the state of
the layer is saved upon calling model.save().
set_weights(self, weights)Sets the values of layer.weights from a list of NumPy arrays.
stateless_call(self, trainable_variables, non_trainable_variables, *args, return_losses=False, **kwargs)Call the layer without any side effects.
- non_trainable_variables: List of non-trainable variables of the
model.
*args: Positional arguments to be passed to
call(). - return_losses: If
True,stateless_call()will return the list of losses created duringcall()as part of its return values. **kwargs: Keyword arguments to be passed tocall().
A tuple. By default, returns `(outputs, non_trainable_variables)`.
If `return_losses = True`, then returns
`(outputs, non_trainable_variables, losses)`.
- Note:
non_trainable_variablesinclude not only non-trainable weights such asBatchNormalizationstatistics, but also RNG seed state (if there are any random operations part of the layer, such as dropout), andMetricstate (if there are any metrics attached to the layer). These are all elements of state of the layer.
model = ...
data = ...
trainable_variables = model.trainable_variables
non_trainable_variables = model.non_trainable_variables
# Call the model with zero side effects
outputs, non_trainable_variables = model.stateless_call(
trainable_variables,
non_trainable_variables,
data,
)
# Attach the updated state to the model
# (until you do this, the model is still in its pre-call state).
for ref_var, value in zip(
model.non_trainable_variables, non_trainable_variables
):
ref_var.assign(value)