jax.export.symbolic_shape

jax.export.symbolic_shape#

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

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

参见 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)。

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

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

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

返回类型:

Sequence[DimSize]

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

内容