# Copyright 2021 The JAX Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Shape polymorphism support. We introduce a set of dimension variables at the top-level of a `jit` function. They are introduced implicitly by way of specifying for each dimension of each argument a symbolic dimension expression in terms of some dimension variables. All dimension variables are assumed to range over integers greater or equal to 1. Symbolic dimensions overload some integer operations, such as add, multiply, divide, equality, etc. The JAX NumPy layer and the LAX layers have been touched up to be sensitive to handling shapes that contain symbolic dimensions. This enables many JAX programs to be traced with symbolic dimensions in some dimensions. A priority has been to enable the batch dimension in neural network examples to be polymorphic. This was built initially for jax2tf, but it is now customizeable to be independent of TF. The best documentation at the moment is in the jax2tf.convert docstring, and the [README](https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md). """ import collections import dataclasses from enum import Enum import functools import itertools import io import math import operator as op import tokenize from typing import (Any, Callable, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union) import numpy as np import opt_einsum import jax from jax import config from jax.interpreters import xla from jax._src import core from jax._src import dtypes from jax._src.interpreters import mlir from jax._src.numpy import lax_numpy from jax._src import tree_util from jax._src import util from jax._src.typing import DimSize, Shape TfVal = Any DimVarEnv = Dict[str, jax.Array] DType = Any class InconclusiveDimensionOperation(core.InconclusiveDimensionOperation): """Raised when we cannot conclusively compute with symbolic dimensions.""" _help_msg = """ This error arises for comparison operations with shapes that are non-constant, and the result of the operation cannot be represented as a boolean value for all values of the symbolic dimensions involved. Please see https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#computing-with-dimension-variables for more details. """ def __init__(self, message: str): error_msg = f"{message}\n{InconclusiveDimensionOperation._help_msg}" # https://github.com/python/mypy/issues/5887 super().__init__(error_msg) # type: ignore class _DimAtom: """Represents an atom in a symbolic dimension expression. Atoms are either variables, or expressions of the form floordiv(E1, E2) or mod(E1, E2). Atoms are multiplied to form monomials (see _DimMon), and monomials are added to form symbolic expressions (see _DimExpr). Args: * var: if specified then the atom is a dimension variable. `operation` must be `None`. * operation: if specified then the atom is an operation applied to `operands`. One of `FLOORDIR` or `MOD`. `var` must be `None` * operands: the operands to which the operation is applied. """ # The supported operations FLOORDIV = "floordiv" MOD = "mod" def __init__(self, *operands: '_DimExpr', var: Optional[str] = None, operation: Optional[str] = None): if var is not None: assert operation is None assert not operands else: assert operation is not None self.var = var self.operation = operation self.operands = operands @classmethod def from_var(cls, v: str) -> '_DimAtom': return _DimAtom(var=v) def to_var(self) -> Optional[str]: return self.var def get_vars(self) -> Set[str]: # All the vars that appear if self.var is not None: return {self.var} else: acc = set() for opnd in self.operands: acc.update(opnd.get_vars()) return acc @classmethod def from_operation(cls, operation: str, *operands: '_DimExpr') -> '_DimAtom': return _DimAtom(*operands, operation=operation) def __str__(self): if self.var is not None: return self.var opnd_str = ", ".join([str(opnd) for opnd in self.operands]) return f"{self.operation}({opnd_str})" __repr__ = __str__ def __hash__(self): return hash((self.var, self.operation, *self.operands)) def __eq__(self, other: Any): # Used only for hashing if not isinstance(other, _DimAtom): return False if (self.var is None) != (other.var is None): return False if self.var is not None: return self.var == other.var else: def symbolic_equal(e1: '_DimExpr', e2: '_DimExpr') -> bool: try: return e1 == e2 except InconclusiveDimensionOperation: return False return (self.operation == other.operation and all(symbolic_equal(self_o, other_o) for self_o, other_o in zip(self.operands, other.operands))) def __lt__(self, other: '_DimAtom'): """ Comparison to another atom in graded reverse lexicographic order. Used only for determining a sorting order, does not relate to the comparison of the values of the atom. """ if self.var is not None and other.var is not None: return self.var < other.var elif self.var is not None: return True elif other.var is not None: return True elif self.operation != other.operation: return self.operation < other.operation # type: ignore else: return id(self) < id(other) def bounds(self) -> Tuple[float, float]: """Returns the lower and upper bounds, or -+ inf.""" if self.var is not None: return (1, np.PINF) # variables are assumed to be >= 1 opnd_bounds = [opnd.bounds() for opnd in self.operands] if self.operation == _DimAtom.FLOORDIV: # a // b (a_l, a_u), (b_l, b_u) = opnd_bounds def math_floor_with_inf(a: float, b: float): # math.floor, but aware of inf assert b != 0 if not np.isinf(b): # divisor is finite return math.floor(a / b) if not np.isinf(a) else np.NINF if (a >= 0) != (b >= 0) else np.PINF elif not np.isinf(a): # dividend is finite and divisor is infinite return -1 if (a >= 0) != (b >= 0) else 0 else: # both dividend and divisor are infinite return np.NINF if (a >= 0) != (b >= 0) else np.PINF # Same reasoning as for multiplication: the bounds are among the cross-product # of the bounds. bound_candidates = [math_floor_with_inf(a_l, b_l), math_floor_with_inf(a_l, b_u), math_floor_with_inf(a_u, b_l), math_floor_with_inf(a_u, b_u)] return (min(*bound_candidates), max(*bound_candidates)) elif self.operation == _DimAtom.MOD: _, (b_l, b_u) = opnd_bounds if b_l > 0: # positive divisor return (0, b_u - 1) elif b_u < 0: # negative divisor return (b_l + 1, 0) else: return (np.NINF, np.PINF) else: assert False def evaluate(self, env: DimVarEnv): if self.var is not None: try: return env[self.var] except KeyError: err_msg = ( f"Encountered dimension variable '{self.var}' that is not appearing in the shapes of the used function arguments.\n" "Please see https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#dimension-variables-must-be-solvable-from-the-input-shapes for more details.") raise KeyError(err_msg) else: operand_values = [opnd.evaluate(env) for opnd in self.operands] div_mod = divmod(*operand_values) # type: ignore if self.operation == _DimAtom.FLOORDIV: return div_mod[0] elif self.operation == _DimAtom.MOD: return div_mod[1] else: assert False, self.operation class _DimMon(dict): """Represents a multiplication of atoms. The representation is a dictionary mapping _DimAtom to exponent. The exponents are integers >= 1. """ def __hash__(self): return hash(frozenset(self.items())) def __str__(self): return "*".join(f"{key}^{exponent}" if exponent != 1 else str(key) for key, exponent in sorted(self.items())) @classmethod def from_var(cls, v: str) -> '_DimMon': return _DimMon({_DimAtom.from_var(v): 1}) @classmethod def from_atom(clscls, a: _DimAtom, aexp: int): return _DimMon({a: aexp}) def to_var(self) -> Optional[str]: """Extract the variable name "x", from a monomial "x". Return None, if the monomial is not a single variable.""" items = self.items() if len(items) != 1: return None (a, aexp), = items if aexp != 1: return None return a.to_var() def get_vars(self) -> Set[str]: # All the vars that appear in the monomial acc = set() for a in self.keys(): acc.update(a.get_vars()) return acc @classmethod def from_operation(cls, operation: str, *operands: '_DimExpr') -> '_DimMon': return _DimMon({_DimAtom.from_operation(operation, *operands): 1}) @property def degree(self): return sum(self.values()) def __lt__(self, other: '_DimMon'): """ Comparison to another monomial in graded reverse lexicographic order. Used only for determining a sorting order, does not relate to the comparison of the values of the monomial. """ self_key = -self.degree, tuple(sorted(self)) other_key = -other.degree, tuple(sorted(other)) return self_key > other_key def mul(self, other: '_DimMon') -> '_DimMon': """ Returns the product with another monomial. Example: (n^2*m) * n == n^3 * m. """ return _DimMon(collections.Counter(self) + collections.Counter(other)) def divide(self, divisor: '_DimMon') -> '_DimMon': """ Divides by another monomial. Raises a InconclusiveDimensionOperation if the result is not a monomial. For example, (n^3 * m) // n == n^2*m, but n // m fails. """ d = collections.Counter(self) for key, exponent in divisor.items(): diff = self.get(key, 0) - exponent if diff < 0: raise InconclusiveDimensionOperation(f"Cannot divide {self} by {divisor}.") elif diff == 0: del d[key] elif diff > 0: d[key] = diff return _DimMon(d) def bounds(self) -> Tuple[float, float]: """Returns the lower and upper bounds, or -+inf.""" # The bounds of a product are among the product of bounds. bounds = [] for a, exp in self.items(): a_l, a_u = a.bounds() assert a_l <= a_u bounds.append((a_l ** exp, a_u ** exp)) candidates = [math.prod(atom_bounds) for atom_bounds in itertools.product(*bounds)] return (min(*candidates), max(*candidates)) # type: ignore def evaluate(self, env: DimVarEnv): prod = lambda xs: functools.reduce(_evaluate_multiply, xs) if xs else core.dim_constant(1) def pow_opt(v, p: int): return v if p == 1 else prod([v] * p) return prod([pow_opt(a.evaluate(env), deg) for a, deg in self.items()]) class _DimExpr(): """Symbolic expression in terms of dimension variables. A dimension expression is an addition of products (_DimMon) of atoms (_DimAtom). We overload integer operations, but we do that soundly, raising :class:`InconclusiveDimensionOperation` when the result is not representable as a _DimExpr. The representation of a _DimExpr is as a dictionary mapping _DimMon to integer coefficients. The special monomial `_DimMon()` is mapped to the free integer coefficient of the expression. """ __array_priority__ = 1000 # Same as tracer, for __radd__ and others on ndarray def __init__(self, coeffs: Dict[_DimMon, int]): # Do not construct _DimExpr directly, unless you are sure that coeffs is # normalized; Use _DimExpr.normalize. # Takes ownership of coeffs self._coeffs = coeffs or {_DimMon(): 0} def monomials(self) -> Iterable[Tuple[_DimMon, int]]: return self._coeffs.items() @classmethod def _add_coeffs(cls, coeffs: Dict[_DimMon, int], mon: _DimMon, coeff: int): """Do `coeffs[mon] += coeff` but remove 0 coefficients.""" old_c = coeffs.get(mon) if old_c is None: if coeff != 0: coeffs[mon] = coeff else: new_c = old_c + coeff if new_c == 0: del coeffs[mon] else: coeffs[mon] = new_c @classmethod def normalize(cls, coeffs: Dict[_DimMon, int]) -> DimSize: """The main constructor for _DimExpr. Ensures that the symbolic dimension is normalized, e.g., it is represented as a Python int if it is known to be a constant. """ # TODO(necula): profile and optimize this has_non_zero_degree = False free_const = 0 new_coeffs: Dict[_DimMon, int] = {} for mon, coeff in coeffs.items(): if coeff == 0: continue if mon.degree == 0: # A constant, there can be a single one free_const = coeff else: has_non_zero_degree = True new_coeffs[mon] = new_coeffs.get(mon, 0) + coeff if has_non_zero_degree: return _DimExpr(new_coeffs) else: return int(free_const) @classmethod def normalize_floordiv_times_divisor(cls, coeffs: Dict[_DimMon, int]) -> DimSize: # Look for floordiv(E, M) * M and turn into E - mod(E, M). This comes # up when handling strided convolution. for dec in _decompose_expr(_DimExpr(coeffs), _DimAtom.FLOORDIV): # e = factor * floordiv(operands)^exp * rest_monomial + rest_expr if dec.exp != 1: continue if dec.rest_monomial == 1 and dec.factor == 1: continue m_trimmed, m_remainder = divmod(dec.factor * dec.rest_monomial, dec.operands[1]) if m_remainder == 0: return m_trimmed * (dec.operands[0] - _DimExpr.from_operation(_DimAtom.MOD, *dec.operands)) + dec.rest_expr return _DimExpr.normalize(coeffs) @classmethod def from_monomial(cls, mon: _DimMon, exp: int): return _DimExpr.normalize({mon: exp}) @classmethod def from_var(cls, v: str) -> '_DimExpr': return _DimExpr({_DimMon.from_var(v): 1}) @classmethod def from_operation(cls, operation: str, *operands: '_DimExpr') -> '_DimExpr': return _DimExpr.from_monomial(_DimMon.from_operation(operation, *operands), 1) def to_var(self) -> Optional[str]: """Extract the variable name "x", from a symbolic expression.""" items = self.monomials() if len(items) != 1: # type: ignore return None (mon, mon_count), = items if mon_count != 1: return None return mon.to_var() def get_vars(self) -> Set[str]: """The variables that appear in a symbolic dimension.""" acc = set() for mon, _ in self.monomials(): acc.update(mon.get_vars()) return acc def eq(self, other: DimSize) -> bool: lb, ub = _ensure_poly(self - other, "eq").bounds() if lb == ub == 0: return True if lb > 0 or ub < 0: return False # See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#comparison-of-symbolic-dimensions-is-partially-supported return False def ge(self, other: DimSize) -> bool: lb, ub = _ensure_poly(self - other, "ge").bounds() if lb >= 0: return True if ub < 0: return False raise InconclusiveDimensionOperation( f"Symbolic dimension comparison '{self}' >= '{other}' is inconclusive.\n" "See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#comparison-of-symbolic0dimensions-is-partially-supported.") def __hash__(self): return hash(tuple(sorted(self.monomials()))) def __str__(self): def _one_monomial(mon, c): if mon.degree == 0: return str(c) if c == 1: return str(mon) return f"{c}*{mon}" return " + ".join(_one_monomial(mon, c) for mon, c in sorted(self.monomials(), reverse=True)) def __repr__(self): return str(self) # We overload +, -, *, because they are fully defined for _DimExpr. def __add__(self, other): if isinstance(other, core.Tracer) or not _convertible_to_poly(other): return self.__jax_array__().__add__(other) other = _ensure_poly(other, "add") coeffs = self._coeffs.copy() for mon, coeff in other.monomials(): _DimExpr._add_coeffs(coeffs, mon, coeff) return _DimExpr.normalize_floordiv_times_divisor(coeffs) def __radd__(self, other): if isinstance(other, core.Tracer) or not _convertible_to_poly(other): return self.__jax_array__().__radd__(other) return _ensure_poly(other, "add").__add__(self) def __sub__(self, other): if isinstance(other, core.Tracer) or not _convertible_to_poly(other): return self.__jax_array__().__sub__(other) return self + -_ensure_poly(other, "sub") def __rsub__(self, other): if isinstance(other, core.Tracer) or not _convertible_to_poly(other): return self.__jax_array__().__rsub__(other) return _ensure_poly(other, "sub").__sub__(self) def __neg__(self) -> '_DimExpr': return _DimExpr({mon: -coeff for mon, coeff in self.monomials()}) def __mul__(self, other): if isinstance(other, core.Tracer) or not _convertible_to_poly(other): return self.__jax_array__().__mul__(other) other = _ensure_poly(other, "mul") coeffs: Dict[_DimMon, int] = {} for mon1, coeff1 in self.monomials(): for mon2, coeff2 in other.monomials(): mon = mon1.mul(mon2) _DimExpr._add_coeffs(coeffs, mon, coeff1 * coeff2) return _DimExpr.normalize_floordiv_times_divisor(coeffs) def __rmul__(self, other): if isinstance(other, core.Tracer) or not _convertible_to_poly(other): return self.__jax_array__().__rmul__(other) return _ensure_poly(other, "mul").__mul__(self) def __pow__(self, power, modulo=None): assert modulo is None try: power = int(power) except: raise InconclusiveDimensionOperation(f"Symblic dimension cannot be raised to non-integer power '{self}' ^ '{power}'") return functools.reduce(op.mul, [self] * power) def __floordiv__(self, divisor): if isinstance(divisor, core.Tracer) or not _convertible_to_poly(divisor): return self.__jax_array__().__floordiv__(divisor) return self.divmod(_ensure_poly(divisor, "floordiv"))[0] def __rfloordiv__(self, other): if isinstance(other, core.Tracer) or not _convertible_to_poly(other): return self.__jax_array__().__rfloordiv__(other) return _ensure_poly(other, "floordiv").__floordiv__(self) def __truediv__(self, divisor): # Used for "/", which always returns a float return self.__jax_array__().__truediv__(divisor) def __rtruediv__(self, dividend): # Used for "/", when dividend is not a _DimExpr return self.__jax_array__().__rtruediv__(dividend) def __mod__(self, divisor): if isinstance(divisor, core.Tracer) or not _convertible_to_poly(divisor): return self.__jax_array__().__mod__(divisor) return self.divmod(_ensure_poly(divisor, "mod"))[1] def __rmod__(self, dividend): if isinstance(dividend, core.Tracer) or not _convertible_to_poly(dividend): return self.__jax_array__().__rmod__(dividend) return _ensure_poly(dividend, "mod").__mod__(self) def __divmod__(self, divisor): if isinstance(divisor, core.Tracer) or not _convertible_to_poly(divisor): return self.__jax_array__().__divmod__(divisor) return self.divmod(_ensure_poly(divisor, "divmod")) def __rdivmod__(self, dividend): if isinstance(dividend, core.Tracer) or not _convertible_to_poly(dividend): return self.__jax_array__().__rdivmod__(dividend) return _ensure_poly(dividend, "divmod").__divmod__(self) def __int__(self): if self.is_constant: return op.index(next(iter(self._coeffs.values()))) else: raise InconclusiveDimensionOperation(f"Symbolic dimension '{self}' used in a context that requires a constant") # We must overload __eq__ and __ne__, or else we get unsound defaults. __eq__ = eq def __ne__(self, other: DimSize) -> bool: return not self.eq(other) __ge__ = ge def __le__(self, other: DimSize): return _ensure_poly(other, "le").__ge__(self) def __gt__(self, other: DimSize): return not _ensure_poly(other, "gt").__ge__(self) def __lt__(self, other: DimSize): return not self.__ge__(other) def divmod(self, divisor: "_DimExpr") -> Tuple[DimSize, int]: """ Floor division with remainder (divmod) generalized to polynomials. If the `divisor` is not a constant, the remainder must be 0. If the `divisor` is a constant, the remainder may be non 0, for consistency with integer divmod. :return: Quotient resulting from polynomial division and integer remainder. """ assert isinstance(divisor, _DimExpr) try: dmon, dcount = divisor.leading_term dividend, quotient = self, 0 # invariant: self = dividend + divisor * quotient # quotient and dividend are changed in the loop; the leading term of # dividend decreases at each iteration. while is_poly_dim(dividend) and not dividend.is_constant: mon, count = dividend.leading_term try: qmon = mon.divide(dmon) except InconclusiveDimensionOperation: raise InconclusiveDimensionOperation("") qcount, rcount = divmod(count, dcount) if rcount != 0: raise InconclusiveDimensionOperation("") q = _DimExpr.from_monomial(qmon, qcount) quotient += q dividend -= q * divisor # type: ignore[assignment] dividend = int(dividend) # type: ignore[assignment] if divisor.is_constant: q, r = divmod(dividend, int(divisor)) # type: ignore quotient += q remainder = r else: if dividend != 0: raise InconclusiveDimensionOperation("") remainder = 0 if config.jax_enable_checks: assert self == divisor * quotient + remainder return quotient, remainder except InconclusiveDimensionOperation: return (_DimExpr.from_operation(_DimAtom.FLOORDIV, self, divisor), # type: ignore _DimExpr.from_operation(_DimAtom.MOD, self, divisor)) def bounds(self) -> Tuple[float, float]: """Returns the lower and upper bounds, or -+inf.""" lb = ub = self._coeffs.get(_DimMon(), 0) # The free coefficient for mon, coeff in self.monomials(): if mon.degree == 0: continue # We already included the free coefficient m_l, m_u = mon.bounds() assert m_l <= m_u and coeff != 0 item_l, item_u = coeff * m_l, coeff * m_u lb = lb + min(item_l, item_u) # type: ignore ub = ub + max(item_l, item_u) # type: ignore if lb != np.NINF or ub != np.PINF: return lb, ub # Watch for special-case: ct*a - ct*mod(b, a) >= 1 when ct >= 0 and a >= 0 # TODO(necula): add more principled support for floordiv and mod # For example, this will miss "1 + a - mod(b, a)" for dec in _decompose_expr(self, _DimAtom.MOD): # E = factor*mod(op1, op2)^exp * rest_monomial + rest_expr if dec.exp == 1 and dec.rest_monomial == 1 and dec.rest_expr == - dec.factor * dec.operands[1]: try: if dec.operands[1] <= 0: continue except InconclusiveDimensionOperation: continue if dec.factor > 0: return (np.NINF, -1) else: return (1, np.PINF) return lb, ub @property def is_constant(self): return len(self._coeffs) == 1 and next(iter(self._coeffs)).degree == 0 @property def leading_term(self) -> Tuple[_DimMon, int]: """Returns the highest degree term that comes first lexicographically.""" return max(self.monomials()) def evaluate(self, env: DimVarEnv): # Evaluates as a value of dtype=core.dim_value_dtype() terms = [_evaluate_multiply(mon.evaluate(env), core.dim_constant(coeff)) for mon, coeff in self.monomials()] return functools.reduce(_evaluate_add, terms) if len(terms) > 1 else terms[0] @staticmethod def get_aval(dim: "_DimExpr"): return core.dim_value_aval() def dimension_as_value(self): """Turns a dimension size into a Jax value that we can compute with.""" return _dim_as_value(self) def __jax_array__(self): # Used for implicit coercions of polynomials as JAX arrays return _dim_as_value(self) @dataclasses.dataclass class _Decomposition: """Decomposition of an expression around an operation atom. E = factor * mod(*operands)^exp * rest_monomial + rest_expr """ factor: int operands: Sequence[_DimExpr] exp: int rest_monomial: _DimExpr rest_expr: _DimExpr def _decompose_expr(e: _DimExpr, operation: str) -> Iterable[_Decomposition]: for m, m_factor in e.monomials(): atoms = [(a, aexp) for a, aexp in m.items() if a.operation == operation] if atoms: e_minus_m_coeffs = e._coeffs.copy() del e_minus_m_coeffs[m] for a, aexp in atoms: yield _Decomposition( factor=m_factor, operands=a.operands, exp=aexp, rest_monomial=_DimExpr({m.divide(_DimMon.from_atom(a, aexp)): 1}), rest_expr=_DimExpr(e_minus_m_coeffs)) core.pytype_aval_mappings[_DimExpr] = _DimExpr.get_aval xla.pytype_aval_mappings[_DimExpr] = _DimExpr.get_aval dtypes._weak_types.append(_DimExpr) def _convertible_to_int(p: DimSize) -> bool: try: op.index(p) return True except: return False def _ensure_poly(p: DimSize, operation_name: str) -> _DimExpr: if isinstance(p, _DimExpr): return p if _convertible_to_int(p): return _DimExpr({_DimMon(): op.index(p)}) raise TypeError(f"Symnbolic dimension {operation_name} not supported for {p}.") def _convertible_to_poly(p: DimSize) -> bool: return isinstance(p, _DimExpr) or _convertible_to_int(p) def is_poly_dim(p: DimSize) -> bool: return isinstance(p, _DimExpr) class DimensionHandlerPoly(core.DimensionHandler): """See core.DimensionHandler. Most methods are inherited. """ def is_constant(self, d: DimSize) -> bool: assert isinstance(d, _DimExpr) return False def symbolic_equal(self, d1: core.DimSize, d2: core.DimSize) -> bool: try: return _ensure_poly(d1, "equal") == d2 except InconclusiveDimensionOperation: return False def greater_equal(self, d1: DimSize, d2: DimSize): return _ensure_poly(d1, "ge") >= d2 def divide_shape_sizes(self, s1: Shape, s2: Shape) -> DimSize: sz1 = math.prod(s1) sz2 = math.prod(s2) if core.symbolic_equal_dim(sz1, sz2): # Takes care also of sz1 == sz2 == 0 return 1 err_msg = f"Cannot divide evenly the sizes of shapes {tuple(s1)} and {tuple(s2)}" try: q, r = _ensure_poly(sz1, "divide_shape").divmod(_ensure_poly(sz2, "divide_shape")) except InconclusiveDimensionOperation as e: raise InconclusiveDimensionOperation(err_msg + f"\nDetails: {e}") if not core.symbolic_equal_dim(r, 0): raise InconclusiveDimensionOperation(err_msg + f"\nRemainder is not zero: {r}") return q # type: ignore[return-value] def stride(self, d: DimSize, window_size: DimSize, window_stride: DimSize) -> DimSize: """Implements `(d - window_size) // window_stride + 1`""" try: # TODO(necula): check for d == 0 or window_size > d and return 0. q, r = _ensure_poly(d - window_size, "stride").divmod(_ensure_poly(window_stride, "stride")) return q + 1 except InconclusiveDimensionOperation as e: raise InconclusiveDimensionOperation( f"Cannot compute stride for dimension '{d}', " f"window_size '{window_size}', stride '{window_stride}'.\nDetails: {e}.") return d def as_value(self, d: DimSize): """Turns a dimension size into a Jax value that we can compute with.""" return _dim_as_value(d) core._SPECIAL_DIMENSION_HANDLERS[_DimExpr] = DimensionHandlerPoly() dtypes.python_scalar_dtypes[_DimExpr] = dtypes.python_scalar_dtypes[int] def _einsum_contract_path(*operands, **kwargs): """Like opt_einsum.contract_path, with support for DimExpr shapes. We use opt_einsum.contract_path to compute the schedule, using a fixed constant for all dimension variables. This is safe because we throw an error if there are more than 1 contractions. Essentially, we just use opt_einsum.contract_path to parse the specification. """ # Replace the polymorphic shapes with some concrete shapes for calling # into opt_einsum.contract_path, because the latter wants to compute the # sizes of operands and intermediate results. fake_ops = [] for operand in operands: # We replace only array operands if not hasattr(operand, "dtype"): fake_ops.append(operand) else: shape = np.shape(operand) def fake_dim(d): if core.is_constant_dim(d): return d else: if not isinstance(d, _DimExpr): raise TypeError(f"Encountered unexpected shape dimension {d}") # It is Ok to replace all polynomials with the same value. We may miss # here some errors due to non-equal dimensions, but we catch them # later. return 8 fake_ops.append(jax.ShapeDtypeStruct(tuple(map(fake_dim, shape)), operand.dtype)) contract_fake_ops, contractions = opt_einsum.contract_path(*fake_ops, **kwargs) contract_operands = [] for operand in contract_fake_ops: idx = tuple(i for i, fake_op in enumerate(fake_ops) if operand is fake_op) assert len(idx) == 1 contract_operands.append(operands[idx[0]]) return contract_operands, contractions lax_numpy._poly_einsum_handlers[_DimExpr] = _einsum_contract_path # A JAX primitive with no array arguments but with a dimension parameter # that is a DimExpr. The value of the primitive is the value of the dimension, # using int64 in x64 mode or int32 otherwise (core.dim_value_dtype()) dim_as_value_p = core.Primitive("dim_as_value") dim_as_value_p.def_abstract_eval(lambda dim: core.dim_value_aval()) def dim_as_value_impl(dim: DimSize): raise NotImplementedError( "Evaluation rule for 'dim_as_value' is not implemented. " "It seems that you are using shape polymorphism outside jax2tf.") dim_as_value_p.def_impl(dim_as_value_impl) def _dim_as_value(dim: DimSize): return dim_as_value_p.bind(dim=dim) def _dim_as_value_lowering(ctx: mlir.LoweringRuleContext, *, dim): res, = mlir.eval_dynamic_shape(ctx, (dim,)) out_type = mlir.aval_to_ir_type(ctx.avals_out[0]) if out_type != res.type: # type: ignore return mlir.hlo.ConvertOp(out_type, res).results else: return [res] mlir.register_lowering(dim_as_value_p, _dim_as_value_lowering) class PolyShape(tuple): """Tuple of polymorphic dimension specifications. See docstring of :func:`jax2tf.convert`. """ def __init__(self, *dim_specs): tuple.__init__(dim_specs) def __new__(cls, *dim_specs): for ds in dim_specs: if not isinstance(ds, (int, str)) and ds != ...: msg = (f"Invalid polymorphic shape element: {repr(ds)}; must be a string " "representing a dimension variable, or an integer, or ...") raise ValueError(msg) return tuple.__new__(PolyShape, dim_specs) def __str__(self): return "(" + ", ".join(["..." if d is ... else str(d) for d in self]) + ")" def _parse_spec(shape_spec: Union[str, PolyShape, None], arg_shape: Sequence[Optional[int]]) -> Sequence[DimSize]: """Parses the shape polymorphic specification for one array argument. We have to be able to parse all strings produced by str(_DimExpr) because sometimes the output polymorphic shapes of one function become the input polymorphic shapes of another. Args: shape_spec: a shape polymorphic specification. None stands for "...". arg_shape: an actual shape, possibly containing unknown dimensions (None). We use `arg_shape` to fill-in the placeholders `_` and `...` in the `shape_spec`. The dimensions of `arg_shape` that are used for filling must be known (not `None`). If a dimension in `arg_shape` is known and the corresponding dimension in `shape_spec` is a constant then they must be equal. See the README.md for usage. """ shape_spec_repr = repr(shape_spec) if shape_spec is None: shape_spec = "..." elif isinstance(shape_spec, PolyShape): shape_spec = str(shape_spec) elif not isinstance(shape_spec, str): raise ValueError("polymorphic shape spec should be None or a string. " f"Found {shape_spec_repr}.") return _Parser(shape_spec, arg_shape, shape_spec_repr).parse() class _Parser: def __init__(self, shape_spec: str, arg_shape: Sequence[Optional[int]], shape_spec_repr: str): self.shape_spec = shape_spec self.shape_spec_repr = shape_spec_repr # For error messages self.arg_shape = arg_shape self.dimensions: List[DimSize] = [] # dimensions we have parsed def parse(self) -> Sequence[DimSize]: self.tokstream = tokenize.tokenize( io.BytesIO(self.shape_spec.encode("utf-8")).readline) tok = self.consume_token(self.next_tok(), tokenize.ENCODING) # Always 1st sh, tok = self.shape(tok) self.expect_token(tok, [tokenize.ENDMARKER]) return sh def add_dim(self, expr: Optional[DimSize], tok: tokenize.TokenInfo): if expr is None: raise self.parse_err(tok, ("unexpected placeholder for unknown dimension " f"for argument shape {self.arg_shape}")) arg_shape_dim = self.arg_shape[len(self.dimensions)] if core.is_constant_dim(expr) and arg_shape_dim is not None: if expr != arg_shape_dim: raise self.parse_err(tok, (f"different size {expr} for known dimension " f"for argument shape {self.arg_shape}")) self.dimensions.append(expr) def parse_err(self, tok: Optional[tokenize.TokenInfo], detail: str) -> Exception: msg = ( f"syntax error in polymorphic shape {self.shape_spec_repr} " f"in dimension {len(self.dimensions)}: {detail}. ") if tok is not None: msg += f"Parsed '{tok.line[:tok.start[1]]}', remaining '{tok.line[tok.start[1]:]}'." return ValueError(msg) def next_tok(self) -> tokenize.TokenInfo: while True: try: t = next(self.tokstream) except StopIteration: raise self.parse_err(None, "unexpected end of string") if t.exact_type not in [tokenize.NEWLINE, tokenize.INDENT, tokenize.DEDENT]: return t def expect_token(self, tok: tokenize.TokenInfo, expected: Sequence[int]) -> None: if tok.exact_type not in expected: msg = ("expecting one of {" + ", ".join(tokenize.tok_name[t] for t in expected) + "} but found " + tokenize.tok_name[tok.exact_type]) raise self.parse_err(tok, msg) def consume_token(self, tok: tokenize.TokenInfo, expected: int) -> tokenize.TokenInfo: self.expect_token(tok, [expected]) return self.next_tok() def integer(self, tok: tokenize.TokenInfo) -> Tuple[int, tokenize.TokenInfo]: self.expect_token(tok, [tokenize.NUMBER]) try: val = int(tok.string) except Exception: raise self.parse_err(tok, f"expecting integer, found {tok.string}") return val, self.next_tok() # What can follow a shape? FOLLOW_SHAPE = [tokenize.ENDMARKER, tokenize.RPAR] def shape(self, tok: tokenize.TokenInfo) -> Tuple[Sequence[DimSize], tokenize.TokenInfo]: # A comma-separated list of _DimExpr, or "_", possibly ended with ... if tok.exact_type == tokenize.LPAR: res, tok = self.shape(self.next_tok()) tok = self.consume_token(tok, tokenize.RPAR) return res, tok while True: if tok.exact_type in self.FOLLOW_SHAPE: break if tok.exact_type == tokenize.ELLIPSIS: to_add = self.arg_shape[len(self.dimensions):] for ad in to_add: self.add_dim(ad, tok) tok = self.next_tok() break if len(self.dimensions) >= len(self.arg_shape): raise self.parse_err(tok, f"too many dimensions, arg_shape has {len(self.arg_shape)}") if tok.exact_type == tokenize.NAME and tok.string == "_": e = self.arg_shape[len(self.dimensions)] tok = self.next_tok() else: e, tok = self.expr(tok) self.add_dim(e, tok) if tok.exact_type in self.FOLLOW_SHAPE: break tok = self.consume_token(tok, tokenize.COMMA) return tuple(self.dimensions), tok # What token can follow a _DimExpr FOLLOW_EXPR = FOLLOW_SHAPE + [tokenize.COMMA] def expr(self, tok: tokenize.TokenInfo) -> Tuple[DimSize, tokenize.TokenInfo]: # A sum of monomials next_m_negated = False acc = 0 while True: m, tok = self.mon(tok) acc = acc + (- m if next_m_negated else m) if tok.exact_type in self.FOLLOW_EXPR: return acc, tok next_m_negated = (tok.exact_type == tokenize.MINUS) self.expect_token(tok, [tokenize.PLUS, tokenize.MINUS]) tok = self.next_tok() FOLLOW_MON = FOLLOW_EXPR + [tokenize.PLUS, tokenize.MINUS] def mon(self, tok: tokenize.TokenInfo) -> Tuple[DimSize, tokenize.TokenInfo]: # A monomial is product of atoms. Each atom may be raised to an integer power. acc = 1 while True: a, tok = self.atom(tok) if tok.exact_type == tokenize.CIRCUMFLEX: tok = self.next_tok() self.expect_token(tok, [tokenize.NUMBER]) power, tok = self.integer(tok) a = a ** power acc = acc * a if tok.exact_type in self.FOLLOW_MON: return acc, tok tok = self.consume_token(tok, tokenize.STAR) def atom(self, tok: tokenize.TokenInfo) -> Tuple[DimSize, tokenize.TokenInfo]: if tok.exact_type == tokenize.NAME: if tok.string == "mod": return self.binary_op(_DimAtom.MOD, self.next_tok()) if tok.string == "floordiv": return self.binary_op(_DimAtom.FLOORDIV, self.next_tok()) return _DimExpr.from_var(tok.string), self.next_tok() number_sign = 1 if tok.exact_type == tokenize.MINUS: # -k are negative constants number_sign = -1 tok = self.next_tok() self.expect_token(tok, [tokenize.NUMBER]) if tok.exact_type == tokenize.NUMBER: v, tok = self.integer(tok) return v * number_sign, tok self.expect_token(tok, [tokenize.NAME, tokenize.MINUS, tokenize.NUMBER]) assert False def binary_op(self, op: str, tok) -> Tuple[DimSize, tokenize.TokenInfo]: tok = self.consume_token(tok, tokenize.LPAR) e1, tok = self.expr(tok) tok = self.consume_token(tok, tokenize.COMMA) e2, tok = self.expr(tok) tok = self.consume_token(tok, tokenize.RPAR) return _DimExpr.from_operation(op, e1, e2), tok # type: ignore def _evaluate_add(v1, v2): try: if op.index(v1) == 0: return v2 except: pass try: if op.index(v2) == 0: return v1 except: pass return v1 + v2 def _evaluate_multiply(v1, v2): try: if op.index(v1) == 1: return v2 except: pass try: if op.index(v2) == 1: return v1 except: pass return v1 * v2 def _is_known_constant(v) -> Optional[int]: try: return int(v) except Exception: # TODO(necula): added this so that in jax2tf, in Eager mode, we can tell # that a tensor is a constant. We should move this dependency into some # jax2tf-specific area. if hasattr(v, "val"): try: vint = int(v.val) if isinstance(vint, int): # In TF, int(tf.Tensor) is tf.Tensor! return vint except Exception: pass return None # dimension_size(operand, dimension=i) get the operand.shape[i] as a # value of type shape_poly.dim_as_value_dtype(). dimension_size_p = core.Primitive("dimension_size") def _dimension_size_abstract_eval(aval: core.AbstractValue, **_) -> core.AbstractValue: return core.dim_value_aval() dimension_size_p.def_abstract_eval(_dimension_size_abstract_eval) def _dimension_size_impl(arg, *, dimension): return core.dim_constant(arg.shape[dimension]) dimension_size_p.def_impl(_dimension_size_impl) def _dimension_size_lowering_rule(ctx, arg, *, dimension): dim_size = mlir.hlo.GetDimensionSizeOp(arg, dimension) dim_type = mlir.aval_to_ir_type(core.dim_value_aval()) if dim_size.result.type != dim_type: dim_size = mlir.hlo.ConvertOp(dim_type, dim_size) return dim_size.results mlir.register_lowering(dimension_size_p, _dimension_size_lowering_rule) def arg_aval( arg_shape: Sequence[Optional[int]], arg_jax_dtype: DType, polymorphic_shape: Optional[Union[str, PolyShape]]) -> core.ShapedArray: """Computes abstract values. Args: arg_shape: the shape for the argument, possibly having None dimensions. arg_dtype: the inferred JAX dtype for the arg. polymorphic_shape: the polymorphic specification for the argument. Returns: the JAX abstract value for the argument. """ aval_shape = _parse_spec(polymorphic_shape, arg_shape) return core.ShapedArray(aval_shape, arg_jax_dtype) def all_dim_vars(args_avals: Sequence[core.AbstractValue]) -> Sequence[str]: dim_vars: Set[str] = set() for a in args_avals: for d in a.shape: if is_poly_dim(d): dim_vars = dim_vars.union(d.get_vars()) return sorted(tuple(dim_vars)) @dataclasses.dataclass(frozen=True) class ShapeConstraint: class Comparator(Enum): EQ = 1 GEQ = 2 comp: Comparator left: DimSize right: DimSize # make_err_msg is invoked with (left_int, right_int) if the constraint fails. make_err_msg: Callable[[int, int], str] def check(self, shapeenv: DimVarEnv) -> None: """Evaluates a constraint statically and raises an error if fails.""" def eval_operand(o: DimSize) -> Union[int, jax.Array]: if core.is_constant_dim(o): return op.index(o) return o.evaluate(shapeenv) # type: ignore try: left1, right1 = eval_operand(self.left), eval_operand(self.right) except KeyError: return None left_int, right_int = _is_known_constant(left1), _is_known_constant(right1) if left_int is not None and right_int is not None: if self.comp == ShapeConstraint.Comparator.EQ: if not (left_int == right_int): raise ValueError(self.make_err_msg(left_int, right_int)) elif self.comp == ShapeConstraint.Comparator.GEQ: if not (left_int >= right_int): raise ValueError(self.make_err_msg(left_int, right_int)) else: assert False else: return None # TODO: evaluate constraint dynamically def __str__(self): return (f"{self.left} {'==' if self.comp == ShapeConstraint.Comparator.EQ else '>='} {self.right}" f" ({self.make_err_msg(self.left, self.right)})") __repr__ = __str__ class ShapeConstraints: def __init__(self): self.constraints: Set[ShapeConstraint] = set() # map DimConstraint to an integer >= 0 def add_constraint(self, comp: ShapeConstraint.Comparator, left: DimSize, right: DimSize, make_err_msg: Callable[[int, int], str]): # Try to evaluate it statically c = ShapeConstraint(comp, left, right, make_err_msg) self.constraints.add(c) def check(self, shapeenv: DimVarEnv) -> None: for constraint in self.constraints: constraint.check(shapeenv) @dataclasses.dataclass class _DimEquation: # Represents dim_expr == dim_value, where `dim_expr` contain unknown dimension # variables, in terms of `dim_value`. dim_expr: _DimExpr dim_value: _DimExpr def __str__(self): return f"{self.dim_expr} == {self.dim_value}" __repr__ = __str__ def args_kwargs_path_to_str(path: tree_util.KeyPath) -> str: # String description of `args` or `kwargs`, assuming the path for a tree for # the tuple `(args, kwargs)`. if path[0] == tree_util.SequenceKey(0): return f"args{tree_util.keystr(path[1:])}" elif path[0] == tree_util.SequenceKey(1): return f"kwargs{tree_util.keystr(path[1:])}" else: assert False def pretty_print_dimension_descriptor( args_kwargs_tree: tree_util.PyTreeDef, flat_arg_idx: int, dim_idx: Optional[int]) -> str: args_kwargs_with_paths, _ = tree_util.tree_flatten_with_path( args_kwargs_tree.unflatten((0,) * args_kwargs_tree.num_leaves)) arg_str = args_kwargs_path_to_str(args_kwargs_with_paths[flat_arg_idx][0]) if dim_idx is not None: arg_str += f".shape[{dim_idx}]" return arg_str @util.cache() def solve_dim_vars( args_avals: Sequence[core.AbstractValue], args_kwargs_tree: tree_util.PyTreeDef, ) -> Tuple[DimVarEnv, ShapeConstraints, Sequence[Tuple[str, int, int]]]: """Solves dimension variables in a called function's avals in terms of actual argument shapes. For example, given: args_avals = [ShapedArray((3, a, a + b), f32)] we introduce fresh "known" dimension variables to represent the actual dimension size of actual arguments for each non-constant dimension. Each known variable has a name, an arg_idx, and a dim_idx, e.g.: known_vars = [("args[0].shape[1]", 0, 1), ("args[0].shape[2]", 0, 2)] and then we express the solution for the unknown dimension variables {a, b} as symbolic expressions in terms of the known variables: dict(a=args[0].shape[1], b=args[0].shape[2] - args[0].shape[1]) Not all equations are solvable. For now, we solve first the linear uni-variate equations, then the solved variables are used to simplify the remaining equations to linear uni-variate equations, and the process continues until all dimension variables are solved. Args: args_avals: the abstract values of the `args`, with shapes that may include unknown dimension variables. args_kwargs_tree: a PyTreeDef that describes the tuple `(args, kwargs)` from which the flat sequence `args_avals` is extracted. Used for describing args and kwargs in known variable names and in error messages. Returns: a 3-tuple with: (a) the solution for the unknown dimension variables (b) a list of constraints that must be satisfied for the solution to be a valid one, and (c) and the list of known variables that may appear in the solution and the constraints. Raises ValueError if it cannot solve some dimension variable. """ dim_equations: List[_DimEquation] = [] known_dimension_vars: List[Tuple[str, int, int]] = [] for arg_idx, aval in enumerate(args_avals): for dim_idx, aval_d in enumerate(aval.shape): if is_poly_dim(aval_d): known_dim_var = pretty_print_dimension_descriptor(args_kwargs_tree, arg_idx, dim_idx) known_dimension_vars.append((known_dim_var, arg_idx, dim_idx)) dim_equations.append( _DimEquation(dim_expr=_ensure_poly(aval_d, "solve_dim_vars"), dim_value=_DimExpr.from_var(known_dim_var))) solution, shape_constraints = _solve_dim_equations(dim_equations) return solution, shape_constraints, known_dimension_vars def compute_dim_vars_from_arg_shapes( args_avals: Sequence[core.AbstractValue], *actual_args: jax.Array, args_kwargs_tree: tree_util.PyTreeDef) -> Sequence[jax.Array]: """Computes values of dimension variables to unify args_avals with actual arguments. Like `solve_dim_vars` except that here we express the solution as JAX arrays that reference the `actual_args`. This function can be used to generate the code for computing the dimension variables. Returns: the values of the dimension variables, in the order determined by `all_dim_vars(args_avals)`. """ dim_vars = all_dim_vars(args_avals) solution, shape_constraints, known_dim_vars = solve_dim_vars( tuple(args_avals), args_kwargs_tree=args_kwargs_tree) # Replace the synthetic vars with the dynamic shape of the actual arg known_env = {vname: dimension_size_p.bind(actual_args[arg_idx], dimension=dim_idx) for (vname, arg_idx, dim_idx) in known_dim_vars} dim_values = [solution[var].evaluate(known_env) for var in dim_vars] shape_constraints.check(known_env) return tuple(dim_values) def _solve_dim_equations( eqns: List[_DimEquation] ) -> Tuple[DimVarEnv, ShapeConstraints]: # Returns a shape environment and the shape constraints if it can solve all # dimension variables. Raises an exception if it cannot. shapeenv: DimVarEnv = {} shape_constraints = ShapeConstraints() def _shapeenv_to_str() -> str: if shapeenv: return (" Partial solution: " + ", ".join([f"{var} = {val}" for var, val in shapeenv.items()]) + ".") else: return "" def process_one_eqn(eqn: _DimEquation) -> bool: # We start with a DimEquation of the form `dim_expr = dim_value` # Try to rewrite the equation as `var * factor_var = dim_value_2` (a linear # uni-variate equation). Returns `False` if this rewrite fails. # Otherwise, compute the `var` value as `dim_value_2 // factor`, add it to # `shapeenv` and return `True`. # # Invariant: # var * factor_var + remaining_monomials_from_dim_expr = dim_value var, factor_var = None, None dim_value = eqn.dim_value for mon, factor in eqn.dim_expr.monomials(): # Perhaps we can already evaluate this monomial (all vars solved) try: mon_value = mon.evaluate(shapeenv) except KeyError: # `mon` still uses some variables not yet solved. We handle only the # case when `mon` is a single variable. v = mon.to_var() if v is not None and var is None: var, factor_var = v, factor continue else: dim_value = dim_value + core.dim_constant(-1) * _evaluate_multiply(mon_value, core.dim_constant(factor)) continue return False # This equation cannot yet be used to solve a variable if var is not None: if factor_var == 1: var_value = dim_value else: var_value, var_remainder = divmod(dim_value, core.dim_constant(factor_var)) # type: ignore shape_constraints.add_constraint( ShapeConstraint.Comparator.EQ, var_remainder, 0, make_err_msg=lambda rem_int, _: ( f"Dimension variable '{var}' must have integer value >= 1. " f"Non-zero remainder {rem_int} for factor {factor_var} when solving " f"{eqn}.{_shapeenv_to_str()}")) shape_constraints.add_constraint( ShapeConstraint.Comparator.GEQ, var_value, 1, make_err_msg=lambda var_int, _: ( f"Dimension variable '{var}' must have integer value >= 1. " f"Found {var_int} when " f"solving {eqn}.{_shapeenv_to_str()}")) if not isinstance(var_value, _DimExpr): assert var_value.dtype == core.dim_value_dtype() shapeenv[var] = var_value # type: ignore return True else: # All variables are resolved for this equation shape_constraints.add_constraint( ShapeConstraint.Comparator.EQ, eqn.dim_value, eqn.dim_expr.evaluate(shapeenv), make_err_msg=lambda val1, val2: ( f"Found inconsistency {val1} != {val2} when solving {eqn}.{_shapeenv_to_str()}")) return True while True: nr_eqns = len(eqns) eqns = [eqn for eqn in eqns if not process_one_eqn(eqn)] if not eqns: return shapeenv, shape_constraints # SUCCESS elif len(eqns) >= nr_eqns: break # We have some equations that we cannot solve further unsolved_vars: Set[str] = set() unsolved_polys: List[_DimExpr] = [] for eqn in eqns: unsolved_vars = unsolved_vars.union(eqn.dim_expr.get_vars()) unsolved_polys.append(eqn.dim_expr) unsolved_vars = unsolved_vars.difference(shapeenv.keys()) eqns_str = "\n ".join([str(eqn) for eqn in eqns]) err_msg = ( f"Cannot solve for values of dimension variables {unsolved_vars} from " f"the remaining dimension polynomials\n {eqns_str}.{_shapeenv_to_str()} " "Dimension variables can be solved only from linear uni-variate polynomials.\n" "\n" "Please see https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#dimension-variables-must-be-solvable-from-the-input-shapes for more details.") raise ValueError(err_msg)