def linear_inference_rule(n: Node, module_instance, symbols, constraints, counter):
"""
Input and output sizes should be the same except for the last dimension
If the input is Dyn, then so should the output
"""
assert isinstance(n.args[0], Node)
linear_output, counter = gen_tvar(counter)
symbols[n] = linear_output
linear_input = symbols[n.args[0]]
input_dyn = BinConstraintT(linear_input, Dyn, op_eq)
output_dyn = BinConstraintT(linear_output, Dyn, op_eq)
c1 = Conj([input_dyn, output_dyn])
c2 = []
for i in range(1, MAX_TENSOR_RANK + 1):
new_dims_rhs_1, counter = gen_tensor_dims(i, counter)
new_dims_rhs_2, counter = gen_tensor_dims(i, counter)
nat_constraints = gen_nat_constraints(new_dims_rhs_1 + new_dims_rhs_2)
c_tensor_i = Conj([BinConstraintT(linear_input, TensorType(new_dims_rhs_1), op_eq),
BinConstraintT(linear_output, TensorType(new_dims_rhs_2), op_eq)] +
add_linear_constraints(new_dims_rhs_1, new_dims_rhs_2, module_instance) +
nat_constraints)
c2.append(c_tensor_i)
return [Disj([c1, Disj(c2)])], counter
def conv2d_inference_rule(n: Node, module_instance, symbols, constraints, counter):
assert isinstance(n.args[0], Node)
my_conv, counter = gen_tvar(counter)
symbols[n] = my_conv
input_var = symbols[n.args[0]]
# dim vars
[d1, d2, d3, d4], counter = gen_tensor_dims(MAX_TENSOR_RANK, counter)
# c1 = Matching(input_var, TensorType([d1, d2, d3, d4]))
c1 = BinConstraintT(input_var, TensorType([d1, d2, d3, d4]), op_matching)
# c2 = DConsistency(module_instance.in_channels, d2)
c2 = BinConstraintD(module_instance.in_channels, d2, op_consistency)
c3 = CalcConv(my_conv, input_var,
module_instance.out_channels,
module_instance.kernel_size,
module_instance.padding,
module_instance.stride,
module_instance.dilation, [d1, d2, d3, d4])
nat_constraints = gen_nat_constraints([d1, d2, d3, d4])
return [c1, c2, c3, *nat_constraints], counter
def cumsum_inference_rule(n: Node, symbols, constraints, counter):
"""
Input and output shapes should be equal
We should verify that the index is valid
"""
assert isinstance(n.args[0], Node)
arg_1 = n.args[1] if len(n.args) > 1 else n.kwargs["dim"]
assert isinstance(arg_1, int)
output, counter = gen_tvar(counter)
symbols[n] = output
input = symbols[n.args[0]]
input_dyn = BinConstraintT(input, Dyn, op_eq)
output_dyn = BinConstraintT(output, Dyn, op_eq)
c1 = Conj([input_dyn, output_dyn])
c2 = []
for i in range(1, MAX_TENSOR_RANK + 1):
new_dims, counter = gen_tensor_dims(i, counter)
nat_constraints = gen_nat_constraints(new_dims)
c_tensor_i = Conj([BinConstraintT(input, TensorType(new_dims), op_eq),
BinConstraintT(output, TensorType(new_dims), op_eq)] +
[range_check(arg_1, i)] + nat_constraints)
c2.append(c_tensor_i)
dyn_or_tensor = Disj([c1, Disj(c2)])
return [dyn_or_tensor], counter
def embedding_inference_rule(n: Node, module_instance, symbols, constraints, counter):
"""
The output shape differs from the input shape in the last dimension
"""
assert isinstance(n.args[0], Node)
embedding_dim = module_instance.embedding_dim # number
embedding_output, counter = gen_tvar(counter)
symbols[n] = embedding_output
embedding_input = symbols[n.args[0]]
input_dyn = BinConstraintT(embedding_input, Dyn, op_eq)
output_dyn = BinConstraintT(embedding_output, Dyn, op_eq)
c1 = Conj([input_dyn, output_dyn])
c2 = []
for i in range(1, MAX_TENSOR_RANK):
new_dims, counter = gen_tensor_dims(i, counter)
nat_constraints = gen_nat_constraints(new_dims)
# we consider all tensor sizes and append embedding_dim to the end of the output dimension in all cases
c_tensor_i = Conj([BinConstraintT(embedding_input, TensorType(new_dims), op_eq),
BinConstraintT(embedding_output, TensorType(new_dims + [embedding_dim]), op_eq)] +
nat_constraints)
c2.append(c_tensor_i)
return [Disj([c1, Disj(c2)])], counter
def gen_embedding_rules(n: Node, symbols, embedding_dim, counter):
embedding_output, counter = gen_tvar(counter)
symbols[n] = embedding_output
embedding_input = symbols[n.args[0]]
input_dyn = BinConstraintT(embedding_input, Dyn, op_eq)
output_dyn = BinConstraintT(embedding_output, Dyn, op_eq)
c1 = Conj([input_dyn, output_dyn])
c2 = []
for i in range(1, MAX_TENSOR_RANK):
new_dims, counter = gen_tensor_dims(i, counter)
nat_constraints = gen_nat_constraints(new_dims)
# we consider all tensor sizes and append embedding_dim to the end of the output dimension in all cases
c_tensor_i = Conj([
BinConstraintT(embedding_input, TensorType(new_dims), op_eq),
BinConstraintT(embedding_output,
TensorType(new_dims + [embedding_dim]), op_eq)
] + nat_constraints)
c2.append(c_tensor_i)
return [Disj([c1, Disj(c2)])], counter
def expand_inference_rule(n: Node, symbols, constraints, counter):
"""
We generate the exact constraints as we do for tensor additions but we constraint
the rank of this expression to be equal to len(n.args[1:]) so that only
those cases get considered for the output
"""
assert isinstance(n.args[0], Node)
# define the output for expand
expand, counter = gen_tvar(counter)
symbols[n] = expand
# since we do not have two nodes here, we will construct an argument variable
e1 = symbols[n.args[0]]
e2, counter = gen_tvar(counter)
e2_nat_constraints = []
for arg in n.args[1:]:
assert isinstance(arg, Node) or isinstance(arg, int)
if isinstance(arg, Node):
assert isinstance(symbols[arg], DVar)
e2_nat_constraints.append(BinConstraintD(0, symbols[arg], op_leq))
e2_constraint = BinConstraintT(e2, TensorType([arg if isinstance(arg, int) else symbols[arg] for arg in n.args[1:]]), op_eq)
constraints, counter = gen_broadcasting_constraints(e1, e2, symbols, counter, expand)
# constraint the output size
dims, counter = gen_tensor_dims(len(n.args[1:]), counter)
nat_constraints = gen_nat_constraints(dims)
c = [BinConstraintT(expand, TensorType(dims), op_eq), *nat_constraints, e2_constraint, *e2_nat_constraints]
constraints += c
return constraints, counter
def gen_consistency_constraints(constraint: Constraint, counter: int):
"""
Args:
constraint: Consistency constraint on tensors
counter: for variable tracking
Returns: Equality and consistency constraints on dimensions
"""
all_constraints = []
for i in range(1, MAX_TENSOR_RANK + 1):
new_dims_rhs_1, counter = gen_tensor_dims(i, counter)
new_dims_rhs_2, counter = gen_tensor_dims(i, counter)
nat_constraints = gen_nat_constraints(new_dims_rhs_1 + new_dims_rhs_2)
c_tensor_i = Conj([
BinConstraintT(constraint.lhs, TensorType(new_dims_rhs_1), op_eq),
BinConstraintT(constraint.rhs, TensorType(new_dims_rhs_2), op_eq)
] + [
BinConstraintD(d1, d2, op_consistency)
for d1, d2 in zip(new_dims_rhs_1, new_dims_rhs_2)
] + nat_constraints)
all_constraints.append(c_tensor_i)
return all_constraints, counter
def transform_get_item(constraint, counter):
"""
generate an equality of the form:
t = [a1, ..., an]
then generate constraints that check if the given index is valid
given this particular tensor size.
If the index is valid, generate a constraint to get the item
Note that we already handled the Dyn input case in the previous
step.
Args:
constraint: GetItem which assumes we are getting an item from a tensor (not Dyn)
counter: variable tracking
Returns: simplified constraints for GetItem
"""
dims, counter = gen_tensor_dims(constraint.tensor_size, counter)
nat_constraints = gen_nat_constraints(dims)
is_valid_index = valid_index(constraint.index, dims)
all_constraints = [
BinConstraintT(constraint.input_var, TensorType(dims), op_eq),
*nat_constraints, is_valid_index
]
# if the index is valid, we generate a constraint for getting an item
# otherwise this clause will have been UNSAT due to the wrong index
if is_valid_index == T():
all_constraints.append(
BinConstraintD(constraint.res, dims[constraint.index], op_eq))
return Conj(all_constraints), counter
def layer_norm_inference_rule(n: Node, module_instance, symbols, constraints, counter):
"""
Input and output shapes should be equal.
Input should be consistent with the normalized_shape
"""
assert isinstance(n.args[0], Node)
output, counter = gen_tvar(counter)
symbols[n] = output
input = symbols[n.args[0]]
input_dyn = BinConstraintT(input, Dyn, op_eq)
output_dyn = BinConstraintT(output, Dyn, op_eq)
c1 = Conj([input_dyn, output_dyn])
c2 = []
for i in range(1, MAX_TENSOR_RANK + 1):
new_dims_rhs, counter = gen_tensor_dims(i, counter)
nat_constraints = gen_nat_constraints(new_dims_rhs)
c_tensor_i = Conj([BinConstraintT(input, TensorType(new_dims_rhs), op_eq),
BinConstraintT(output, TensorType(new_dims_rhs), op_eq)] +
add_layer_norm_constraints(new_dims_rhs, list(module_instance.normalized_shape)) +
nat_constraints)
c2.append(c_tensor_i)
return [Disj([c1, Disj(c2)])], counter
def transform_get_item_tensor(constraint, counter):
"""
When the index is a tuple, then the output will be a tensor
TODO: we have to check if this is the case for all HF models
The cases we are covrering here are a tuple with one of:
- slice with default argument
- None
None appends 1 to the input tensor dimensions
so each occurrence of 'None' increases the rank by 1
slice with default arguments does not change the rank
"""
assert isinstance(constraint.index_tuple, tuple)
# generate a result tensor of the expected size
dims, counter = gen_tensor_dims(constraint.tensor_size, counter)
nat_constraints = gen_nat_constraints(dims)
# generate a place-holder list of the right rank
# where "slice" does not contribute to the rank and "None" does
none_c = constraint.index_tuple.count(None)
resulting_tensor_dims = (none_c + len(dims)) * [None]
dim_index = 0
for i in range(len(constraint.index_tuple)):
# append 1 to the right location of the resulting tensor
if constraint.index_tuple[i] is None:
resulting_tensor_dims[i] = 1
elif constraint.index_tuple[i] == slice(None, None, None):
pass
else:
raise NotImplementedError('Method not yet implemented')
# append the remaining dimensions to the right location
dim_index = 0
for i in range(len(resulting_tensor_dims)):
if resulting_tensor_dims[i] is None:
resulting_tensor_dims[i] = dims[dim_index]
dim_index += 1
# check if the index is valid
is_valid_index = valid_index_tensor(constraint.index_tuple, dims)
# check if the resulting tensor is within bounds
if len(resulting_tensor_dims) > 4:
return F(), counter
else:
constraints = [
BinConstraintT(constraint.input_var, TensorType(dims), op_eq),
BinConstraintT(constraint.res, TensorType(resulting_tensor_dims),
op_eq), *nat_constraints, is_valid_index
]
return Conj(constraints), counter
def generate_flatten_constraints(start_dim, end_dim, input, flattened, n, counter):
d, counter = gen_tensor_dims(n, counter)
c1 = BinConstraintT(input, TensorType(d), op_eq)
start_dim = n if start_dim == -1 else abs(start_dim)
end_dim = n + end_dim + 1 if end_dim < 0 else end_dim + 1
c2 = CalcProduct(start_dim, end_dim, flattened, d)
nat_constraints = gen_nat_constraints(d)
return Conj([c1, c2, *nat_constraints]), counter
def apply_padding(e1_var: TVar, e11: BinConstraintT, e2: BinConstraintT,
e12: BinConstraintT, d2: List[DVar], d11: List[DVar],
d12: List[DVar], counter: int):
"""
We are considering the possibility where one input has less dimensions than
another input, so we apply padding to the broadcasted results
Args:
e1_var: Variable representing the first input where padding will be
e11: constraint of the form e11 = Tensortype[d1, ..., dn]
e2: constraint of the form e2 = Tensortype[d1, ..., dn]
e12: constraint of the form e11 = Tensortype[d1, ..., dn]
d2: Tensor variables for the second input
d11: Tensor variables for the broadcasted first input
d12: Tensor variables for the broadcasted second input
counter: variable tracking
Returns: A new constraint whose goal is to apply padding to the broadcasted result
"""
res = []
# pad the shorter input with None so we can pass it to the broadcasting helper function
for i in range(1, len(d2)):
d1, counter = gen_tensor_dims(i, counter)
nat_constraints = gen_nat_constraints(d1 + d2 + d11 + d12)
e1 = BinConstraintT(e1_var, TensorType(d1), op_eq)
simulate_padding = [None] * (len(d2) - i)
assert len(simulate_padding + d1) == len(d2)
broadcast_padding = []
# for every padding size, we also consider broadcasting
for j in range((len(d2) - i)):
broadcast_padding.append(
broadcast_dim(simulate_padding, d2, d11, d12, j, True))
# we consider the possibilities for broadcasting for every dimension. Since we already
# padded d1, we do not consider it while broadcasting
all_broadcasting_possibilities = generate_all_broadcasting_possibilities_no_padding(
d1, d2[(len(d2) - i):], d11[(len(d2) - i):], d12[(len(d2) - i):])
# combine all constraints into a conjunction
c = Conj([
e1, e11, e2, e12, *broadcast_padding,
all_broadcasting_possibilities, *nat_constraints
])
res.append(c)
return Disj(res), counter
def gen_broadcasting_constraints(e1: TVar, e2: TVar, e11: TVar, e12: TVar,
i: int, counter: int):
"""
Simulates broadcasting on e1 and e2 and returns the results
respectively in e11 and e12. Because of gradual types,
e1 and e2 may not be equal. Similarly, e11 and e12 may not
be equal. e11 and e12 should be guaranteed to be consistent
as they represent the shapes of the tensors to be added after
broadcasting.
Args:
e1: TVar representing the type of input 1
e2: TVar representing the type of input 2
e11: TVar representing the representing broadcasted input 1
e12: TVar representing the representing broadcasted input 2
i: The rank of the resulting type of addition
counter: for variable tracking
Returns: Simplified broadcasting constraints
"""
dims, counter = gen_lists_of_dims(4, i, counter)
[d1, d2, d3, d4] = dims
nat_dims_i = gen_nat_constraints(list(itertools.chain(*dims)))
initialize_tensors_constraints = create_equality_constraints_for_broadcasting(
e1, e2, e11, e12, d1, d2, d3, d4)
[e1_tensor, e11_tensor, e2_tensor,
e12_tensor] = initialize_tensors_constraints
# without padding, broadcast all possibilities for tensors of size i
final_tensor_constraint_no_padding = Conj([
*initialize_tensors_constraints,
generate_all_broadcasting_possibilities_no_padding(d1, d2, d3, d4)
])
# with padding, broadcast all possibilities for tensors of size i
final_tensor_constraint_padding_arg1, counter = \
apply_padding(e1, e11_tensor, e2_tensor, e12_tensor, d2, d3, d4, counter)
final_tensor_constraint_padding_arg2, counter = \
apply_padding(e2, e12_tensor, e1_tensor, e11_tensor, d1, d4, d3, counter)
return final_tensor_constraint_no_padding, \
final_tensor_constraint_padding_arg1, \
final_tensor_constraint_padding_arg2, nat_dims_i, counter
def maxpool_inference_rule(n: Node, module_instance, symbols, constraints, counter):
assert isinstance(n.args[0], Node)
maxpool, counter = gen_tvar(counter)
symbols[n] = maxpool
input_var = symbols[n.args[0]]
# dim vars
[d1, d2, d3, d4], counter = gen_tensor_dims(MAX_TENSOR_RANK, counter)
c1 = BinConstraintT(input_var, TensorType([d1, d2, d3, d4]), op_matching)
c2 = CalcMaxPool(maxpool, input_var, module_instance.kernel_size, module_instance.padding,
module_instance.stride, module_instance.dilation, [d1, d2, d3, d4])
nat_constraints = gen_nat_constraints([d1, d2, d3, d4])
return [c1, c2, *nat_constraints], counter
def adaptive_inference_rule(n: Node, module_instance, symbols, constraints, counter):
assert isinstance(n.args[0], Node)
avg_pool, counter = gen_tvar(counter)
symbols[n] = avg_pool
input_var = symbols[n.args[0]]
# dim vars
d1, counter = gen_dvar(counter)
d2, counter = gen_dvar(counter)
d3, counter = gen_dvar(counter)
d4, counter = gen_dvar(counter)
nat_constraints = gen_nat_constraints([d1, d2, d3, d4])
c1 = BinConstraintT(input_var, TensorType([d1, d2, d3, d4]), op_matching)
c2 = BinConstraintT(avg_pool, TensorType([d1, d2, module_instance.output_size[0], module_instance.output_size[1]]), op_eq)
return [c1, c2, *nat_constraints], counter
def batchnorm_inference_rule(n: Node, module_instance, symbols, constraints, counter):
assert isinstance(n.args[0], Node)
# generate the new variable
batchnorm_output, counter = gen_tvar(counter)
symbols[n] = batchnorm_output
batchnorm_input = symbols[n.args[0]]
# dim vars
d1, counter = gen_dvar(counter)
d2, counter = gen_dvar(counter)
d3, counter = gen_dvar(counter)
d4, counter = gen_dvar(counter)
nat_constraints = gen_nat_constraints([d1, d2, d3, d4])
c1 = BinConstraintT(batchnorm_input, TensorType([d1, d2, d3, d4]), op_matching)
c2 = BinConstraintT(batchnorm_input, batchnorm_output, op_eq)
return [c1, c2, *nat_constraints], counter
def transform_transpose(constraint, counter):
"""
Similar to a sequence of two index-selects
"""
dims, counter = gen_tensor_dims(constraint.tensor_size, counter)
is_valid_index1 = valid_index(constraint.index1, dims)
is_valid_index2 = valid_index(constraint.index2, dims)
new_dims = copy.deepcopy(dims)
nat_constraints = gen_nat_constraints(dims)
if is_valid_index1 == T() and is_valid_index2 == T():
new_dims[constraint.index1] = dims[constraint.index2]
new_dims[constraint.index2] = dims[constraint.index1]
transformed_constraint = Conj([
BinConstraintT(constraint.input_var, TensorType(dims), op_eq),
*nat_constraints, is_valid_index1, is_valid_index2,
BinConstraintT(constraint.output, TensorType(new_dims), op_eq)
])
return transformed_constraint, counter
def gen_layer_norm_constraints(n: Node, normalized_shape, symbols, counter):
output, counter = gen_tvar(counter)
symbols[n] = output
input = symbols[n.args[0]]
input_dyn = BinConstraintT(input, Dyn, op_eq)
output_dyn = BinConstraintT(output, Dyn, op_eq)
c1 = Conj([input_dyn, output_dyn])
c2 = []
for i in range(1, MAX_TENSOR_RANK + 1):
new_dims_rhs, counter = gen_tensor_dims(i, counter)
nat_constraints = gen_nat_constraints(new_dims_rhs)
c_tensor_i = Conj([
BinConstraintT(input, TensorType(new_dims_rhs), op_eq),
BinConstraintT(output, TensorType(new_dims_rhs), op_eq)
] + add_layer_norm_constraints(new_dims_rhs, list(normalized_shape)) +
nat_constraints)
c2.append(c_tensor_i)
return [Disj([c1, Disj(c2)])], counter
def gen_greatest_upper_bound(constraint: TGreatestUpperBound, counter: int):
"""
Args:
constraint: Greatest upper bound on tensors
counter: variable tracking
Returns: A set of equality constraints and DGreatestUpperBound constraints
"""
all_constraints = []
for i in range(1, MAX_TENSOR_RANK + 1):
c = []
dims1, counter = gen_tensor_dims(i, counter)
c1tensor = TensorType(dims1)
dims2, counter = gen_tensor_dims(i, counter)
c2tensor = TensorType(dims2)
dims3, counter = gen_tensor_dims(i, counter)
c3tensor = TensorType(dims3)
c += [BinConstraintT(constraint.rhs1, c1tensor, op_eq),
BinConstraintT(constraint.rhs2, c2tensor, op_eq),
BinConstraintT(constraint.res, c3tensor, op_eq)] + \
gen_nat_constraints(dims1 + dims2 + dims3)
assert len(c3tensor.__args__) == len(c1tensor.__args__) == len(
c2tensor.__args__)
for i in range(len(c3tensor.__args__)):
c.append(
DGreatestUpperBound(c3tensor.__args__[i], c1tensor.__args__[i],
c2tensor.__args__[i]))
all_constraints.append(Conj(c))
return all_constraints, counter
def linear_constraints(n: Node, in_features, out_features, symbols, counter):
linear_output, counter = gen_tvar(counter)
symbols[n] = linear_output
linear_input = symbols[n.args[0]]
input_dyn = BinConstraintT(linear_input, Dyn, op_eq)
output_dyn = BinConstraintT(linear_output, Dyn, op_eq)
c1 = Conj([input_dyn, output_dyn])
c2 = []
for i in range(1, MAX_TENSOR_RANK + 1):
new_dims_rhs_1, counter = gen_tensor_dims(i, counter)
new_dims_rhs_2, counter = gen_tensor_dims(i, counter)
nat_constraints = gen_nat_constraints(new_dims_rhs_1 + new_dims_rhs_2)
c_tensor_i = Conj([
BinConstraintT(linear_input, TensorType(new_dims_rhs_1), op_eq),
BinConstraintT(linear_output, TensorType(new_dims_rhs_2), op_eq)
] + add_linear_constraints(new_dims_rhs_1, new_dims_rhs_2, in_features,
out_features) + nat_constraints)
c2.append(c_tensor_i)
return [Disj([c1, Disj(c2)])], counter
def transform_index_select(constraint, counter):
"""
The constraints consider the given tensor size, checks if the index is valid
and if so, generates a constraint for replacing the input dimension
with the required dimension
"""
dims, counter = gen_tensor_dims(constraint.tensor_size, counter)
is_valid_index = valid_index(constraint.index, dims)
nat_constraints = gen_nat_constraints(dims)
# if the index is valid then replace the input dimension with the new dimension
# otherwise the dimension will not be replaced and the clause will contain False
if is_valid_index == T():
new_dims = copy.deepcopy((dims))
new_dims[constraint.index] = constraint.dim_replace
transformed_constraint = Conj([
BinConstraintT(constraint.input_var, TensorType(dims), op_eq),
*nat_constraints, is_valid_index,
BinConstraintT(constraint.output, TensorType(new_dims), op_eq)
])
# print(constraints)
return transformed_constraint, counter
