Data dependent/unknown shapes

[shoyer]

Some libraries, particularly those with a graph-based computational model (e.g., Dask and TensorFlow), have support for "unknown" or "data dependent" shapes, e.g., due to boolean indexing such as x[y > 0] (#84). Other libraries (e.g., JAX and Dask in some cases) do not support some operations because they would produce such data dependent shapes.

We should consider a standard way to represent these shapes in shape attributes, ideally some extension of the "tuple of integer" format used for fully known shapes. For example, TensorFlow and Dask currently use different representations:

• TensorFlow uses a custom TensorShape object (which acts very similarly to tuple), where some values may be None
• Dask uses tuples, where some values may be nan integer of integers

[rgommers]

That makes sense. A custom tuple-like object, with None for not-yet-know dimensions, seems a bit more flexible and appropriate to me.

[rgommers]

We actually already allow for either a tuple or a custom shape object (https://data-apis.github.io/array-api/latest/API_specification/array_object.html#shape), so it's None vs. nan that needs deciding.

[shoyer]

Remind me why we allow the custom object? Is it just because that's what TensorFlow currently does or is there a more specific use case?

…

On Thu, Dec 3, 2020 at 1:30 PM Ralf Gommers ***@***.***> wrote: We actually already allow for either a tuple or a custom shape object ( https://data-apis.github.io/array-api/latest/API_specification/array_object.html#shape), so it's None vs. nan that needs deciding. — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#97 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJJFVRG5AFMXEEYJPLZ5R3SS77N3ANCNFSM4UMQRVAQ> .

[rgommers]

IIRC because using Python objects (ints, tuples, etc.) rather than custom ones can force extra synchronizations. PyTorch also uses a custom object:

>>> t = torch.tensor([[1, 2], [3, 4]])                                                         
>>> t.shape                                                                                    
torch.Size([2, 2])

[shoyer]

IIRC because using Python objects (ints, tuples, etc.) rather than custom ones can force extra synchronizations

I guess this is for cases where some parts of shape may require on-device computation? E.g., z = x[:, y > 0], then len(z.shape) == 2 and z.shape[0] == x.shape[0] (no synchronization required), but z.shape[1] may be expensive to calculate.

Interestingly, I think TensorFlow has this for different reasons, possibly just for preserving backwards compatibility at this point. It has two ways of accessing shape information: tf.shape(x) which returns the dynamic shape as a Tensor and x.shape which returns the static shape (possibly including None for dynamic sizes).

[rgommers]

It may be due to legacy reasons in PyTorch as well. From pytorch/pytorch#20823 (comment):
Maybe it's time to consider removing torch.Size and make tensor.shape return a tuple? I believe torch.Size was necessary to support legacy constructors, like torch.Tensor(a.size()), which overloaded torch.Tensor with both a size and a data constructor.

I remember a discussion about this when we reviewed the array object API, and someone (probably @alextp or @apaszke) suggested adding the custom object.

[alextp]

I think the custom object is only useful when you want to deal with staged computation like @shoyer mentions above. TF2 stopped using the wrapped Dimension type and just has its TensorShape be a tuple of either integers or None to indicate shapes which we do not know at tracing time. So I think saying shapes are tuples here is probably a good idea.

[agarwal-ashish]

TensorFlow (and TF NumPy) will be returning a TensorShape object and tensor.shape is None will not work when checking for unknown rank.

Perhaps we should reconsider allowing custom objects.
Alternatively users could check for ndarray.ndim is None leaving ndarray.shape undefined for unknown rank cases.

[alextp]

I think this is a bug in tf.numpy, then.

…

On Wed, Dec 9, 2020 at 11:08 PM token ***@***.***> wrote: TensorFlow (and TF NumPy) will be returning a TensorShape object and tensor.shape is None will not work when checking for unknown rank. Perhaps we should reconsider allowing custom objects. Alternatively users could check for ndarray.ndim is None leaving ndarray.shape undefined for unknown rank cases. — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#97 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAABHRP6KUAHRFPNCVZRBODSUBXVHANCNFSM4UMQRVAQ> .

-- - Alex

[rgommers]

Alternatively users could check for ndarray.ndim is None leaving ndarray.shape undefined for unknown rank cases.

I don't think we included anything in the standard that could cause this. Indexing may make the size of one or more dimensions unknown, but ndim should always be known because for data-dependent operations there's no squeezing:

>>> x = np.ones((3,2))
>>> x[:0, np.zeros(2, dtype=bool)]
array([], shape=(0, 0), dtype=float64)

After writing this: we did include squeeze() so a combination of indexing and explicit squeezing could indeed cause this.

Leaving any property (like .shape) undefined seems unhealthy. I'd say in this case, use ndim = None and shape is None. If dimensionality is known but exact shape isn't, use None for the unknown dimensions (e.g., shape = (3, None, 100)).

and just has its TensorShape be a tuple of either integers or None to indicate shapes which we do not know at tracing time. So I think saying shapes are tuples here is probably a good idea.

Maybe we should say "shape is tuple" and add a note that if for backwards compat reasons an implementation is using a custom object, then it should make sure that it is a subtype of tuple so that it works when users annotate code using .shape with Tuple.

[shoyer]

After writing this: we did include squeeze() so a combination of indexing and explicit squeezing could indeed cause this.

This might be an argument for requiring an axis argument for squeeze(). In my opinion, that's pretty much always a good idea for library code, anyways. I've found cases where libraries use squeeze (like matplotlib's plt.subplots) to be quite annoying/inconsistent.

Leaving any property (like .shape) undefined seems unhealthy. I'd say in this case, use ndim = None and shape is None. If dimensionality is known but exact shape isn't, use None for the unknown dimensions (e.g., shape = (3, None, 100)).

Agreed!

[agarwal-ashish]

If we insist that rank is always statically known, then shape being None is moot and we can define shape to be a tuple (or tuple-like) object.

Regarding custom objects, if the object represents a union of tuple and None, then it doesn't sound correct to subclass this union from tuple: a subclass of tuple should not satisfy is None, I'd think.

TensorFlow's TensorShape object implements most of the tuple methods but is not subclassed from it since it represents an optional tuple. Instead it raises exception in certain methods when rank is unknown.

[rgommers]

Regarding custom objects, if the object represents a union of tuple and None, then it doesn't sound correct to subclass this union from tuple

Agreed. I said subtype rather than subclass (as in structural subtyping). To clarify, I'd like it to be possible for users to use type annotation Tuple or Optional[Tuple] which also works with the custom object, rather than have to use some more complicated construct.

This implies that the custom object must behave tuple-like; you can index it, += it with another tuple, use len on it, etc.

[rgommers]

This might be an argument for requiring an axis argument for squeeze(). In my opinion, that's pretty much always a good idea for library code, anyways. I've found cases where libraries use squeeze (like matplotlib's plt.subplots) to be quite annoying/inconsistent.

I like that idea, I'll open a PR so we can discuss there.