    def _infer_shapes(self, a: Tensor, b: Tensor, bias: Tensor):
        """Infers output shapes from input tensors."""

        bias_shape = bias._attrs["shape"]
        if len(bias_shape) != 1:
            raise RuntimeError("Bias should be 1D vector ")
        bias_shape_value = bias_shape[0]._attrs["values"]
        if len(bias_shape_value) != 1:
            raise RuntimeError("Bias should be fixed 1D vector")
        bias_dim = bias_shape_value[0]
        outshape = super()._infer_shapes(a, b)
        if outshape[1]._attrs["values"][0] != bias_dim:
            raise RuntimeError("GEMM/Bias shape doesn't match")
        return outshape
Infers output shapes from input tensors.
python/aitemplate/compiler/ops/gemm_epilogue_vistor/gemm_rcr_bias_softmax.py