    def _infer_shapes(self, inputs: List[Tensor], dim) -> List[IntVar]:
        """Infers shapes for concatenate."""
        concatenate.check_rank(inputs, dim)
        rank = concatenate.get_rank(inputs)
        # all inputs are empty
        if rank is None:
            return [IntImm(0)]

        ref_input, _ = concatenate.get_first_non_empty_input_if_any(inputs)
        # reference shape should come from a non-empty tensor
        ref_input_shape = ref_input._attrs["shape"]
        input_shapes = []
        for t in inputs:
            if shape_utils.is_empty_rank1_tensor(t._attrs["shape"]):
                shape = deepcopy(ref_input_shape)
                shape[dim] = IntImm(0)
            else:
                shape = t._attrs["shape"]
            input_shapes.append(shape)
        output_shape = []
        input_shape_values = [
            [d._attrs["values"] for d in shape] for shape in input_shapes
        ]
        for idx, lst in enumerate(zip(*input_shape_values)):
            if idx == dim:
                min_value_sum = sum(value[0] for value in lst)
                max_value_sum = sum(value[-1] for value in lst)
                sym_val = reduce(
                    lambda x, y: x + y,
                    [
                        input_shape[idx]._attrs["symbolic_value"]
                        for input_shape in input_shapes
                    ],
                )
                shape_var = shape_utils.gen_int_var(
                    [min_value_sum, max_value_sum], symbolic_value=sym_val
                )
                output_shape.append(shape_var)
            else:
                output_dim = ref_input_shape[idx]
                for shape in input_shapes:
                    # the corresponding input tensor is empty
                    if shape_utils.is_empty_rank1_tensor(shape):
                        continue
                    # if output_dim != shape[idx]:
                    if output_dim._attrs["values"] != shape[idx]._attrs["values"]:
                        raise RuntimeError(
                            "tensors expected to have the same dimensions "
                            "except concat_dim! dim: {}, shape1: {}, shape2: {}, inputs: {}".format(
                                idx, output_dim, shape[idx], inputs
                            )
                        )
                output_shape.append(output_dim)
        return output_shape
Infers shapes for concatenate.
python/aitemplate/compiler/ops/tensor/concatenate.py