jax.export.symbolic_shape#

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

Constructs a symbolic shape from a string representation.

See https://docs.jax.dev/en/latest/export/shape_poly.html for examples.

Parameters:

  • shape_spec (str | None) – a symbolic shape specification. None stands for “…”. A shape specification is the string representation of a tuple (the parentheses are optional) with comma-separated dimension expressions. A dimension expression can be either: an integer constant, a dimension variable (alphanumeric starting with a letter), e1 + e2, e1 - e2, e1 * e2, floordiv(e1, e2), mod(e1, e2), max(e1, e2), or min(e1, e2).

  • constraints (Sequence[str]) – a sequence of constraints on symbolic dimension expressions, of the form e1 >= e2 or e1 <= e2, or e1 == e2. See [the documentation](https://docs.jax.dev/en/latest/export/shape_poly.html#user-specified-symbolic-constraints) for usage.

  • scope (SymbolicScope | None | None) – optionally, you can specify that the parsed symbolic expressions be created in the given scope. If this is missing, then a new SymbolicScope is created with the given constraints. You cannot specify both a scope and constraints. See [the documentation](https://docs.jax.dev/en/latest/export/shape_poly.html#user-specified-symbolic-constraints) for usage.

  • like (Sequence[int | None] | None | None) – when shape_spec contains placeholders (“_”, “…”), use this shape to fill in the placeholders. The dimensions of like that are used for filling must be not None. If a dimension in like is not None and the corresponding dimension in shape_spec is a constant then they must be equal.

Return type:

Sequence[DimSize]

Returns: a tuple with integers or symbolic expressions involving dimension variables.