Formas e layout

O protótipo de Shape do XLA (xla_data.proto) descreve a classificação, o tamanho e o tipo de dados de uma matriz de N-dimensional (abreviada como matriz).

Terminologia, notação e convenções

  • A classificação de uma matriz é igual ao número de dimensões. A classificação verdadeira de uma matriz é o número de dimensões com tamanho maior que 1.

  • As dimensões são numeradas de 0 a N-1 em uma matriz dimensional N. Os números de dimensão são rótulos arbitrários por conveniência. A ordem desses números de dimensão não implica uma ordem secundária/principal específica no layout da forma. O layout é determinado pelo protótipo Layout.

  • Por convenção, as dimensões são listadas em ordem crescente de número de dimensão. Por exemplo, para uma matriz tridimensional de tamanho [A x B x C], a dimensão 0 tem o tamanho A, a 1 tem o tamanho B e a 2 tem o tamanho C.

    Alguns utilitários no XLA também oferecem suporte à indexação negativa como Python: a dimensão 1 é a última dimensão (equivalente a N-1 para uma matriz dimensional N). Por exemplo, para a matriz tridimensional descrita acima, a dimensão -1 tem o tamanho C, a dimensão -2 tem o tamanho B e assim por diante.

  • Matrizes de duas, três e quatro dimensões geralmente têm letras específicas associadas a dimensões. Por exemplo, para uma matriz 2D:

    • dimensão 0: y
    • dimensão 1: x

    No caso de uma matriz 3D:

    • dimensão 0: z
    • dimensão 1: y
    • dimensão 2: x

    Para uma matriz 4D:

    • dimensão 0: p
    • dimensão 1: z
    • dimensão 2: y
    • dimensão 3: x
  • As funções na API do XLA que usam dimensões fazem isso em ordem crescente de número de dimensão. Isso corresponde à ordem usada ao transmitir dimensões como um initializer_list. Por exemplo:

    ShapeUtil::MakeShape(F32, {A, B, C, D})

    criará um polígono cuja matriz de tamanho de dimensão consiste na sequência [A, B, C, D].

Layout

O protótipo Layout descreve como uma matriz é representada na memória. O protótipo Layout inclui os campos abaixo:

message Layout {
  repeated int64 minor_to_major = 1;
  repeated int64 padded_dimensions = 2;
  optional PaddingValue padding_value = 3;
}

Ordenação das dimensões menor a maior

O único campo obrigatório é minor_to_major. Esse campo descreve a ordem da menor para a maior das dimensões em uma forma. Os valores em minor_to_major são uma ordenação das dimensões da matriz (de 0 a N-1 para uma matriz dimensional N), em que o primeiro valor é a dimensão mais secundária até o último, que é a dimensão mais principal. A dimensão mais secundária é aquela que muda mais rapidamente ao percorrer os elementos da matriz disposto na memória linear.

Por exemplo, considere a seguinte matriz 2D de tamanho [2 x 3]:

a b c
d e f

Aqui, a dimensão 0 é de tamanho 2, e a dimensão 1 é de 3. Se o campo minor_to_major no layout for [0, 1], a dimensão 0 será a menor e a 1 será a mais principal. Isso corresponde ao seguinte layout na memória linear:

a d b e c f

Essa ordem de dimensão menor a maior de 0 até N-1 é semelhante à coluna principal (na classificação 2). Presumindo uma ordem monotônica de dimensões, outra maneira de fazer referência a esse layout no código é simplesmente "dim 0 é menor".

Por outro lado, se o campo minor_to_major for [1, 0], o layout na memória linear será:

a b c d e f

Uma ordem de dimensão menor a maior de N-1 até 0 para uma matriz dimensional N é semelhante à linha principal (na classificação 2). Presumindo uma ordem monotônica de dimensões, outra maneira de nos referirmos a esse layout no código é simplesmente "dim 0 é principal".

Ordem padrão menor para maior

O layout padrão para Formas recém-criadas é "a ordem de dimensão é maior para menor" (semelhante à linha maior na classificação 2).

Padding

O padding é definido nos campos opcionais padded_dimensions e padding_value. O campo padded_dimensions descreve os tamanhos (larguras) aos quais cada dimensão é preenchida. Se presente, o número de elementos em padded_dimensions precisa ser igual à classificação da forma.

Por exemplo, considerando a matriz [2 x 3] definida acima, se padded_dimensions for [3, 5], a dimensão 0 será preenchida com uma largura de 3, e a dimensão 1 será preenchida com uma largura de 5. O layout na memória linear (supondo um valor de preenchimento de 0 e um layout principal da coluna) é:

a d 0 b e 0 c f 0 0 0 0 0 0 0

Isso equivale ao layout da matriz a seguir com a mesma ordem de dimensões de menor para maior:

a b c 0 0
d e f 0 0
0 0 0 0 0

Como indexar em matrizes

A classe IndexUtil em index_util.h oferece utilitários para conversão entre índices multidimensionais e índices lineares de acordo com uma forma e um layout. Índices multidimensionais incluem um índice int64 para cada dimensão. Índices lineares são um único valor int64 que é indexado no buffer que contém a matriz. Consulte shape_util.h e layout_util.h no mesmo diretório para ver utilitários que simplificam a criação e a manipulação de formas e layouts.