tf.layers.PoolerLayer module

class tf.layers.PoolerLayer.PoolerLayer(*args: Any, **kwargs: Any)

Bases: modelzoo.common.tf.layers.BaseLayer.BaseLayer

The pooler layer.

Currently supports the following pooler types:

• "mean": Mean reduction.

• "max": Max reduction.

• "first": First slice in the axis dimension.

• "last": Last slice in the axis dimension.

• "sum": Takes the sum over the axis dimension. Defaults to the entire Tensor.

• None: No pooling (output=input).

call(inputs, padding_mask=None, **kwargs)

Apply pooler of a given type.

Takes in a padding mask with 1s for tokens and 0s for padding.