[mgoldfarb-nvidia]

Summary

Adds jax.export.symbolic_dim_bounds, a public API for querying conservative
bounds that JAX can prove for a symbolic dimension or a derived expression.

The API exposes JAX's existing symbolic bounds decision procedure through a
supported public interface. Bounds are inclusive and may not be tight; infinite
endpoints mean JAX could not establish a finite bound.

This supports consumers that need a safe upper bound for intermediate symbolic
expressions, such as sizing workspace buffers during ahead-of-time kernel
compilation. This change does not add new constraint reasoning.

Testing

Added public API, typing, constraint-composition, and solver soundness coverage,
tested with both default and x64 configurations.

[gnecula]

Thank you for your contribution!

I am definitely interested in the fixes in the decision procedure. This made me look more closely at the logic for computing the bounds for a product of factors and I can see the bug that you found. I think that the fix in shape_poly_decision.py is a good one, and it seems to be sufficient. Let's split a smaller PR that contains only the changes to shape_poly_decision.py. Please include in this separate PR the corresponding unit test for the sign-crossing-case. I would like to also include a separate unit test along these lines:

def test_bounds_product(self):
a, _ = shape_poly.symbolic_shape("a, b")
# Interval [-2, 1] squared (even power) should have bounds [0, 4]
x = core.max_dim(-2, core.min_dim(1, a - 5))
self.assertEqual((0, 4), _bounds(x * x))

# Interval [-2, 1] cubed (odd power) should have bounds [-8, 1]
self.assertEqual((-8, 1), _bounds(x * x * x))

