How should I practice Tensor Shapes Through a Neural Network?

Start by reading the full statement, solve the visible examples, then compare the pattern against related NVIDIA phone screen problems in FastPrep.

Is Tensor Shapes Through a Neural Network associated with NVIDIA?

FastPrep tags this as a NVIDIA phone screen coding practice problem when the problem catalog has that company and stage context.

Tensor Shapes Through a Neural Network

NVIDIA● EasyFULLTIMEPHONE SCREEN

Learn

Problem statement

Given the input tensor shape and an ordered sequence of neural-network layers, return the tensor shape at every layer boundary. Compute dimensions only; tensor values and weights are not supplied.

This exercise uses a sequential model. The initial shape initialShape is [N, C, H, W]: batch size, channel count, height and width. The arrays layerTypes and layerParameters describe corresponding layers in execution order.

Layer rules

CONV uses parameters [outChannels, kernelHeight, kernelWidth, strideHeight, strideWidth, paddingHeight, paddingWidth]. It accepts [N,C,H,W] and returns [N,outChannels,outputHeight,outputWidth]. Padding is applied equally on both sides of each axis; dilation is 1.
POOL uses parameters [kernelHeight, kernelWidth, strideHeight, strideWidth]. It accepts and returns a four-dimensional shape, preserving N and C. It uses padding 0 and dilation 1.
For either spatial layer, an axis has output length 1 + floor((inputLength + 2 * padding - kernelLength) / stride). Only complete windows count; incomplete trailing windows are discarded.
FLATTEN uses parameters [] and changes [N,C,H,W] into [N,C * H * W].
LINEAR uses parameters [outFeatures] and changes [N,features] into [N,outFeatures].

All inputs describe a valid model: zero or more spatial layers, optionally followed by one FLATTEN and then zero or more LINEAR layers. A model may end at any of these boundaries. No invalid-shape handling is required.

If there are L layers, return L + 1 rows. Row 0 is initialShape, and row i + 1 is the output of layer i. Thus row i and row i + 1 give that layer's input and output shapes. For an empty model, return only the initial shape. Rows have four entries before flattening and two entries afterward.

What the interview report shared

The report describes calculating each layer's input and output tensor shape in a deep learning model. It does not specify the model, layers or shape conventions.

Function

tensorShapeTrace(initialShape: int[], layerTypes: String[], layerParameters: int[][]) → int[][]

Examples

Example 1

initialShape = [2,3,8,9]layerTypes = ["CONV","POOL","FLATTEN","LINEAR"]layerParameters = [[4,3,3,2,2,1,1],[2,2,2,2],[],[5]]return = [[2,3,8,9],[2,4,4,5],[2,4,2,2],[2,16],[2,5]]

The convolution produces 4 channels and spatial dimensions 4 by 5. Pooling keeps the batch and channels, producing 2 by 2. Flattening gives 4 * 2 * 2 = 16 features per sample, and the linear layer produces 5 features.

Example 2

initialShape = [1,2,10,13]layerTypes = ["CONV","CONV"]layerParameters = [[5,3,4,3,4,0,0],[2,2,1,2,2,1,0]]return = [[1,2,10,13],[1,5,3,3],[1,2,2,2]]

The first layer uses complete windows, so the height and width are both 3. The second layer receives that output, uses height padding 1 and width padding 0, and produces shape [1,2,2,2].

Example 3

initialShape = [3,2,2,4]layerTypes = ["FLATTEN","LINEAR","LINEAR"]layerParameters = [[],[7],[2]]return = [[3,2,2,4],[3,16],[3,7],[3,2]]

Flattening combines only the channel and spatial axes: 2 * 2 * 4 = 16. Each linear layer replaces the feature count while preserving batch size 3. The output includes the initial shape and all three layer outputs.

Constraints

initialShape.length = 4, and 0 <= layerTypes.length = layerParameters.length <= 80.
1 <= N, C, outChannels <= 64.
Every spatial input and output dimension is between 1 and 128.
Every kernel size and stride is between 1 and 128; convolution padding is between 0 and 64.
Each kernel fits within the current padded input axis, so both spatial outputs are positive.
1 <= outFeatures <= 1048576. Flattening produces at most 64 * 128 * 128 = 1048576 features.
Layer names, parameter counts and layer ordering follow the stated rules.