def METHOD_NAME(bb: BlockBuilder, call: Call) -> Expr:
    def te_expand_dims(data, axis):
        data_relax = relax.Var("data", relax.TensorStructInfo(data.shape))
        f_infer_sinfo = call.op.get_attr("FInferStructInfo")
        output_shape = f_infer_sinfo(relax.op.expand_dims(data_relax, axis), bb).shape
        output_ndim = len(output_shape)
        data_dims = []
        for i in range(output_ndim):
            if i not in axis and (i - output_ndim) not in axis:
                data_dims.append(i)
        return te.compute(
            output_shape,
            lambda *idx: data(*[idx[dim] for dim in data_dims]),
            name="expand_dims",
        )
    return bb.call_te(
        te_expand_dims, call.args[0], call.attrs.axis, primfunc_name_hint="expand_dims"
    )