    def _infer_shapes(self, *args: Tensor) -> List[IntVar]:
        """Offline shape inference."

        Parameters
        ----------
        args : input tensors.

        Returns
        -------
        List[IntVar] : output tensor shape.
        """

        if len(args) == 0:
            raise RuntimeError(
                "Elementwise op {} doesn't have inputs!".format(self._attrs["func"])
            )

        _discover_implicit_jagged_inputs(args)

        dense_shapes = [arg._attrs["shape"] for arg in args if not arg.is_jagged()]
        jagged_shapes = [arg._attrs["shape"] for arg in args if arg.is_jagged()]

        max_dense_shape = _broadcast_dense_shapes(dense_shapes)
        if not jagged_shapes:
            return max_dense_shape

        max_jagged_shape = _broadcast_jagged_shapes(jagged_shapes)
        if not dense_shapes:
            return max_jagged_shape

        return _broadcast_dense_and_jagged_shape(max_dense_shape, max_jagged_shape)
Offline shape inference."

Parameters
----------
args : input tensors.

Returns
-------
List[IntVar] : output tensor shape.
python/aitemplate/compiler/ops/common/elementwise.py