Introducing dims syntax?

[srush]

I'm finding this part a bit clunky in my code, and it seems like a great source of errors if you mess up this order.

batch, channel, width, height = dims(4)
input_positional = torch.rand(2, 3, 224, 224)
input = input_positional[batch, channel, width, height]

Kind of curious whether something like this is possible? I guess it is important to start with unordered for various reasons, but it is a bit awkward.

batch, channel, width, height = dims(size=[2, 3, 224, 224])
input_positional = torch.rand(batch, channel, width, height)

[Balandat]

I came here to post the same feedback / suggestion, glad to see someone was faster. It would be great if something along those lines would be possible, I think that would clean things up quite a bit (and help avoid creating multiple tensors).

[zdevito]

Yeah, that looks quite nice. I think the semantics are pretty clear: in any tensor constructor where we currently specify the size with an integer, a dim can be substituted and this should have the effect of constructing a tensor with that dim.

A related piece of functionality is being able to create tensors with a similar to shape to others that have dimensions. Things like rand_like already work, but it would nice to get a single object that holds both the dims and the sizes for positional dimensions. Given the above semantics, that could just be a list with a mix of both (e.g. [batch, 1024]) but it couldn't replace something like input.shape because one would expect input.shape[0] to be the first positional dimension and not batch.

[srush]

Huh rand_like is quite interesting, I hadn't thought of that. Part of me thinks that function shouldn't exist. If the batch dim is out of scope, you shouldn't be able to manipulate it by creating a new randomness along that dim. Jax for instance would not allow you to do this, You could create it, but it would not have a different random seed for each batch. One_like seems more okay if you don't give the user the names back.

Another issue with named constructors is that they make .grad kind of weird. If .grad is always ordered, and you require_grad, you would need a mapping somehow to get the order.

[zdevito]

Right -- autograd and the .grad property are really only defined right now on purely positional tensors since I haven't done a more complete integration. As long as a program starts with positional tensors and creates a positional loss, then autograd can work through any dim-ful operations because they just decompose to existing operations with gradient formulas. Constructors that immediately contain dims have no way to get their gradient at the moment because of this. Doesn't seem like an issue long term if gradient APIs worked directly with dim Tensors.

For randomness, the statefulness of torch's random ops makes it hard to give a behavior that is desirable in all cases. My rationale for having different random numbers across a dim, though, is based on thinking of the attached dims as 'implicit loop nests', in which case a call to rand_like with a batch dim would behave like:

  for batch in range(...):
      result[batch] = torch.rand_like(x[batch])

In other words, batch is out of lexical scope, but it is still part of the dynamic scope of execution.

[srush]

Cool, I agree with most of this.

Yeah random is a hard case, and rand_like is a particularly hard example because it combines scoping and state. It is interesting to think about how this works in Jax. I would say the default assumption is that if you vmap a dimension it would be strange to have effects along that dimension. But I guess you could argue that in Torch something like the second example is implicitly going on.

import jax
import jax.numpy as np
key = jax.random.PRNGKey(1)
def scoped_function(x, key):
    return jax.random.uniform(key, (5,))

print(jax.vmap(scoped_function, in_axes=(0, None))(np.array([1, 2, 3]), key))
print()
print(jax.vmap(scoped_function, in_axes=(0, 0))(np.array([1, 2, 3]), jax.random.split(key, 3)))

[[0.83079386 0.1660409  0.45611215 0.6428884  0.56865394]
 [0.83079386 0.1660409  0.45611215 0.6428884  0.56865394]
 [0.83079386 0.1660409  0.45611215 0.6428884  0.56865394]]

[[0.19773066 0.83156574 0.77199304 0.13508439 0.07068253]
 [0.79785824 0.73702896 0.4746939  0.32407498 0.84053826]
 [0.34117317 0.8346112  0.31955123 0.50551903 0.6800754 ]]