tf.PoolerLayer module

tf.PoolerLayer module

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

Bases: modelzoo.common.layers.tf.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.