形状多态性#
当 JAX 在 JIT 模式下使用时，将为每个输入类型和形状组合跟踪、降低到 StableHLO 并编译一个函数。导出函数并在另一个系统上反序列化后，我们不再拥有 Python 源代码，因此无法重新跟踪和重新降低它。**形状多态性** 是 JAX 导出的一项功能，允许一些导出的函数用于整个输入形状族。这些函数在导出期间只跟踪和降低一次，并且 Exported 对象包含能够在许多具体输入形状上编译和执行函数所需的信息。我们通过在导出时指定包含维度变量（符号形状）的形状来实现这一点，如下例所示
>>> import jax
>>> from jax import export
>>> from jax import numpy as jnp
>>> def f(x): # f: f32[a, b]
... return jnp.concatenate([x, x], axis=1)
>>> # We construct symbolic dimension variables.
>>> a, b = export.symbolic_shape("a, b")
>>> # We can use the symbolic dimensions to construct shapes.
>>> x_shape = (a, b)
>>> x_shape
(a, b)
>>> # Then we export with symbolic shapes:
>>> exp: export.Exported = export.export(jax.jit(f))(
... jax.ShapeDtypeStruct(x_shape, jnp.int32))
>>> exp.in_avals
(ShapedArray(int32[a,b]),)
>>> exp.out_avals
(ShapedArray(int32[a,2*b]),)
>>> # We can later call with concrete shapes (with a=3 and b=4), without re-tracing `f`.
>>> res = exp.call(np.ones((3, 4), dtype=np.int32))
>>> res.shape
(3, 8)
请注意，此类函数在调用时仍会根据每个具体输入形状按需重新编译。只有跟踪和降低被保存。
上述示例中使用了 jax.export.symbolic_shape() 将符号形状的字符串表示解析为维度表达式对象（类型为 _DimExpr），这些对象可用于代替整数常量来构建形状。维度表达式对象重载了大多数整数运算符，因此您可以在大多数情况下像使用整数常量一样使用它们。有关详细信息，请参阅 使用维度变量计算。
此外，我们提供 jax.export.symbolic_args_specs()，它可用于基于多态形状规范构造 jax.ShapeDtypeStruct 对象的 pytrees。
>>> def f1(x, y): # x: f32[a, 1], y : f32[a, 4]
... return x + y
>>> # Assuming you have some actual args with concrete shapes
>>> x = np.ones((3, 1), dtype=np.int32)
>>> y = np.ones((3, 4), dtype=np.int32)
>>> args_specs = export.symbolic_args_specs((x, y), "a, ...")
>>> exp = export.export(jax.jit(f1))(* args_specs)
>>> exp.in_avals
(ShapedArray(int32[a,1]), ShapedArray(int32[a,4]))
请注意，多态形状规范 "a, ..." 包含占位符 ...，用于从参数 (x, y) 的具体形状的具体形状中填充。占位符 ... 代表 0 个或多个维度，而占位符 _ 代表一个维度。 jax.export.symbolic_args_specs() 支持参数的 pytrees，这些参数用于填充数据类型和任何占位符。该函数将构造一个参数规范 (jax.ShapeDtypeStruct) 的 pytree，与传递给它的参数的结构匹配。多态形状规范可以是 pytree 前缀，在多个参数应应用一个规范的情况下，如上例所示。请参阅 如何将可选参数与参数匹配。
一些形状规范示例
("(b, _, _)", None)可用于具有两个参数的函数，第一个参数为 3D 数组，具有一个符号批处理前导维度。第一个参数的其他维度以及第二个参数的形状根据实际参数进行专门化。请注意，如果第一个参数是 3D 数组的 pytree，则相同的规范将起作用，所有这些数组都具有相同的领先维度，但尾部维度可能不同。第二个参数的None值表示该参数不是符号。等效地，可以使用...。("(batch, ...)", "(batch,)")指定两个参数具有匹配的前导维度，第一个参数的秩至少为 1，第二个参数的秩为 1。
形状多态性的正确性#
我们希望相信导出的程序在针对任何适用的具体形状编译和执行时会产生与原始 JAX 程序相同的结果。更准确地说
对于任何 JAX 函数 f 和任何包含符号形状的参数规范 arg_spec，以及任何形状与 arg_spec 匹配的具体参数 arg
如果 JAX 本机执行在具体参数上成功：
res = f(arg),并且如果导出使用符号形状成功：
exp = export.export(f)(arg_spec),那么编译和运行导出将成功，结果相同：
res == exp.call(arg)
必须了解，f(arg) 有权重新调用 JAX 跟踪机制，实际上它对每个不同的具体 arg 形状都这样做，而 exp.call(arg) 的执行不再可以使用 JAX 跟踪（此执行可能发生在 f 的源代码不可用的环境中）。
确保这种形式的正确性很困难，在最困难的情况下，导出会失败。本章的其余部分描述了如何处理这些失败。
使用维度变量计算#
JAX 会跟踪所有中间结果的形状。当这些形状取决于维度变量时，JAX 会将它们计算为包含维度变量的符号维度表达式。维度变量代表大于或等于 1 的整数值。符号表达式可以表示对维度表达式和整数（int、np.int 或可通过 operator.index 转换的任何内容）应用算术运算符（加、减、乘、地板除、模，包括 NumPy 变体 np.sum、np.prod 等）的结果。然后，这些符号维度可以在 JAX 原语和 API 的形状参数中使用，例如，在 jnp.reshape、jnp.arange、切片索引等中。
例如，在以下用于扁平化 2D 数组的代码中，计算 x.shape[0] * x.shape[1] 将符号维度 4 * b 计算为新的形状
>>> f = lambda x: jnp.reshape(x, (x.shape[0] * x.shape[1],))
>>> arg_spec = jax.ShapeDtypeStruct(export.symbolic_shape("b, 4"), jnp.int32)
>>> exp = export.export(jax.jit(f))(arg_spec)
>>> exp.out_avals
(ShapedArray(int32[4*b]),)
可以使用 jnp.array(x.shape[0]) 或甚至 jnp.array(x.shape) 将维度表达式显式转换为 JAX 数组。这些操作的结果可以用作常规 JAX 数组，但不能再用作形状中的维度。
>>> exp = export.export(jax.jit(lambda x: jnp.array(x.shape[0]) + x))(
... jax.ShapeDtypeStruct(export.symbolic_shape("b"), np.int32))
>>> exp.call(jnp.arange(3, dtype=np.int32))
Array([3, 4, 5], dtype=int32)
>>> exp = export.export(jax.jit(lambda x: x.reshape(jnp.array(x.shape[0]) + 2)))(
... jax.ShapeDtypeStruct(export.symbolic_shape("b"), np.int32))
Traceback (most recent call last):
TypeError: Shapes must be 1D sequences of concrete values of integer type, got [Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>].
当符号维度在与非整数的算术运算中使用时，例如 float、np.float、np.ndarray 或 JAX 数组，它将自动使用 jnp.array 转换为 JAX 数组。例如，在下面的函数中，所有 x.shape[0] 的出现都将隐式转换为 jnp.array(x.shape[0])，因为它们与非整数标量或 JAX 数组的运算有关
>>> exp = export.export(jax.jit(
... lambda x: (5. + x.shape[0],
... x.shape[0] - np.arange(5, dtype=jnp.int32),
... x + x.shape[0] + jnp.sin(x.shape[0]))))(
... jax.ShapeDtypeStruct(export.symbolic_shape("b"), jnp.int32))
>>> exp.out_avals
(ShapedArray(float32[], weak_type=True),
ShapedArray(int32[5]),
ShapedArray(float32[b], weak_type=True))
>>> exp.call(jnp.ones((3,), jnp.int32))
(Array(8., dtype=float32, weak_type=True),
Array([ 3, 2, 1, 0, -1], dtype=int32),
Array([4.14112, 4.14112, 4.14112], dtype=float32, weak_type=True))
另一个典型示例是计算平均值（注意 x.shape[0] 如何自动转换为 JAX 数组）
>>> exp = export.export(jax.jit(
... lambda x: jnp.sum(x, axis=0) / x.shape[0]))(
... jax.ShapeDtypeStruct(export.symbolic_shape("b, c"), jnp.int32))
>>> exp.call(jnp.arange(12, dtype=jnp.int32).reshape((3, 4)))
Array([4., 5., 6., 7.], dtype=float32)
形状多态性存在时的错误#
大多数 JAX 代码假设 JAX 数组的形状是整数元组，但在形状多态性中，某些维度可能是符号表达式。这会导致许多错误。例如，我们可以有常见的 JAX 形状检查错误
>>> v, = export.symbolic_shape("v,")
>>> export.export(jax.jit(lambda x, y: x + y))(
... jax.ShapeDtypeStruct((v,), dtype=np.int32),
... jax.ShapeDtypeStruct((4,), dtype=np.int32))
Traceback (most recent call last):
TypeError: add got incompatible shapes for broadcasting: (v,), (4,).
>>> export.export(jax.jit(lambda x: jnp.matmul(x, x)))(
... jax.ShapeDtypeStruct((v, 4), dtype=np.int32))
Traceback (most recent call last):
TypeError: dot_general requires contracting dimensions to have the same shape, got (4,) and (v,).
我们可以通过指定参数具有形状 (v, v) 来修复上述矩阵乘法示例。
符号维度的比较部分支持#
在 JAX 内部，有很多涉及形状的相等和不相等比较，例如，用于进行形状检查，甚至用于为某些原语选择实现。比较支持如下
相等性支持，但有警告：如果两个符号维度在维度变量的所有赋值下都表示相同的值，则相等性评估为
True，例如，对于b + b == 2*b；否则，相等性评估为False。有关此行为的重要后果的讨论，请参阅下文。不相等性始终是相等性的否定。
不等式部分支持，类似于部分相等性。但是，在这种情况下，我们考虑维度变量在严格的正整数范围内。例如，
b >= 1、b >= 0、2 * a + b >= 3为True，而b >= 2、a >= b、a - b >= 0不确定，导致异常。
在无法将比较操作解析为布尔值的情况下，我们将引发 InconclusiveDimensionOperation。例如，
import jax
>>> export.export(jax.jit(lambda x: 0 if x.shape[0] + 1 >= x.shape[1] else 1))(
... jax.ShapeDtypeStruct(export.symbolic_shape("a, b"), dtype=np.int32)) # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
jax._src.export.shape_poly.InconclusiveDimensionOperation: Symbolic dimension comparison 'a + 1' >= 'b' is inconclusive.
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.
如果您确实获得了 InconclusiveDimensionOperation，您可以尝试几种策略
如果您的代码使用内置的
max或min，或np.max或np.min，那么您可以将它们替换为core.max_dim和core.min_dim，它们的作用是将不等式比较延迟到编译时，此时形状变为已知。尝试使用
core.max_dim和core.min_dim重新编写条件语句，例如，而不是d if d > 0 else 0，您可以编写core.max_dim(d, 0)。尝试重写代码，使其对维度必须为整数的依赖性降低，并依靠符号维度在大多数算术运算中与整数具有相同类型的特点。例如，不要写
int(d) + 5，而是写d + 5。指定符号约束，如下所述。
用户指定的符号约束#
默认情况下，JAX 假设所有维度变量的值都大于或等于 1，并尝试从该假设中推导出其他简单的不等式，例如：
a + 2 >= 3,a * 2 >= 1,a + b + c >= 3,a // 4 >= 0，a**2 >= 1，等等。
如果您将符号形状规范更改为为维度大小添加**隐式**约束，则可以避免一些不等式比较失败。例如：
您可以使用
2*b来表示维度，将其约束为偶数且大于或等于 2。您可以使用
b + 15来表示维度，将其约束为至少为 16。例如，以下代码如果没有+ 15部分将失败，因为 JAX 会想要验证切片大小是否不超过轴大小。
>>> _ = export.export(jax.jit(lambda x: x[0:16]))(
... jax.ShapeDtypeStruct(export.symbolic_shape("b + 15"), dtype=np.int32))
这种隐式符号约束用于决定比较，并在编译时进行检查，如下所述。
您还可以指定**显式**符号约束
>>> # Introduce dimension variable with constraints.
>>> a, b = export.symbolic_shape("a, b",
... constraints=("a >= b", "b >= 16"))
>>> _ = export.export(jax.jit(lambda x: x[:x.shape[1], :16]))(
... jax.ShapeDtypeStruct((a, b), dtype=np.int32))
这些约束与隐式约束一起构成一个合取。您可以指定 >=，<= 和 == 约束。目前，JAX 对处理符号约束的支持有限
当约束形式为变量大于或等于或小于或等于常数时，您可以最大限度地利用这些约束。例如，从约束
a >= 16和b >= 8，我们可以推断出a + 2*b >= 32。当约束涉及更复杂的表达式时，您的能力将受到限制，例如，从
a >= b + 8，我们可以推断出a - b >= 8，但不能推断出a >= 9。我们可能会在将来对这方面进行一些改进。等式约束被视为重写规则：每当遇到
==左侧的符号表达式时，它都会被重写为右侧的表达式。例如，floordiv(a, b) == c通过将所有出现的floordiv(a, b)替换为c来实现。等式约束的左侧不能包含顶层的加法或减法。有效的左侧示例包括a * b，或4 * a，或floordiv(a + c, b)。
>>> # Introduce dimension variable with equality constraints.
>>> a, b, c, d = export.symbolic_shape("a, b, c, d",
... constraints=("a * b == c + d",))
>>> 2 * b * a
2*d + 2*c
>>> a * b * b
b*d + b*c
符号约束还有助于解决 JAX 推理机制中的限制。例如，在下面的代码中，JAX 将尝试证明切片大小 x.shape[0] % 3（即符号表达式 mod(b, 3)）小于或等于轴大小（即 b）。对于 b 的所有严格正值，这都是正确的，但 JAX 的符号比较规则无法证明这一点。因此，以下代码会引发错误
from jax import lax
>>> b, = export.symbolic_shape("b")
>>> f = lambda x: lax.slice_in_dim(x, 0, x.shape[0] % 3)
>>> export.export(jax.jit(f))(
... jax.ShapeDtypeStruct((b,), dtype=np.int32)) # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
jax._src.export.shape_poly.InconclusiveDimensionOperation: Symbolic dimension comparison 'b' >= 'mod(b, 3)' is inconclusive.
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.
这里一个选择是限制代码仅在轴大小为 3 的倍数时工作（通过在形状中将 b 替换为 3*b）。然后，JAX 将能够将模运算 mod(3*b, 3) 简化为 0。另一个选择是添加一个符号约束，包含 JAX 试图证明的准确的不确定不等式
>>> b, = export.symbolic_shape("b",
... constraints=["b >= mod(b, 3)"])
>>> f = lambda x: lax.slice_in_dim(x, 0, x.shape[0] % 3)
>>> _ = export.export(jax.jit(f))(
... jax.ShapeDtypeStruct((b,), dtype=np.int32))
与隐式约束一样，显式符号约束在编译时进行检查，使用与如下所述相同的机制。
符号维度范围#
符号约束存储在一个 αn jax.export.SymbolicScope 对象中，该对象是为每个 jax.export.symbolic_shapes() 调用隐式创建的。您必须小心不要混合使用不同范围的符号表达式。例如，以下代码将失败，因为 a1 和 a2 使用不同的范围（由不同的 jax.export.symbolic_shape() 调用创建）。
>>> a1, = export.symbolic_shape("a,")
>>> a2, = export.symbolic_shape("a,", constraints=("a >= 8",))
>>> a1 + a2
Traceback (most recent call last):
ValueError: Invalid mixing of symbolic scopes for linear combination.
Expected scope 4776451856 created at <doctest shape_poly.md[31]>:1:6 (<module>)
and found for 'a' (unknown) scope 4776979920 created at <doctest shape_poly.md[32]>:1:6 (<module>) with constraints:
a >= 8
源自单个 jax.export.symbolic_shape() 调用的符号表达式共享一个范围，可以在算术运算中混合使用。结果也将共享相同的范围。
您可以重复使用范围
>>> a, = export.symbolic_shape("a,", constraints=("a >= 8",))
>>> b, = export.symbolic_shape("b,", scope=a.scope) # Reuse the scope of `a`
>>> a + b # Allowed
b + a
您还可以显式创建范围
>>> my_scope = export.SymbolicScope()
>>> c, = export.symbolic_shape("c", scope=my_scope)
>>> d, = export.symbolic_shape("d", scope=my_scope)
>>> c + d # Allowed
d + c
JAX 追踪使用部分由形状作为键的缓存，如果符号形状打印相同，但使用不同的范围，它们将被视为不同。
等式比较的警告#
等式比较对于 b + 1 == b 或 b == 0（在这种情况下可以确定维度对于维度变量的所有值都不同）返回 False，但对于 b == 1 和 a == b 也返回 False。这是不合理的，我们应该引发 core.InconclusiveDimensionOperation，因为在某些赋值下结果应该为 True，而在其他赋值下结果应该为 False。我们选择使等式成为全等式，从而允许不合理性，因为否则我们在对维度表达式或包含它们的物体（形状、core.AbstractValue、core.Jaxpr）进行哈希时，可能会遇到哈希冲突而导致错误。除了哈希错误外，等式的部分语义会导致以下表达式出现错误 b == a or b == b 或 b in [a, b]，即使我们更改了比较顺序，也能避免错误。
形式为 if x.shape[0] != 1: raise NiceErrorMessage 的代码即使在这种等式处理方式下也是合理的，但形式为 if x.shape[0] != 1: return 1 的代码是不合理的。
维度变量必须能够从输入形状中求解#
目前，在调用导出对象时传递维度变量值的唯一方法是通过数组参数的形状间接进行。例如，b 的值可以在调用站点从类型为 f32[b] 的第一个参数的形状中推断出来。这适用于大多数用例，并且它反映了 JIT 函数的调用约定。
有时您可能想要导出一个由整数参数化的函数，该整数确定程序中的某些形状。例如，我们可能想要导出下面定义的函数 my_top_k，它由 k 的值参数化，该值确定结果的形状。以下尝试将导致错误，因为维度变量 k 无法从输入 x: i32[4, 10] 的形状中推断出来
>>> def my_top_k(k, x): # x: i32[4, 10], k <= 10
... return lax.top_k(x, k)[0] # : i32[4, 3]
>>> x = np.arange(40, dtype=np.int32).reshape((4, 10))
>>> # Export with static `k=3`. Since `k` appears in shapes it must be in `static_argnums`.
>>> exp_static_k = export.export(jax.jit(my_top_k, static_argnums=0))(3, x)
>>> exp_static_k.in_avals[0]
ShapedArray(int32[4,10])
>>> exp_static_k.out_avals[0]
ShapedArray(int32[4,3])
>>> # When calling the exported function we pass only the non-static arguments
>>> exp_static_k.call(x)
Array([[ 9, 8, 7],
[19, 18, 17],
[29, 28, 27],
[39, 38, 37]], dtype=int32)
>>> # Now attempt to export with symbolic `k` so that we choose `k` after export.
>>> k, = export.symbolic_shape("k", constraints=["k <= 10"])
>>> export.export(jax.jit(my_top_k, static_argnums=0))(k, x)
Traceback (most recent call last):
UnexpectedDimVar: "Encountered dimension variable 'k' that is not appearing in the shapes of the function arguments
将来，我们可能会添加一种额外的机制来传递维度变量的值，除了通过输入形状隐式传递之外。与此同时，针对上述用例的解决方法是用一个形状为 (0, k) 的数组来替换函数参数 k，这样就可以从数组的输入形状中推断出 k。第一维为 0 以确保整个数组为空，并且在调用导出函数时不会产生任何性能损耗。
>>> def my_top_k_with_dimensions(dimensions, x): # dimensions: i32[0, k], x: i32[4, 10]
... return my_top_k(dimensions.shape[1], x)
>>> exp = export.export(jax.jit(my_top_k_with_dimensions))(
... jax.ShapeDtypeStruct((0, k), dtype=np.int32),
... x)
>>> exp.in_avals
(ShapedArray(int32[0,k]), ShapedArray(int32[4,10]))
>>> exp.out_avals[0]
ShapedArray(int32[4,k])
>>> # When we invoke `exp` we must construct and pass an array of shape (0, k)
>>> exp.call(np.zeros((0, 3), dtype=np.int32), x)
Array([[ 9, 8, 7],
[19, 18, 17],
[29, 28, 27],
[39, 38, 37]], dtype=int32)
您可能还会在某些维度变量出现在输入形状中，但以 JAX 目前无法解决的非线性表达式出现时，遇到错误
>>> a, = export.symbolic_shape("a")
>>> export.export(jax.jit(lambda x: x.shape[0]))(
... jax.ShapeDtypeStruct((a * a,), dtype=np.int32))
Traceback (most recent call last):
ValueError: Cannot solve for values of dimension variables {'a'}.
We can only solve linear uni-variate constraints.
Using the following polymorphic shapes specifications: args[0].shape = (a^2,).
Unprocessed specifications: 'a^2' for dimension size args[0].shape[0].
形状断言错误#
JAX 假设维度变量的取值范围是严格的正整数，当代码被编译成具体输入形状时，会检查这一假设。
例如，给定符号输入形状 (b, b, 2*d)，JAX 将在使用实际参数 arg 调用时生成代码来检查以下断言
arg.shape[0] >= 1arg.shape[1] == arg.shape[0]arg.shape[2] % 2 == 0arg.shape[2] // 2 >= 1
例如，当我们在形状为 (3, 3, 5) 的参数上调用导出函数时，会遇到以下错误。
>>> def f(x): # x: f32[b, b, 2*d]
... return x
>>> exp = export.export(jax.jit(f))(
... jax.ShapeDtypeStruct(export.symbolic_shape("b, b, 2*d"), dtype=np.int32))
>>> exp.call(np.ones((3, 3, 5), dtype=np.int32))
Traceback (most recent call last):
ValueError: Input shapes do not match the polymorphic shapes specification.
Division had remainder 1 when computing the value of 'd'.
Using the following polymorphic shapes specifications:
args[0].shape = (b, b, 2*d).
Obtained dimension variables: 'b' = 3 from specification 'b' for dimension args[0].shape[0] (= 3), .
Please see https://jax.net.cn/en/latest/export/shape_poly.html#shape-assertion-errors for more details.
这些错误出现在编译之前的预处理步骤中。
对符号维度的除法运算仅部分支持#
JAX 将尝试简化除法和取模运算，例如，(a * b + a) // (b + 1) == a 和 (6 * a + 4) % 3 == 1。特别地，当 (a) 没有余数，或 (b) 除数为常数时，JAX 将处理这些情况，在这种情况下可能会有一个常数余数。
例如，以下代码在尝试计算 reshape 操作的推断维度时会导致除法错误。
>>> b, = export.symbolic_shape("b")
>>> export.export(jax.jit(lambda x: x.reshape((2, -1))))(
... jax.ShapeDtypeStruct((b,), dtype=np.int32))
Traceback (most recent call last):
jax._src.core.InconclusiveDimensionOperation: Cannot divide evenly the sizes of shapes (b,) and (2, -1).
The remainder mod(b, - 2) should be 0.
注意，以下操作将成功。
>>> b, = export.symbolic_shape("b")
>>> # We specify that the first dimension is a multiple of 4
>>> exp = export.export(jax.jit(lambda x: x.reshape((2, -1))))(
... jax.ShapeDtypeStruct((4*b,), dtype=np.int32))
>>> exp.out_avals
(ShapedArray(int32[2,2*b]),)
>>> # We specify that some other dimension is even
>>> exp = export.export(jax.jit(lambda x: x.reshape((2, -1))))(
... jax.ShapeDtypeStruct((b, 5, 6), dtype=np.int32))
>>> exp.out_avals
(ShapedArray(int32[2,15*b]),)
调试#
首先，请参阅 调试 文档。此外，您还可以调试形状细化，它是在编译时为具有维度变量或多平台支持的模块调用的。
如果形状细化过程中出现错误，您可以设置 JAX_DUMP_IR_TO 环境变量，以便在形状细化之前查看 HLO 模块的转储（命名为 ..._before_refine_polymorphic_shapes.mlir）。此模块应该已经具有静态输入形状。
要启用所有形状细化阶段的日志记录，您可以在 OSS 中设置环境变量 TF_CPP_VMODULE=refine_polymorphic_shapes=3（在 Google 内部，您传递 --vmodule=refine_polymorphic_shapes=3）。
# Log from python
JAX_DUMP_IR_TO=/tmp/export.dumps/ TF_CPP_VMODULE=refine_polymorphic_shapes=3 python tests/shape_poly_test.py ShapePolyTest.test_simple_unary -v=3