class DimensionType(Type):
    pass

# I guess all dimensions have the same type?
DimType = DimensionType()

class Dimension(Variable):
    def __init__(self, name=None, length=None):
        pass

    def length(self) -> TensorVariable:
        ...

class DimOp(Op):
    pass

class DimConstant(DimensionVariable):
    def __init__(self, name, *, length=None):
        pass

# The result of stacking two dims
class Product(DimOp):
    __props__ = ("name",)

    def make_node(self, *inputs: DimensionVariable):
        if self.name is not None:
            name = self.name
        else:
            name = f"product[{','.join(input.name for input in inputs)}]"
        output = Dimension(name=name, length=prod(input.length for input in inputs))
        return Apply(self, inputs, [output])

def dim(name, *, length=None):
    return DimConstant(name, length=length)

def stack(variable, *, dims, name=None):
    dim_op = Product(name)
    stacked_dim = dim_op(dims)
    ...

country = pt.dim("country")
treatment = pt.dim("treatment")

# We can talk about the stacked dim directly:
interaction = pt.stacked_dim(country, treatment)
effect = xtensor(dim=interaction)
assert effect.unstack(interaction).dims == (country, treatment)
assert pt.stack(effect, [country, treatment]).dims == interaction

# Same for indices and slices...
subset = pt.dim_slice(country, slice("A", "B"))
sub_effect = pt.xtensor(dims=subset)
effect = pt.xtensor(dims=country)
assert effect.sel(country=slice("A", "B")).dims == (subset,)
effect = effect.at.sel(country=slice("A", "B")).add(sub_effect)