Tensoring the "technical" interview

Ah yes, Leetcode Easy, where even the brightest of minds have struggled. But you should be able to overcome this, sleep deprived or not, if only with the help of a certain energy drink (which you frequently wish would sponsor you).

You start (as any good piece of code should) with a dash of boilerplate:

import jax.numpy as jnp
from jax.nn import one_hot
from jax import jit, lax

The problem, inspired as it is, compels you to bind the laws of time to the medium of mathematics:

A = jnp.zeros((7, 7, 10), dtype="int32")
A = A.at[:, 1, 0].set(one_hot(2, 7))
A = A.at[:, 1, 1].set(one_hot(2, 7))
A = A.at[:, 1, 2].set(one_hot(3, 7))
for i in range(10):
    A = A.at[:, 2, i].set(one_hot(4, 7))
    A = A.at[:, 5, i].set(one_hot(6, 7))
for i in range(6):
    A = A.at[:, 4, i].set(one_hot(5, 7))
for i in range(4):
    A = A.at[:, 3, i].set(one_hot(4, 7))

The language of time, cryptic to some, but regular (finite too, though it matters little right now).

Conjure an entry point, capturing the essense of the input string, but leave the main code still unwritten:

def solve(inp):
    chrs = jnp.array([ord(inp[i]) for i in (0, 1, 3, 4)], dtype="int32")
    return impl(chrs)

And finish it, only now having to use a single nontrivial control flow construct.

@jit
def impl(chrs):
    vs = lax.map(
        lambda sym: lax.select(
            sym != ord("?"),
            one_hot(sym - ord("0"), 10, dtype="int32"),
            jnp.ones(10, dtype="int32"),
        ),
        chrs,
    )
    return lax.fori_loop(
        0, vs.shape[0], lambda i, w: w @ A @ vs[i], one_hot(1, 7, dtype="int32")
    )[6]

Think of a million parallel executions of the kernel, sharded between thousands of GPUs, all doing precisely what they excel at -- linear algebra.

Surely this is the solution they had in mind?

tail modulo кто?

Одна из самых важных функций на контейнере это инстанс функтора на нем map. На типичных структурах нам хочется, чтобы она работала быстро, а для простых структур было бы странным, если её реализация будет сложной.

Возьмем List:

let rec map f lst =
  match lst with
  | [] -> []
  | x :: xs -> f x :: map f xs

Это наиболее естественный способ написания ~~~~, и в коммите от 4 мая 1995 года (уф) со звучным сообщением Passage a la version bootstrappee (franchissement du Rubicon) эта функция была реализована в OCaml base практически так же, и не менялась следующие 27 лет и 2 месяца (спойлеры). Но он не подходит для стандартной библиотеки промышленного языка 🙂, хаскелистам увы ¹.

Посмотрим, что получалось у людей, которым платят деньги за код, уже лет 15 назад:

(* stdlib/list.ml *)
let rev_map f l =
  let rec rmap_f accu = function
    | [] -> accu
    | a::l -> rmap_f (f a :: accu) l
  in
  rmap_f [] l

let rec rev_append l1 l2 =
  match l1 with
    [] -> l2
  | a :: l -> rev_append l (a :: l2)

let rev l = rev_append l []

(* js Base/src/list.ml *)
let map_slow l ~f = rev (rev_map ~f l)

let rec count_map ~f l ctr =
  match l with
  | [] -> []
  | [x1] ->
    let f1 = f x1 in
    [f1]
  ...
  (* ручной разбор случаев *)
  ...
  | x1 :: x2 :: x3 :: x4 :: x5 :: tl ->
    let f1 = f x1 in
    let f2 = f x2 in
    let f3 = f x3 in
    let f4 = f x4 in
    let f5 = f x5 in
    f1 :: f2 :: f3 :: f4 :: f5 ::
    (if ctr > 1000
     then map_slow ~f tl
     else count_map ~f tl (ctr + 1))

let map l ~f = count_map ~f l 0

👍️️️️️️👍️️️️️️👍️️️️️️

Всё из-за того, что изначальный map использовал \( \mathcal{O}(n) \) стека, так как рекурсивный вызов находился не в хвостовой позиции.

В текущей же версии мы увидим² значительно более простое

(* binder для списка тут перенесен направо и анонимен,
   потому что мы мгновенно деконструируем его *)
let[@tail_mod_cons] rec map f = function
    [] -> []
  | x::l -> let y = f x in y::map f l

Класс функций, для которых (единственный) рекурсивный вызов находится внутри конструктора, который возвращается из нее, называется tail-recursive modulo constructor, и его тоже можно скомпилировать так, чтобы использовать \(\mathcal{O}(1)\) стека:

(* фейковый синтаксис из документации, самому написать такое нельзя *)
let rec map f l =
  match l with
  | [] -> []
  | x :: xs ->
    let y = f x in
    let dst = y ::{mutable} Hole in
    map_dps f xs dst 1;
    dst

and map_dps f l dst idx =
  match l with
  | [] -> dst.idx <- []
  | x :: xs ->
    let y = f x in
    let dst' = y ::{mutable} Hole in
    dst.idx <- dst';
    map_dps f xs dst' 1

Достаточно сконструировать результат, заполнив поле хоть каким-нибудь значением, и дальше исправить его на корректное значение в следующем хвостовом вызове. По-моему, это забавно.

Но еще забавнее следующее³, документация приводит

let[@tail_mod_cons] rec length l =
  match l with
  | [] -> 0
  | _ :: xs -> 1 + length xs

как пример функции, которая в которой эта техника не работает -- (* does *not* work: addition is not a data constructor *). Но фанаты Lean, например, знают, что 1 + x для x : Nat это каноническая форма для Nat.suc x, а это буквально конструктор Nat в индуктивном смысле:

inductive Nat where
  | zero : Nat
  | succ (n : Nat) : Nat

Известно, конечно, что Lean в этом случае нагло лжет, внутри себя он не представляет натуральные числа как связные списки, однако утверждения выше все еще верны. Вот это уже действительно заставляет задуматься.

думайте

алгебраическая бебрология подпишись

1. хотя, справедливости ради, то, что мы будем обсуждать, там некорректно из-за ленивости ↩

2. на самом деле случай одного элемента разобран там отдельно, но это не принципиально ↩

3. и это то, почему я вообще написал этот пост ↩