6.13 Folds, Reductions and Expansions

6.13 Folds, Reductions and Expansions🔗ℹ

6.13.1 Axis Folds🔗ℹ

procedure
(array-axis-fold arr k f) → (Array A)
  arr : (Array A)
  k : Integer
  f : (A A -> A)
(array-axis-fold arr k f init) → (Array B)
  arr : (Array A)
  k : Integer
  f : (A B -> B)
  init : B

Folds a binary function f over axis k of arr. The result has the shape of arr but with axis k removed.

The three-argument variant uses the first value of each row in axis k as init. It therefore requires axis k to have positive length.

Examples:

> (define arr (index-array #(3 4)))
> arr
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Index))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[#[0 1 2 3] #[4 5 6 7] #[8 9 10 11]])
> (array-axis-fold arr 0 +)
- : #(struct:Array
      (Indexes
       Index
       (Boxof Boolean)
       (-> Void)
       (-> Indexes Nonnegative-Integer))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[12 15 18 21])
> (array-axis-fold arr 1 (inst cons Index (Listof Index)) empty)
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes (Listof Index)))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #['(3 2 1 0) '(7 6 5 4) '(11 10 9 8)])

Notice that the second example returns an array of reversed lists. This is therefore a left fold; see foldl.

If you need control over the order of evaluation in axis k’s rows, see array-axis-reduce.

syntax
(array-axis-sum arr k)
(array-axis-sum arr k init)
syntax
(array-axis-prod arr k)
(array-axis-prod arr k init)

   arr : (Array Number)
   k : Integer
   init : Number

syntax
(array-axis-min arr k)
(array-axis-min arr k init)
syntax
(array-axis-max arr k)
(array-axis-max arr k init)

   arr : (Array Real)
   k : Integer
   init : Real

Some standard per-axis folds, defined in terms of array-axis-fold. The two-argument variants require axis k to have positive length.

Examples:

> (define arr (index-array #(3 4)))
> arr
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Index))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[#[0 1 2 3] #[4 5 6 7] #[8 9 10 11]])
> (array-axis-fold arr 1 +)
- : #(struct:Array
      (Indexes
       Index
       (Boxof Boolean)
       (-> Void)
       (-> Indexes Nonnegative-Integer))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[6 22 38])
> (array-axis-sum arr 1)
- : #(struct:Array
      (Indexes
       Index
       (Boxof Boolean)
       (-> Void)
       (-> Indexes Nonnegative-Integer))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[6 22 38])
> (array-axis-sum arr 0 0.0)
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Nonnegative-Flonum))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[12.0 15.0 18.0 21.0])

procedure
(array-axis-count arr k pred?) → (Array Index)
  arr : (Array A)
  k : Integer
  pred? : (A -> Any)

Counts the elements x in rows of axis k for which (pred? x) is true.

Examples:

> (define arr (index-array #(3 3)))
> arr
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Index))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[#[0 1 2] #[3 4 5] #[6 7 8]])
> (array-axis-count arr 1 odd?)
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Index))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[1 2 1])

procedure
(array-axis-and arr k) → (Array (U A Boolean))
  arr : (Array A)
  k : Integer
