def gemm_stride_checker(
    original_ta: TensorAccessor, dim: int, get_stride_at_dim: int = None
) -> bool:
    """
    Checks whether it's possible to get gemm stride correctly from original_ta
    given an input stride dim.

    This function should be called before actually invoking TensorAccessor.update_base_tensor()
    for gemm ops. If it returns False, we should avoid calling update_base_tensor().

    This is necessary because CUTLASS gemm ops have special "stride" params:
    batch_stride and row_stride. We need to make sure that after stride dim is updated
    in the TensorAccessor, these strides could still be populated successfully.
    """

    if not _dynamic_shape_checker(original_ta.original_shapes, dim):
        return False

    # Need to make sure that the new stride dim doesn't break
    # last dim's continuity.
    # This is because CUTLASS GEMM API assumes that GEMM stride
    # only operates on the last dim for row-major output.
    # For example, concatenations of GEMMs along dimensions to the right of the
    # original shape can't be fused. A particular case of this is when GEMM
    # output of shape (M, N) is unsqueezed to (M, N, 1) and concatenated with
    # another (M, N, 1).
    if not original_ta.is_rightmost_dim_contiguous(dim):
        return False

    if get_stride_at_dim is None:
        # The dim before the last dim
        get_stride_at_dim = len(original_ta.original_shapes) - 2

    tmp_ta = copy.deepcopy(original_ta)
    tmp_shape = copy.deepcopy(
        original_ta.actual_shapes
        if original_ta.actual_shapes is not None
        else original_ta.original_shapes
    )
    # Make sure new shape is different from the original shape.
    tmp_shape[dim] = IntImm(tmp_shape[dim].value() + 1)
    tmp_tensor = Tensor(shape=tmp_shape)
    tmp_ta.update_base_tensor(tmp_tensor, dim, stride_dim_offset=0)

    # TODO: Make this configurable for different gemms, bmms, etc.
    stride_strs = tmp_ta.try_get_stride_strs(get_stride_at_dim)
    if stride_strs is None:
        _LOGGER.debug(
            f"Failed in gemm_stride_checker: "
            f"dim: {dim}, "
            f"original_shapes length: {len(original_ta.original_shapes)}"
        )
        return False
    else:
        return True
Checks whether it's possible to get gemm stride correctly from original_ta
given an input stride dim.

This function should be called before actually invoking TensorAccessor.update_base_tensor()
for gemm ops. If it returns False, we should avoid calling update_base_tensor().

This is necessary because CUTLASS gemm ops have special "stride" params:
batch_stride and row_stride. We need to make sure that after stride dim is updated
in the TensorAccessor, these strides could still be populated successfully.
python/aitemplate/compiler/transform/transform_strided_ops_utils.py