FastPrepImplement Explicit Integer Einsum

Implement Explicit Integer Einsum

NVIDIA logoNVIDIA● HardFULLTIMEPHONE SCREEN
Learn

Problem statement

For this exercise, implement a bounded explicit integer einsum evaluator without calling a tensor library. An equation has the form term1,term2,...->output. Each input term lists one lowercase label per axis of its operand. The output labels are distinct and may be empty.

shapes[r] gives operand r's dimensions. values[r] contains that operand's entries in row-major order: the last axis varies fastest. Every occurrence of one label has the same dimension, including repeated labels inside a single term.

For each assignment of coordinates to the distinct labels, read one entry from each operand using its term's labels, multiply those entries, and add the product to the output cell selected by the output labels. Thus, a label absent from the output is summed; a repeated input label selects a diagonal; and retained labels align across operands.

Output axes follow the right-hand label order exactly. Return the row-major flattened output as64-bit integers. An empty output term represents a scalar and returns a one-element array. For example, ii->i keeps a matrix diagonal, whereas ii-> sums that diagonal.

All inputs are valid for this explicit subset. Implicit output, ellipsis, broadcasting, scalar input operands and malformed equations are outside the accepted domain.

Function

einsum(equation: String, shapes: int[][], values: int[][]) → long[]

Examples

Example 1

equation = "ij,jk->ik"shapes = [[2,3],[3,2]]values = [[1,2,3,4,5,6],[7,8,9,10,11,12]]return = [58,64,139,154]

Contract j: the first output cell is1*7+2*9+3*11=58. The output shape is[2,2].

Example 2

equation = "ii->"shapes = [[3,3]]values = [[1,2,3,4,5,6,7,8,9]]return = [15]

Repeated i selects the diagonal1,5,9. An empty output-label list reduces to one scalar, returned as[15].

Example 3

equation = "ab->ba"shapes = [[2,3]]values = [[1,2,3,4,5,6]]return = [1,4,2,5,3,6]

Output axes are b then a, so row-major output is the matrix transpose.

Constraints

  • 1 <= shapes.length = values.length <= 3, matching the number of input terms.
  • Each input term contains 1 through 4 labels, matching its shape rank; labels are lowercase ASCII letters, and equations contain no spaces.
  • There are at most 6 distinct input labels. Each dimension is from 1 through 8, and all occurrences of a label have equal dimensions.
  • Each output label occurs in an input term and appears at most once in the output.
  • Each flattened input has the product of its shape dimensions and at most 5000 entries; the output has at most 5000 entries.
  • The product of dimensions of all distinct labels is at most 100000.
  • -100 <= values[r][i] <= 100. Products and sums must use signed64-bit integers.

More NVIDIA problems

See NVIDIA hiring insights
public long[] einsum(String equation, int[][] shapes, int[][] values) {
    return new long[0];
}
equation"ij,jk->ik"
shapes[[2,3],[3,2]]
values[[1,2,3,4,5,6],[7,8,9,10,11,12]]
expected[58,64,139,154]
Checking account…