Method output_shape

axlearn/common/layers.py:822–836

(self, *, input_shape: Sequence[Optional[int]])

        cfg = self.config
        if len(input_shape) != 4:
            raise ValueError(f"We expect len(input_shape) = 4, but got {len(input_shape)}.")
        input_height, input_width = input_shape[1:3]

        if input_height is not None:
            output_height = max(input_height - cfg.window[0], 0) // cfg.strides[0] + 1
        else:
            output_height = None
        if input_width is not None:
            output_width = max(input_width - cfg.window[1], 0) // cfg.strides[1] + 1
        else:
            output_width = None
        return [input_shape[0], output_height, output_width, input_shape[3]]


class Embedding(BaseLayer):
    """Implements an embedding lookup function.

    Batched map for int in [0, <num_embeddings>) -> <dim> float vector.