# Product of interval (-inf, -1] and [0, inf) should have bounds (-inf, 0]
# This exposes the 0 * inf -> NaN bug
u = (a // -2) * core.max_dim(0, a - 5)
self.assertEqual((-np.inf, 0), _bounds(u))

# Product of interval [-1, 1] and [1, inf) should have bounds (-inf, inf)
z = core.max_dim(-1, core.min_dim(1, a - 5))
self.assertEqual((-np.inf, np.inf), _bounds(z * a))

Once you split that other PR, this PR will contain only the additions to the APIs, which we should discuss further. I must say that I am unconvinced that it is worth extending the API for symbolic shapes. We care a lot in JAX about having the smallest API footprint possible to give us more flexibility to evolve JAX. What are the important use cases for having this extended API, and to what extent this is not possible with the existing API?

[mgoldfarb-nvidia]

Thank you for your contribution!

I am definitely interested in the fixes in the decision procedure. This made me look more closely at the logic for computing the bounds for a product of factors and I can see the bug that you found. I think that the fix in shape_poly_decision.py is a good one, and it seems to be sufficient. Let's split a smaller PR that contains only the changes to shape_poly_decision.py. Please include in this separate PR the corresponding unit test for the sign-crossing-case. I would like to also include a separate unit test along these lines:

def test_bounds_product(self): a, _ = shape_poly.symbolic_shape("a, b") # Interval [-2, 1] squared (even power) should have bounds [0, 4] x = core.max_dim(-2, core.min_dim(1, a - 5)) self.assertEqual((0, 4), _bounds(x * x))

# Interval [-2, 1] cubed (odd power) should have bounds [-8, 1]
self.assertEqual((-8, 1), _bounds(x * x * x))

# Product of interval (-inf, -1] and [0, inf) should have bounds (-inf, 0]
# This exposes the 0 * inf -> NaN bug
u = (a // -2) * core.max_dim(0, a - 5)
self.assertEqual((-np.inf, 0), _bounds(u))

# Product of interval [-1, 1] and [1, inf) should have bounds (-inf, inf)
z = core.max_dim(-1, core.min_dim(1, a - 5))
self.assertEqual((-np.inf, np.inf), _bounds(z * a))

Once you split that other PR, this PR will contain only the additions to the APIs, which we should discuss further. I must say that I am unconvinced that it is worth extending the API for symbolic shapes. We care a lot in JAX about having the smallest API footprint possible to give us more flexibility to evolve JAX. What are the important use cases for having this extended API, and to what extent this is not possible with the existing API?

@gnecula thank you for your review comments. I have split out the fix here: #40054

Regarding this new API. I have proposed it because we have cases where a CuteDSL kernel is compiled to an artifact and serialized into a custom call. During this process we need to allocate workspace tensors but to know the size we need concrete shapes that are resolved through custom compiled methods. When the user exports pure concrete shapes this works perfectly fine. For symbolic shapes we need to determine an upper bound for the workspace and for that we need to have access to an upper bound and other constraints "e.g. M <= 1024, M mod 8 == 0" etc. If we can not determine a size we of course have to error out.

Thank you for your contribution!

I am definitely interested in the fixes in the decision procedure. This made me look more closely at the logic for computing the bounds for a product of factors and I can see the bug that you found. I think that the fix in shape_poly_decision.py is a good one, and it seems to be sufficient. Let's split a smaller PR that contains only the changes to shape_poly_decision.py. Please include in this separate PR the corresponding unit test for the sign-crossing-case. I would like to also include a separate unit test along these lines:

def test_bounds_product(self): a, _ = shape_poly.symbolic_shape("a, b") # Interval [-2, 1] squared (even power) should have bounds [0, 4] x = core.max_dim(-2, core.min_dim(1, a - 5)) self.assertEqual((0, 4), _bounds(x * x))

# Interval [-2, 1] cubed (odd power) should have bounds [-8, 1]
self.assertEqual((-8, 1), _bounds(x * x * x))

# Product of interval (-inf, -1] and [0, inf) should have bounds (-inf, 0]
# This exposes the 0 * inf -> NaN bug
u = (a // -2) * core.max_dim(0, a - 5)
self.assertEqual((-np.inf, 0), _bounds(u))

# Product of interval [-1, 1] and [1, inf) should have bounds (-inf, inf)
z = core.max_dim(-1, core.min_dim(1, a - 5))
self.assertEqual((-np.inf, np.inf), _bounds(z * a))

Once you split that other PR, this PR will contain only the additions to the APIs, which we should discuss further. I must say that I am unconvinced that it is worth extending the API for symbolic shapes. We care a lot in JAX about having the smallest API footprint possible to give us more flexibility to evolve JAX. What are the important use cases for having this extended API, and to what extent this is not possible with the existing API?

@gnecula thank you for your review comments. I have split out the fix here: #40054

Regarding this new API. I have proposed it because we have cases where a CuteDSL kernel is compiled to an artifact and serialized into a custom call. During this process we need to allocate workspace tensors but to know the size we need concrete shapes that are resolved through custom compiled methods. When the user exports pure concrete shapes this works perfectly fine. For symbolic shapes we need to determine an upper bound for the workspace and for that we need to have access to an upper bound and other constraints "e.g. M <= 1024, M mod 8 == 0" etc.

My main concern with the current API is that is does not appear to have any publicly available stable function or class to make these queries. If you have a suggestion how we could be handle that I have no problem writing such functionality on my side. But I still need a stable structure and representation beyond the strings.

[gnecula]

The use case you describe justifies a symbolic_dim_bounds API. I have not tried to read carefully yet the implementation of this API, but you must have found cases where the existing _bounds_decision(e, BEST) was not sufficient?

For the upper bounds use case you described, do you need to compute it only for the inputs and outputs, or also for intermediate results? For inputs and outputs, I assume that the symbolic expressions are relatively simple, right?

[gnecula]

For the other introspection APIs, can't they be implemented in terms of the existing APIs? E.g.,

def dim_at_least(e1, e2):
   try:
       return e1 >= e2
   except InconclusiveDimensionOperation:
      return None

[mgoldfarb-nvidia]

@gnecula I have updated the MR to reduce down to a single new method: symbolic_dim_bounds that exposes the private _bounds_decision. I agree that for this use case we don't need to extra tightened lower bound since this is for workspace tensors.

[gnecula]

@mgoldfarb-nvidia can you please look at the documentation test failure?

[mgoldfarb-nvidia]

@mgoldfarb-nvidia can you please look at the documentation test failure?

Fixed the type reference and ran doc generation locally and it worked.