procedure
(array-axis-or arr k) → (Array (U A #f))
  arr : (Array A)
  k : Integer

Apply and or or to each row in axis k of array arr. Evaluation is short-cut as with the and and or macros, which is only observable if arr is nonstrict.

In the following example, computing the second array element sets second? to #t:

> (define second? (ann #f Boolean))
> (define arr
    (parameterize ([array-strictness #f])
      (build-array #(2) (λ: ([js : Indexes])
                          (cond [(zero? (vector-ref js 0))  #f]
                                [else  (set! second? #t)
                                       #t])))))

Printing arr causes (set! second? #t) to be evaluated, but applying array-axis-and does not:

> (array-axis-and arr 0)
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Boolean))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #f)
> second?
- : Boolean
#f

However, if arr were strict, (set! second? #t) would be evaluated when arr was created.

6.13.2 Whole-Array Folds🔗ℹ

procedure
(array-fold arr g) → (Array A)
arr : (Array A)
g : ((Array A) Index -> (Array A))

Folds g over each axis of arr, in reverse order. The arguments to g are an array (initially arr) and the current axis. It should return an array with one fewer dimension than the array given, but does not have to.

Examples:

> (define arr (index-array #(3 4)))
> arr
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Index))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[#[0 1 2 3] #[4 5 6 7] #[8 9 10 11]])
> (array-fold arr (λ: ([arr : (Array Integer)] [k : Index])
                    (array-axis-sum arr k)))
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Integer))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array 66)
> (apply + (array->list arr))
- : Integer [more precisely: Nonnegative-Integer]
66
> (array-ref (array-fold arr (inst array->list-array (Listof* Integer)))
             #())
- : (U (Listof (Rec g8423864 (U (Listof g8423864) Integer))) Index)
'((0 1 2 3) (4 5 6 7) (8 9 10 11))

procedure
(array-all-fold arr f) → A
  arr : (Array A)
  f : (A A -> A)
(array-all-fold arr f init) → A
  arr : (Array A)
  f : (A A -> A)
  init : A

Folds f over every element of arr by folding f over each axis in reverse order. The two-argument variant is equivalent to

(array-ref (array-fold arr (λ: ([arr : (Array A)] [k : Index])
(array-axis-fold arr k f)))
#())

and the three-argument variant is similar. The two-argument variant requires every axis to have positive length.

Examples:

> (define arr (index-array #(3 4)))
> arr
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Index))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[#[0 1 2 3] #[4 5 6 7] #[8 9 10 11]])
> (array-all-fold arr +)
- : Integer [more precisely: Nonnegative-Integer]
66
> (array-all-fold (array #[]) + 0.0)
- : Flonum [more precisely: Nonnegative-Flonum]
0.0

Because f is folded over the last axis first, it receives arr’s elements (as its first argument) in row-major order.

syntax
(array-all-sum arr)
(array-all-sum arr init)
syntax
(array-all-prod arr)
(array-all-prod arr init)

arr : (Array Number)
init : Number

syntax
(array-all-min arr)
(array-all-min arr init)
syntax
(array-all-max arr)
(array-all-max arr init)

arr : (Array Real)
init : Real

Some standard whole-array folds, defined in terms of array-all-fold. The one-argument variants require each axis in arr to have positive length.

Examples:

> (define arr (index-array #(3 4)))
> arr
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Index))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[#[0 1 2 3] #[4 5 6 7] #[8 9 10 11]])
> (array-all-fold arr +)
- : Integer [more precisely: Nonnegative-Integer]
66
> (array-all-sum arr)
- : Integer [more precisely: Nonnegative-Integer]
66
> (array-all-sum arr 0.0)
- : Real [more precisely: Nonnegative-Real]
66.0

procedure
(array-all-and arr) → (U A Boolean)
arr : (Array A)
procedure
(array-all-or arr) → (U A #f)
arr : (Array A)

Apply and or or to arr’s elements using short-cut evaluation in row-major order.

Examples:

> (define arr (index-array #(3 3)))
> (array-all-and (array= arr arr))
- : Boolean
#t
> (define brr (array+ arr (array 1)))
> (array-all-and (array= arr brr))
- : Boolean
#f
> (array-all-or (array= arr (array 0)))
- : Boolean
#t

(array-all-and arr) is defined as

(parameterize ([array-strictness #f])
(array-ref (array-fold arr array-axis-and) #()))

and array-all-or is defined similarly, using array-axis-or.

syntax
(array-count pred? arrs ...)

arrs : (Array Ts)
pred? : (Ts ... -> Any)

When given one array arr, returns the number of elements x in arr for which (pred? x) is true. When given multiple arrays, array-count does the same with the corresponding elements from any number of arrays. If the arrays’ shapes are not the same, they are broadcast first.

Examples:

> (array-count zero? (array #[#[0 1 0 2] #[0 3 -1 4]]))
- : Integer [more precisely: Index]
3
> (array-count equal?
(array #[#[0 1] #[2 3] #[0 1] #[2 3]])
(array #[0 1]))
- : Integer [more precisely: Index]
4

(array-count pred? arrs ...) is like

(array-all-sum (array-map (λ (x ...) (if (pred? x ...) 1 0)) arrs ...)
0)

but does not create intermediate (strict) arrays, and always returns an Index.

syntax
(array-andmap pred? arrs ...)
syntax
(array-ormap pred? arrs ...)

arrs : (Array Ts)
pred? : (Ts ... -> Any)

Like andmap and ormap, but for arrays. Evaluation is short-cut, in row-major order. If the arrays’ shapes are not the same, they are broadcast first.

Determining whether each row is equal to (array #[0 1]):

> (array-andmap equal?
(array #[#[0 1] #[0 1] #[0 1] #[0 1]])
(array #[0 1]))
- : Boolean
#t

Determining whether any row has 0 as its first element or 1 as its second:

> (array-ormap equal?
(array #[#[0 2] #[2 3] #[1 1] #[2 3]])
(array #[0 1]))
- : Boolean
#t

Determining whether any row is equal to (array #[0 1]):

> (array-ormap equal?
(array->list-array (array #[#[0 2] #[2 3] #[1 1] #[2 3]]))
(array->list-array (array #[0 1])))
- : Boolean
#f

(array-andmap pred? arrs ...) is defined as

(parameterize ([array-strictness #f])
(array-all-and (array-map pred? arrs ...)))

and array-ormap is defined similarly, using array-all-or.

6.13.3 General Reductions and Expansions🔗ℹ

procedure
(array-axis-reduce arr k h) → (Array B)
  arr : (Array A)
  k : Integer
  h : (Index (Integer -> A) -> B)

Like array-axis-fold, but allows evaluation control (such as short-cutting and and or) by moving the loop into h. The result has the shape of arr, but with axis k removed.

The arguments to h are the length of axis k and a procedure that retrieves elements from that axis’s rows by their indexes in axis k. It should return the elements of the resulting array.

For example, summing the squares of the rows in axis 1:

> (define arr (index-array #(3 3)))
> arr
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Index))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[#[0 1 2] #[3 4 5] #[6 7 8]])
> (array-axis-reduce
   arr 1
   (λ: ([dk : Index] [proc : (Integer -> Real)])
     (for/fold: ([s : Real 0]) ([jk  (in-range dk)])
       (+ s (sqr (proc jk))))))
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Real))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[5 50 149])
> (array-axis-sum (array-map sqr arr) 1)
- : #(struct:Array
      (Indexes
       Index
       (Boxof Boolean)
       (-> Void)
       (-> Indexes Nonnegative-Integer))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[5 50 149])

Transforming an axis into a list using array-axis-fold and array-axis-reduce:

> (array-map (inst reverse Index)
             (array-axis-fold arr 1
                              (inst cons Index (Listof Index))
                              empty))
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes (Listof Index)))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #['(0 1 2) '(3 4 5) '(6 7 8)])
> (array-axis-reduce arr 1 (inst build-list Index))
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes (Listof Index)))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #['(0 1 2) '(3 4 5) '(6 7 8)])

The latter is essentially the definition of array->list-array.

Every fold, including array-axis-fold, is ultimately defined using array-axis-reduce or its unsafe counterpart.

procedure
(array-axis-expand arr k dk g) → (Array B)
  arr : (Array A)
  k : Integer
  dk : Integer
  g : (A Index -> B)

Inserts a new axis number k of length dk, using g to generate values; k must be no greater than the dimension of arr, and dk must be nonnegative.

Conceptually, g is applied dk times to each element in each row of axis k, once for each nonnegative index jk < dk.

Turning vector elements into rows of a new last axis using array-axis-expand and vector-ref:

> (define arr (array #['#(a b c) '#(d e f) '#(g h i)]))
> (array-axis-expand arr 1 3 vector-ref)
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Any))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[#['a 'b 'c] #['d 'e 'f] #['g 'h 'i]])

Creating a vandermonde-matrix:

> (array-axis-expand (list->array '(1 2 3 4)) 1 5 expt)
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Positive-Integer))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[#[1 1 1 1 1] #[1 2 4 8 16] #[1 3 9 27 81] #[1 4 16 64 256]])

This function is a dual of array-axis-reduce in that it can be used to invert applications of array-axis-reduce. To do so, g should be a destructuring function that is dual to the constructor passed to array-axis-reduce. Example dual pairs are vector-ref and build-vector, and list-ref and build-list.

(Do not pass list-ref to array-axis-expand if you care about performance, though. See list-array->array for a more efficient solution.)

procedure
(array->list-array arr [k]) → (Array (Listof A))
arr : (Array A)
k : Integer = 0

Returns an array of lists, computed as if by applying list to the elements in each row of axis k.

Examples:

> (define arr (index-array #(3 3)))
> arr
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Index))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[#[0 1 2] #[3 4 5] #[6 7 8]])
> (array->list-array arr 1)
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes (Listof Index)))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #['(0 1 2) '(3 4 5) '(6 7 8)])
> (array-ref (array->list-array (array->list-array arr 1) 0) #())
- : (Listof (Listof Index))
'((0 1 2) (3 4 5) (6 7 8))

See mean for more useful examples, and array-axis-reduce for an example that shows how array->list-array is implemented.

procedure
(list-array->array arr [k]) → (Array A)
arr : (Array (Listof A))
k : Integer = 0

Returns an array in which the list elements of arr comprise a new axis k. Equivalent to (array-axis-expand arr k n list-ref) where n is the length of the lists in arr, but with O(1) indexing.

Examples:

> (define arr (array->list-array (index-array #(3 3)) 1))
> arr
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes (Listof Index)))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #['(0 1 2) '(3 4 5) '(6 7 8)])
> (list-array->array arr 1)
- : #(struct:Array
      (Indexes Index (Boxof Boolean) (-> Void) (-> Indexes Index))
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:56:13 prop:equal+hash>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:55:13 prop:custom-write>
      #<syntax:build/user/8.15/pkgs/math-lib/math/private/array/typed-array-struct.rkt:54:13 prop:custom-print-quotable>)
(array #[#[0 1 2] #[3 4 5] #[6 7 8]])

For fixed k, this function and array->list-array are mutual inverses with respect to their array arguments.

1	Constants and Elementary Functions
2	Flonums
3	Special Functions
4	Number Theory
5	Arbitrary-Precision Floating-Point Numbers (Bigfloats)
6	Arrays
7	Matrices and Linear Algebra
8	Statistics Functions
9	Probability Distributions
10	Stuff That Doesn’t Belong Anywhere Else

6.1	Quick Start
6.2	Definitions
6.3	Broadcasting
6.4	Slicing
6.5	Nonstrict Arrays
6.6	Types, Predicates and Accessors
6.7	Construction
6.8	Conversion
6.9	Comprehensions and Sequences
6.10	Pointwise Operations
6.11	Indexing and Slicing
6.12	Transformations
6.13	Folds, Reductions and Expansions
6.14	Other Array Operations
6.15	Subtypes
6.16	Strictness

6.13.1	Axis Folds
6.13.2	Whole-Array Folds
6.13.3	General Reductions and Expansions