tfm.uplift.layers.uplift_networks.two_tower_output_head.TwoTowerOutputHead

Two tower training and inference output computation.

tfm.uplift.layers.uplift_networks.two_tower_output_head.TwoTowerOutputHead(
    treatment_indicator_feature_name: str,
    uplift_network: base_uplift_networks.BaseTwoTowerUpliftNetwork,
    inverse_link_fn: (Callable[[tf.Tensor], tf.Tensor] | None) = None,
    **kwargs
)

This layer is intended to be used in conjunction with a two tower uplift network layer to compute training and inference outputs. It passes the input dictionary of feature tensors to the uplift network and computes training and inference tensors from the uplift network's outputs.

The control, treatment, and uplift predictions are computed from the uplift network's control and treatment logits. Additionally, if the treatment indicator tensor is part of the inputs, the true logits are also computed to be used for loss computation. The treatment indicator tensor should therefore be present in the inputs dictionary during training and evaluation.

treatment_indicator_feature_name the name of the feature representing the treatment indicator tensor, which should be castable to a boolean tensor (False for control and True for treatment). When this tensor is part of the layer inputs, additional true logits are returned for the loss computation. In this case the layer returns a TwoTowerTrainingOutputs instance, otherwise a TwoTowerPredictionOutputs instance is returned.
uplift_network a layer for computing control and treatment logits. Its input is expected to be a dictionary of feature tensors and its output is exptected to be a TwoTowerNetworkOutputs instance.
inverse_link_fn a function for computing the control and treatment predictions from their respective logits. If left as None it is functionally equivalent to the idendity function.
**kwargs base layer keyword arguments.

Methods

call

View source

call(
    inputs: types.DictOfTensors,
    training: (bool | None) = None,
    mask: (tf.Tensor | None) = None
) -> (types.TwoTowerPredictionOutputs | types.TwoTowerTrainingOutputs)

Computes two tower inference and training outputs.

Args
inputs feature tensors to be passed to the uplift network. For training and evaluation, the treatment indicator tensor must be given as part of the inputs.
training optional boolean training flag to pass to the uplift network.
mask optional tensor mask to pass to the uplift network.

Returns
TwoTowerTrainingOutputs when the treatment indicator tensor is part of the input features, otherwise TwoTowerPredictionOutputs.