jax.export.symbolic_shape#

jax.export.symbolic_shape(shape_spec, *, constraints=(), scope=None, like=None)[源代码]#

从字符串表示形式构造一个符号形状。

有关示例，请参阅https://jax.ac.cn/en/latest/export/shape_poly.html。

参数:

  • shape_spec (str | None) – 符号形状规范。None 表示 “…” 。形状规范是一个元组（括号是可选的）的字符串表示形式，其中包含以逗号分隔的维度表达式。维度表达式可以是：整数常量、维度变量（以字母开头的字母数字）、e1 + e2、e1 - e2、e1 * e2、floordiv(e1, e2)、mod(e1, e2)、max(e1, e2) 或 min(e1, e2)。

  • constraints (Sequence[str]) – 一系列关于符号维度表达式的约束，形式为 e1 >= e2 或 e1 <= e2 或 e1 == e2。有关用法，请参阅 [文档](https://jax.ac.cn/en/latest/export/shape_poly.html#user-specified-symbolic-constraints)。

  • scope (SymbolicScope | None | None) – （可选）您可以指定在给定作用域中创建解析的符号表达式。如果缺少此项，则会使用给定的 constraints 创建一个新的 SymbolicScope。您不能同时指定 scope 和 constraints。 有关用法，请参阅 [文档](https://jax.ac.cn/en/latest/export/shape_poly.html#user-specified-symbolic-constraints)。

  • like (Sequence[int | None] | None | None) – 当 shape_spec 包含占位符（“_”, “…”）时，使用此形状来填充占位符。like 中用于填充的维度必须不是 None。如果 like 中的维度不是 None，并且 shape_spec 中对应的维度是常量，则它们必须相等。

返回类型:

Sequence[DimSize]

返回: 包含整数或涉及维度变量的符号表达式的元组。