    def __init__(
        self,
        total_length: IntVar,
        batch_dim: IntVar,
        jagged_dims: List[JaggedDim],
    ):
        """Initializes a JaggedIntVar.

        Parameters
        ----------
        total_length : IntVar
            The existing IntVar defining the total length sum_B(N_B) of the
            JaggedIntVar. The "name" and "values" attributes of the JaggedIntVar
            are the same as those of the total_length. This allows transparent
            treatment of the jagged Tensor as dense by non-jagged-aware ops.
            Must be a dynamic dim (IntVar, not IntImm).
        batch_dim : IntVar
            The batch dimension B in the sum_B(N_B) representation of the
            JaggedIntVar. Specifies the number of (outermost) variable-length
            sequences encoded within the JaggedIntVar. Must be a dynamic dim
            (IntVar, not IntImm).
        jagged_dims : List[JaggedDim]
            One or more jagged dimension encoded in the JaggedIntVar. Each
            JaggedDim specifies the bounds of one level of nested jaggedness
            of the JaggedIntVar. See the class docstring for details.
            The list must contain at least one JaggedDim. All JaggedDims
            in the list must have their offsets already set to the
            corresponding rank-1 Tensors.
        """
        if total_length is None or type(total_length) != IntVar:
            raise TypeError(
                "total_length must be dynamic (IntVar), "
                f"but given {type(total_length).__name__}."
            )
        if not jagged_dims or not all(
            isinstance(dim, JaggedDim) for dim in jagged_dims
        ):
            raise TypeError(
                "jagged_dims must be a non-empty list of JaggedDims, "
                f"but given {jagged_dims}."
            )
        offsets_types = set()
        for i, dim in enumerate(jagged_dims):
            if dim.offsets() is None:
                raise ValueError(
                    f"JaggedDim {i} in the jagged_dims list has no associated offsets. "
                    "This probably means that the JaggedIntVar is instantiated directly. "
                    "Instead, jagged Tensor must be created by calling the make_jagged op."
                )
            else:
                offsets_type = dim.offsets()._attrs["dtype"]
                if offsets_type not in ["int32", "int64"]:
                    raise TypeError(
                        "The offsets Tensors can be either int32 or int64, "
                        f"but given the Tensor of type {offsets_type}."
                    )
                offsets_types.add(offsets_type)
        if len(offsets_types) > 1:
            raise TypeError(
                "All offsets Tensors must be of the same type,"
                f" but given the Tensors of different types: {offsets_types}."
            )

        super().__init__(
            values=total_length._attrs["values"],
            name=total_length._attrs["name"],
            symbolic_value=total_length._attrs["symbolic_value"],
        )

        self._attrs["batch_dim"] = batch_dim
        self._attrs["jagged_dims"] = jagged_dims
        self._attrs["offsets_type"] = f"{offsets_types.pop()}_t"
        self._total_length = total_length
Initializes a JaggedIntVar.

Parameters
----------
total_length : IntVar
The existing IntVar defining the total length sum_B(N_B) of the
JaggedIntVar. The "name" and "values" attributes of the JaggedIntVar
are the same as those of the total_length. This allows transparent
treatment of the jagged Tensor as dense by non-jagged-aware ops.
Must be a dynamic dim (IntVar, not IntImm).
batch_dim : IntVar
The batch dimension B in the sum_B(N_B) representation of the
JaggedIntVar. Specifies the number of (outermost) variable-length
sequences encoded within the JaggedIntVar. Must be a dynamic dim
(IntVar, not IntImm).
jagged_dims : List[JaggedDim]
One or more jagged dimension encoded in the JaggedIntVar. Each
JaggedDim specifies the bounds of one level of nested jaggedness
of the JaggedIntVar. See the class docstring for details.
The list must contain at least one JaggedDim. All JaggedDims
in the list must have their offsets already set to the
corresponding rank-1 Tensors.
python/aitemplate/compiler/base.py