def _get_alignments_and_broadcast_sizes(
    dtype: str,
    input_accessors: List[TensorAccessor],
    output_accessors: List[TensorAccessor],
    mixed_jagged_dense_indexing: bool,
    output_volume: Optional[List[IntVar]],
) -> Tuple[List[int], List[Optional[List[IntVar]]]]:
    """
    Returns Tuple(input_alignments, input_broadcast_sizes)
    """
    # Handle input broadcast.
    output_shape = output_accessors[0].original_shapes

    input_broadcast_sizes = _get_input_broadcast_sizes(
        input_accessors, output_accessors, mixed_jagged_dense_indexing, output_volume
    )

    # In the mixed jagged / dense indexing case, the number of the
    # rightmost non-broadcated static dimensions of the dense inputs
    # to be considered for vectorization can't be larger than the
    # number of the jagged output's inner dimensions (i.e., the
    # dimensions following the JaggedIntVar). Otherwise, there may
    # be an overlap with the jagged dimensions, in which case the
    # vectorization can break.
    max_num_rightmost_dims_considered_for_alignments = (
        len(output_shape) - 1 if mixed_jagged_dense_indexing else len(output_shape)
    )

    # We do not support mixed input / output alignments except for last dim broadcast.
    # The global_max_alignment is the min value of:
    #     1) input shape alignments (with input broadcast in consideration);
    #     2) input tensor accessor alignments (strides, offsets);
    #     3) output shape alignments;
    #     4) output tensor accessor alignments (strides, offsets);

    # Now calculate global_max_alignment based on 2), 3) and 4) first.
    global_max_alignment = min(
        alignment_utils.find_max_alignment(
            shape_utils.get_num_rightmost_static_elements(
                output_shape, max_num_rightmost_dims_considered_for_alignments
            ),
            dtype,
        ),
        tensor_accessor_codegen.find_max_alignment_for_accessors(
            dtype, input_accessors
        ),
        tensor_accessor_codegen.find_max_alignment_for_accessors(
            dtype, output_accessors
        ),
    )

    # Now calculate global_max_alignment based on 1).
    # Also calculate input alignments.
    input_alignments = _get_input_alignments(
        input_accessors,
        input_broadcast_sizes,
        max_num_rightmost_dims_considered_for_alignments,
        output_shape,
        dtype,
        global_max_alignment,
    )
    return input_alignments, input_broadcast_sizes
Returns Tuple(input_alignments, input_broadcast_sizes)
python/aitemplate/backend/common/elementwise_common.py