def _discover_implicit_jagged_inputs(inputs: List[Tensor]):
    """
    Convert implicit jagged Tensor inputs into explicit jagged Tensors.

    There may be cases when elementwise has both explicit jagged Tensor
    inputs (i.e. with a JaggedIntVar as the first dimension in the shape)
    and "implicit" jagged Tensor inputs (i.e. dense Tensors with the first
    dimension == the JaggedIntVar.total_length() in the jagged Tensor
    inputs). Here we detect such implicit jagged Tensor inputs and replace
    the total_length: IntVar in the dense input's shape by the corresponding
    JaggedIntVar from the jagged input's shape. Importantly, this must be
    done before the mixed jagged / dense broadcasting takes place.
    """
    total_length_map = {}
    for tensor in inputs:
        if tensor.is_jagged():
            jagged_int_var = tensor._attrs["shape"][0]
            total_length = jagged_int_var.total_length()
            total_length_map[total_length] = jagged_int_var

    if total_length_map:
        # there are explicit jagged Tensors among the inputs:
        # we check if there are implict ones and make them explicit
        for tensor in inputs:
            shape = tensor._attrs["shape"]
            if not tensor.is_jagged() and shape and not isinstance(shape[0], IntImm):
                if shape[0] in total_length_map:
                    # the dense Tensor input's first dimension is the total_length
                    # dimension in the JaggedIntVar of one of the jagged Tensor
                    # inputs: we replace the dense Tensor input's first dimension
                    # by the corresponding JaggedIntVar, hence giving it a
                    # jagged Tensor semantics for further processing.
                    shape[0] = total_length_map[shape[0]]
Convert implicit jagged Tensor inputs into explicit jagged Tensors.

There may be cases when elementwise has both explicit jagged Tensor
inputs (i.e. with a JaggedIntVar as the first dimension in the shape)
and "implicit" jagged Tensor inputs (i.e. dense Tensors with the first
dimension == the JaggedIntVar.total_length() in the jagged Tensor
inputs). Here we detect such implicit jagged Tensor inputs and replace
the total_length: IntVar in the dense input's shape by the corresponding
JaggedIntVar from the jagged input's shape. Importantly, this must be
done before the mixed jagged / dense broadcasting takes place.
python/aitemplate/compiler/ops/common/elementwise.py