tfm.uplift.utils.split_by_treatment

Splits a tensor into control and treatment tensors.

tfm.uplift.utils.split_by_treatment(
    values: tf.Tensor, is_treatment: tf.Tensor
) -> Tuple[tf.Tensor, tf.Tensor]

Args
values	a tf.Tensor of shape (D0, D1, ..., DN).
is_treatment	a tf.Tensor of shape (D0,) or (D0, 1) castable to boolean indicating if the example belongs to the treatment group (True) or control group (False).

Returns
A tuple with control and treatment values sliced by the is_treatment tensor.