Ariadne: Analysis for Machine Learning | Julian Dolby, IBM Thomas J. Watson Research Center

Analysis for Machine Learning

Julian Dolby, Avraham Shinnar

Allison Allain, Jenna Reinen

IBM Thomas J. Watson Research Center

https://juliandolby.github.io/mapl/talk/2018/06/14/mapl-ariadne-analysis.html

MAPL, PLDI, Philadelphia, June 2018

Elucidate Opaque Code

ML code manipulates complex, unspecified tensors

Ariadne tracks and surfaces tensor information

Outline

Describing tensors
Tensor analysis
Implementation in WALA
Initial results
Future work

Tensor Type Example

Symbolic dimension for runtime value
Compound dimension to capture structure

Tensor Types

$\newcommand{\trec}[1]{\{#1\}} \newcommand{\ttensor}[2]{[{#1} \; \text{of} \; {#2}]} \newcommand{\ttop}{\top} \begin{array}{lrcl} (\mbox{Python}) &\pi& ::= & \trec{f_1 : \pi_1 , \ldots , f_n:\pi_n}\\ &&\mid&(f_1 : \pi_1, \ldots, f_n:\pi_n)\to\pi\\ &&\mid& \tau\mid l\mid\ttop\\ (\mbox{Tensor})&\tau & ::=&\ttensor{d_1 , \ldots , d_n}{\pi}\\ (\mbox{Dim})&d & ::=& l \mid n \mid l(n) \mid d*d \\ \end{array}$

integrate with optional Python Types
capture structure within dimensions

Tensor Analysis

Track tensors through the program
Make information available everywhere
E.g., clear feed_dict has appropriate value for x

Tensor Analysis

Sometimes even inputs can be inferred
set_shape imputes tensor shapes to input values
Inputs can also be declared explicitly

Tensor Analysis

$T(y) \subseteq \left\{ \begin{array}{ll} \left\{\mathfrak{I}\right\} & \hbox{$y$ is input} \\ T(x) & y \prec x\\ z & \begin{array}{l} y \prec {\tt reshape(x, z)} \wedge\\ \exists_{z_i \in S(z)} T(x) \doteq z_i\\ \end{array}\\ \ldots & y \prec \hbox{other Tensorflow APIs} \end{array}\right.$

Input data is specified if necessary
Tensor operation semantics encoded
Tensor information follows data flow

Implementation in WALA

Exploit WALA support for translating ASTs
WALA Common AST handles much translation
Some work remained for idiosyncratic constructs
Python ASTs created with Jython
Much shared translation with JavaScript
Rich source position information vital
Greatly eases handling myriad popular languages

Simple Translation Example

public CAstNode visitIf(If arg0) throws Exception {
  return Ast.makeNode(CAstNode.IF_STMT,
    arg0.getInternalTest().accept(this),
    block(arg0.getInternalBody()),
    block(arg0.getInternalOrelse()));
}

IF_STMT provided by WALA
Jython accessors for node components

Evaluated programs

conv_network builds CNN to classify MNIST data.
mnist_deep classifies with convolutional layers.
mnist_max simple classifier from Tensorflow.
mnist_max_xla simple classifier from Tensorflow.
mnist_sum simple classifier from Tensorflow.
neuroimage classifies 3D brain images

Tensorflow constructs analyzed

${\tiny \begin{array}{|r|l|l|l|l|} \hline {\rm{program}} & {\rm{reshape}} & {\rm{conv2d}} & {\rm{conv3d}} & {\rm{placeholder}}\\ \hline conv\_network & ✓ & ✓ & ✗ & ✗ \\ mnist\_deep & ✓ & ✓ & ✗ & ✓\\ mnist\_max & ✗ & ✗ & ✗ & ✓ \\ mnist\_max\_xla & ✗ & ✗ & ✗ & ✓ \\ mnist\_sum & ✓ & ✗ & ✗ &✓ \\ neuroimage & ✓ & ✗& ✓ & ✗\\ \hline \end{array}}$

no false positives
limited modeling so far

See our poster

Poster at MAPL has more on IDE experience