Skip to content

Instantly share code, notes, and snippets.

@wanchaol
Last active December 20, 2018 23:35
Show Gist options
  • Save wanchaol/3d2545fcf72fb01750a00cdbb080e59e to your computer and use it in GitHub Desktop.
Save wanchaol/3d2545fcf72fb01750a00cdbb080e59e to your computer and use it in GitHub Desktop.

ATen/JIT Function Schema

This doc is to talk about the difference in ATen/JIT function schema and find a way to align each other.

List syntax

Currently in ATen function schema when we want to define a list, we have:

  • TensorList (ArrayRef in C++)
  • IntList (ArrayRef in C++)

But this currently cannot support more list syntaxes, like list of optional tensors, list of list of ints, even list of doubles. So first we would like to expand the list syntax in ATen schema, to make list syntax more general. For example we can use [] to denote a list, the above two(TensorList and IntList) can become:

  • Tensor[]
  • int[]

Using [] as a list decorator is one possibility that we can enable other types inside the list container (we can also think of other syntax as well), so the list of optional tensors, or the list of list of ints, and list of doubles can become:

  • Tensor?[]
  • int[][]
  • double[]

The [] is currently only used in JIT schema, we would need to think of a way to align with ATen schema, we can also not use [] in the schema, but do something like type annotation, then the above three become:

  • List[Optional[Tensor]]
  • List[List[int]]
  • List[float]

I think either [] or type annotation will give us the ability to generalize the list syntax declaration, I am not sure if we want to get rid of TensorList or IntList because they are still valid C++ types in ATen, but we should at least expand the list syntax in the ATen function schema.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment