Matrix operators as stack languages 2: closing maps @ Código para llevar [blog]

On the last post we saw how some matrix operations can be considered as forth-like operations over a stack of dimensions.

To do a swap add a transpose(mat) and continue on the left of it.
To move from dimension stack to the auxiliary one, open a map( (don't close it) and continue on inside the map.
New operations can be performed inside the map as long as they are done over the map lambda parameter.

Now, let's try to extend the process to allow exiting a MAP (moving from the auxiliary stack back into the dimensions one).

For starters let's update the code from last post so we can see how the operations are converted to code.

import re
import sys

def print_status(text, cursor, target):
    print(' '.join(text[:cursor]) + ' *CURSOR* ' + ' '.join(text[cursor:]), f"| target={target}")

def compose(commands, target):
    text = []
    cursor = 0

    for command in commands:
        if command.upper() == 'SWAP':
            target = f"transpose({target})"
        elif command.upper() == 'MAP':
            text.insert(cursor, r'map (\next -> ')
            text.insert(cursor + 1, f') ({target})')
            cursor += 1
            target = "next"
        else:
            if not command.startswith(':'):
                raise Exception(f"Found “{command}”, expected “SWAP”, “MAP” or something starting with ‘:’")
            function = command[1:]
            target = f'map ({function}) ({target})'

        print(f"Command: {command}")
        print_status(text, cursor, target)
        print()

    text.insert(cursor, target)
    return ' '.join(text)

def convert(stack_code, target_matrix):
    lines = stack_code.split('\n')
    no_comments = []
    for line in lines:
        if '--' in line:
            line = line[:line.index("--")]
        no_comments.append(line.strip())

    commands = re.split(r'\s+', " ".join(no_comments).strip())
    return compose(commands, target_matrix)

if __name__ == '__main__':
    if len(sys.argv) < 3:
        print(f"Usage: {sys.argv[0]} <base matrix> <operations>")
        exit(0)

    print("Result:", convert(' '.join(sys.argv[2:]), sys.argv[1]))

If we launch it we can see how the code is built and how the cursor moves:

$ python converter.py 'mat' 'SWAP MAP SWAP MAP SWAP'
Command: SWAP
 *CURSOR*  | target=transpose(mat)

Command: MAP
map (\next ->  *CURSOR* ) (transpose(mat)) | target=next

Command: SWAP
map (\next ->  *CURSOR* ) (transpose(mat)) | target=transpose(next)

Command: MAP
map (\next ->  map (\next ->  *CURSOR* ) (transpose(next)) ) (transpose(mat)) | target=next

Command: SWAP
map (\next ->  map (\next ->  *CURSOR* ) (transpose(next)) ) (transpose(mat)) | target=transpose(next)

Result: map (\next ->  map (\next ->  transpose(next) ) (transpose(next)) ) (transpose(mat))

Now, we see that as of now we can't return from a map without completing the whole procedure, as the cursor can only go deeper. For the return to happen we have to use the value of map we're returning from as the parameter of the following functions.

A minimal example of a program that returned from a map would be something that applied multiple reduce-like operations. For example to sum all elements on different dimensions.

While we're at it, let's remove the map when applying custom functions (see line 23 above).

:i32.sum      \ generates: i32.sum (mat1d)
MAP :i32.sum  \ generates: map (i32.sum) (mat2d)

\ Thus, adding a RETurn operation we could do
MAP :i32.sum RET :i32.sum  \ generates: i32.sum (map (i32.sum) (mat2d))

\ For 3 dimensions
MAP MAP :i32.sum RET :i32.sum RET :i32.sum
\ generates:  i32.sum (map (\next -> i32.sum (map (\next -> (i32.sum) (next))
\                                                 (next)))
\                          (mat3d))

\ And finally, for 4 dimensions
MAP MAP MAP :i32.sum RET :i32.sum RET :i32.sum RET :i32.sum
\ generates:  i32.sum (map (\next -> (i32.sum (map (\next -> i32.sum (map (\next -> (i32.sum) (next))
\                                                                (next)))
\                                                  (next))))
\                          (mat4d))

Note: You might notice in the following script that each innermost map (i32.sum) (mat) will get converted into a (map (\next -> ((i32.sum) (next)) ) (mat)). This is equivalent and allows for more regularity on the converter script.

To do this we might add two elements.

A cursor stack, which points to the point where the cursor is in before each map. This can be done with a simple integer list as the cursor only increases, and nothing is ever added before it. This removes the need for recalculating the map cursors.
Delimiters for map scopes. So we can extract a whole map operation (and it's parameters) from the text list. This has to be implemented as an element inside such list, as elements can be added before the end of the maps, making it necessary to recalculate this positions.

The final script is this:

import re
import sys

def print_status(text, cursor, target):
    print(' '.join(map(str, text[:cursor])) + ' *CURSOR* ' + ' '.join(map(str, text[cursor:])), f"| target={target}")

def compose(commands, target):
    text = []
    cursor = 0
    stack = []
    SCOPE_END_DELIMITER = ('SCOPE_END',)

    for command in commands:
        if command.upper() == 'SWAP':
            target = f"transpose({target})"
        elif command.upper() == 'MAP':
            stack.append(cursor)
            text.insert(cursor, r'(map (\next -> ')
            text.insert(cursor + 1, f') ({target}))')
            text.insert(cursor + 2, SCOPE_END_DELIMITER)
            cursor += 1
            target = "next"
        elif command.upper() == 'RET':
            text.insert(cursor,f"({target})")

            if len(stack) < 1:
                raise Exception("Cannot RETurn, not inside a MAP")

            cursor = stack.pop()

            # Remove what's between the cursor and the scope delimiter and use that as target
            SCOPE_END = text.index(SCOPE_END_DELIMITER, cursor)
            text.pop(SCOPE_END)

            target = ' '.join(text[cursor:SCOPE_END])
            text = text[:cursor] + text[SCOPE_END:]  # text[SCOPE_END] was pop'ed just before
        else:
            if not command.startswith(':'):
                raise Exception(f"Found “{command}”, expected “SWAP”, “MAP” or something starting with ‘:’")
            function = command[1:]
            target = f'(({function}) ({target}))'

        # # Uncomment this to track the generated code
        # print(f"Command: {command}")
        # print_status(text, cursor, target)
        # print()

    text.insert(cursor, target)
    return ' '.join([chunk for chunk in text if chunk != SCOPE_END_DELIMITER])


def convert(stack_code, target_matrix):
    lines = stack_code.split('\n')
    no_comments = []
    for line in lines:
        if '--' in line:
            line = line[:line.index("--")]
        no_comments.append(line.strip())

    commands = re.split(r'\s+', " ".join(no_comments).strip())
    return compose(commands, target_matrix)


if __name__ == '__main__':
    if len(sys.argv) < 3:
        print(f"Usage: {sys.argv[0]} <base matrix> <operations>")
        exit(0)

    print(convert(' '.join(sys.argv[2:]), sys.argv[1]))

And this is an example of the output:

$ python converter.py 'mat4d' 'MAP MAP MAP :i32.sum RET :i32.sum RET :i32.sum RET :i32.sum'
((i32.sum) ((map (\next ->  (((i32.sum) ((map (\next ->  (((i32.sum) ((map (\next ->  (((i32.sum) (next))) ) (next))))) ) (next))))) ) (mat4d))))
$ futhark repl
[0]> let mat4d = [[[[1,2], [3,4]],[[5,6],[7,8]]],[[[10,11],[12,13]],[[14,15],[16,17]]]]
[1]> ((i32.sum) ((map (\next ->  (((i32.sum) ((map (\next ->  (((i32.sum) ((map (\next ->  (((i32.sum) (next))) ) (next))))) ) (next))))) ) (mat4d))))
144i32

Ok, that looks like enough about this matrix dimension operators, let's go back to the usual :)

Codigo para llevar [blog]

Matrix operators as stack languages 2: closing maps