    def __call__(self, x: Tensor, dim: int = None) -> Tensor:
        """call the op

        Parameters
        ----------
        x : Tensor
            input tensor
        dim : int
            the dimension to be normalized.
            (default: None, in this case the input tensor will be treated as
             a 1-D tensor)

        Returns
        ----------
            Tensor
        """
        if dim is None:
            raise NotImplementedError(
                "flattening input tensor before normalization is not supported yet"
            )
        dim = wrap_dim(dim, x._rank())

        inner_dims = x.shape()[dim + 1 :]
        if not all(isinstance(d, IntImm) for d in inner_dims):
            raise NotImplementedError(
                "inner dims must all be static; {dim=}, {x.shape()=}"
            )

        self._attrs["inputs"] = [x]
        self._attrs["dim"] = dim
        self._set_depth()
        output_shape = self._infer_shapes(x)
        output = Tensor(output_shape, src_ops={self}, dtype=x.dtype())
        self._attrs["outputs"] = [output]
        return output
call the op

Parameters
----------
x : Tensor
input tensor
dim : int
the dimension to be normalized.
(default: None, in this case the input tensor will be treated as
a 1-D tensor)

Returns
----------
Tensor
python/aitemplate/compiler/ops/softmax/softmax.py