Katkat
Back

Debug Tensor Conversion Bugs

WaymoMachine Learning Engineer · Technical Screen
MediumMachine Learning EngineerTechnical ScreenCoding & Algorithms

Last updated

This question evaluates debugging and implementation-level skills in tensor and array manipulation, focusing on aliasing behavior, conversion between column lists and ndarrays, correct handling of non-divisible chunking, and writing tests for these issues.

You are given a small tensor utility library used in distributed numerical code. The implementation has several bugs related to aliasing, array conversion, and dropping remainder elements. Identify the bugs, fix the implementation, and add tests.

import numpy as np

class Matrix:
 def __init__(self, data):
 self.data = data

 @staticmethod
 def zeros(rows, cols):
 # Intended: create a rows by cols matrix where each cell is independent.
 return Matrix([[0] * cols] * rows)

def to_ndarray(columns):
 # columns is a list of one-dimensional arrays of equal length.
 # Intended output shape: (n, k), where k is the number of columns.
 return np.stack(columns)

def from_ndarray(arr, chunk_size):
 # Intended: split arr along axis 0 into chunks of size chunk_size.
 # The final chunk should contain the remaining rows if arr.shape[0]
 # is not divisible by chunk_size.
 chunks = []
 for i in range(arr.shape[0] // chunk_size):
 chunks.append(arr[i * chunk_size:(i + 1) * chunk_size])
 return chunks

Requirements:

  1. Mutating one row of Matrix.zeros(rows, cols) must not mutate other rows.
  2. to_ndarray([a, b, c]) should return an array whose columns are a, b, and c.
  3. from_ndarray must not silently drop remainder rows.
  4. Add tests for aliasing, shape correctness, and non-divisible chunk sizes.