# A machine learning study to identify collective flow in small and large
colliding systems

Han-Sheng Wang  Key Laboratory of Nuclear Physics and Ion-beam Application
(MOE), Institute of Modern Physics, Fudan University, Shanghai 200433, China
Shanghai Research Center for Theoretical Nuclear Physics, NSFC and Fudan
University, Shanghai $200438$, China  Shuang Guo  Key Laboratory of Nuclear
Physics and Ion-beam Application (MOE), Institute of Modern Physics, Fudan
University, Shanghai 200433, China  Shanghai Research Center for Theoretical
Nuclear Physics, NSFC and Fudan University, Shanghai $200438$, China  Kai Zhou
<EMAIL_ADDRESS>Frankfurt Institute for Advanced Studies, Ruth-
Moufang-Str. $1$, $60438$ Frankfurt am Main, Germany  Guo-Liang Ma
<EMAIL_ADDRESS>Key Laboratory of Nuclear Physics and Ion-beam Application
(MOE), Institute of Modern Physics, Fudan University, Shanghai 200433, China
Shanghai Research Center for Theoretical Nuclear Physics, NSFC and Fudan
University, Shanghai $200438$, China

###### Abstract

Collective flow has been found similar between small colliding systems (p+p
and p+A collisions) and large colliding systems (peripheral A+A collisions) at
the CERN Large Hadron Collider (LHC). In order to study the difference of
collective flow between small and large colliding systems, we employ a point
cloud network to identify $p$ $+$ Pb collisions and peripheral Pb $+$ Pb
collisions at $\sqrt{s_{NN}}=$ 5.02 TeV from a multiphase transport model
(AMPT). After removing the discrepancies in the pseudorapidity distribution
and the $p_{\rm T}$ spectra, we capture the discrepancy in anisotropic flow.
Although the verification accuracy of our PCN is limited due to similar event-
by-event distributions of elliptic and triangular flow, we demonstrate that
collective flow between $p$ $+$ Pb collisions and peripheral Pb $+$ Pb
collisions becomes more different with increasing final hadron multiplicity
and parton scattering cross section.

##  I INTRODUCTION

Quark-gluon plasma (QGP) at extreme conditions of high temperature and density
is thought to be a form of the early universe, which has been produced in the
laboratory by relativistic heavy-ion collisions at the BNL Relativistic Heavy
Ion Collider (RHIC) and the CERN Large Hadron Collider (LHC) Adams et al.
(2005); Adcox et al. (2005); Aamodt et al. (2008); Bzdak et al. (2020); Luo
and Xu (2017). The experimental results have shown that this new type of
strongly interacting matter behaves as a perfect fluid, because it translates
initial spatial geometry or initial energy density fluctuations into momentum
anisotropy of the final particles through pressure gradients in hydrodynamics
with the smallest ratio of shear viscosity over entropy Heinz and Snellings
(2013); Song and Heinz (2008); Jeon and Heinz (2015); Shen and Yan (2020);
Gale et al. (2013); Yan (2018); Alver and Roland (2010); Ma and Wang (2011);
Yan et al. (2014); Qin et al. (2010), thus generating strong collective flow
Ollitrault (1992); Stoecker (2005).

Over the last decade, the measurements of collective flow in various colliding
systems, such as $p$$+$$p$, $p$$+$Pb collisions at the LHC Khachatryan et al.
(2010); Chatrchyan et al. (2013); Abelev et al. (2013a); Aad et al. (2013a);
Sirunyan et al. (2018); Acharya et al. (2019), $p$$+$Au, $d$$+$Au, and
3He$+$Au collisions at RHIC Adare et al. (2015); Adamczyk et al. (2015);
Aidala et al. (2019); Acharya et al. (2022); Abdulameer et al. (2023), have
been performed. Surprisingly, similar collective flow is found in peripheral
A+A and high multiplicity $p$+A collisions at the same multiplicity, raising
doubts about whether QGP droplets can be also generated in small colliding
systems. Many theoretical efforts have been made to understand the origin of
collective flow in small colliding system Dusling et al. (2016); Loizides
(2016); Nagle and Zajc (2018). Basically, they can be divided into two
categories depending on whether the origin is from the initial or final state.
The color glass condensate (CGC) in the initial state has also been proposed
as a possible mechanism to interpret the experimentally measured “flow” in
small colliding systems Dumitru et al. (2011); Dusling and Venugopalan (2013);
Skokov (2015); Schenke et al. (2015); Schlichting and Tribedy (2016); Kovner
et al. (2017); Iancu and Rezaeian (2017); Mace et al. (2018); Nie et al.
(2019); Shi et al. (2021). Similarly to large colliding systems, hydrodynamics
in final sate also can transform the initial geometric asymmetry into the
final momentum anisotropic flow through the pressure gradient of the QGP in
small colliding systems Bozek (2012); Bzdak et al. (2013); Shuryak and Zahed
(2013); Qin and Müller (2014); Bozek and Broniowski (2013); Bozek et al.
(2015); Song et al. (2017). It is generally believed that the transport model
will behave more like hydrodynamics as the multiplicity or scattering cross
section increases, i.e., from non-equilibrium to equilibrium. It has been
demonstrated that a Multi-Phase Transport (AMPT) model Lin et al. (2005) is
capable of describing the experimental data in both large and small colliding
systems Bzdak and Ma (2014); Orjuela Koop et al. (2015); Ma and Bzdak (2016).
Since most of partons are not scattered in the small colliding systems at RHIC
and the LHC, a parton escape mechanism has been proposed to explain the
formation of azimuthal anisotropies in the transport model He et al. (2016);
Lin et al. (2016). It has been shown that parton collisions are crucial for
the generation of anisotropic flows Ma and Bzdak (2016); Ma et al. (2021). By
using a new test-particle method, we recently proved that collectivity
established by final state parton collisions is much stronger in large
colliding systems compared to that in small colliding systems Wang and Ma
(2022).

On the other hand, the event-averaged flow $v_{n}$ mainly reflects the
averaged hydrodynamic response to the initial collision geometry of the
produced matter. More information, such as the event-by-event (EbyE)
fluctuations of the overlap region Alver and Roland (2010), can be obtained by
measuring event-by-event $v_{n}$ distribution $P(v_{n})$ for charged hadrons,
which has been measured by the ATLAS Collaboration using unfolding method in
Pb$+$Pb collisions at $\sqrt{s_{NN}}$ = 2.76 TeV Aad et al. (2013b), which
gives a good constraint to the initial condition in A+A collisions. However,
the corresponding experimental measurements in small colliding systems are not
yet available.

In addition to the similarity of the anisotropic flow $v_{n}$, the differences
of other observables are worthy of investigation between large and small
colliding systems, such as pseudorapidity distribution and $p_{\rm T}$
spectra. For example, previous experimental studies have shown power law-
shaped $p_{\rm T}$ spectra in small colliding systems Adams et al. (2006);
Abelev et al. (2014); Adam et al. (2015); Acharya et al. (2020); Abelev et al.
(2009, 2013b), unlike exponential-shaped $p_{\rm T}$ spectra in large
colliding systems Adcox et al. (2004); Abelev et al. (2009, 2013b).

With the advancement of computational hardware and algorithms, machine
learning (ML) techniques have become a popular and powerful tool for
extracting information from big data, and has been applied in high energy
nuclear physics for various topics Pang et al. (2018); Zhou et al. (2019);
Graczykowski et al. (2022); ZHOU et al. (2022); WANG et al. (2022); HE et al.
(2022); DU et al. (2022); LI et al. (2023). Among many deep neural network
architectures within ML, the point cloud network (PCN) is a highly efficient
and effective architecture that can solve many problems by involving point
cloud structured record, such as 3D object segmentation and scene semantic
parsing Qi et al. (2016). It has also been used in high energy nuclear
physics, such as reconstructing the impact parameter of collisions precisely
Steinheimer et al. (2019); Omana Kuttan et al. (2020a); Omana Kuttan et al.
(2021), classifying the equation of state (EOS) Omana Kuttan et al. (2020b),
and identifying weak intermittency signals associated with critical phenomena
Huang et al. (2022) by learning the data for heavy-ion collisions. In this
work, by adopting PCN in supervised training, we will find out if it is
possible to distinguish between $p$$+$Pb collisions and peripheral Pb$+$Pb
collisions from the multiphase transport (AMPT) model, and reveal the
differences between small and large colliding systems.

The paper is organized as follows. First, we introduce the AMPT model, which
generates the data for $p$$+$Pb collisions and peripheral Pb$+$Pb collisions,
and the relevant observables in Sec. II. In Sec. III, we describe the details
of the PCN. In Sec. IV, the validation accuracy of PCN is presented and the
relevant physics is discussed. Finally, we summarize the implications of our
results in Sec. V.

##  II MODEL AND METHOD

###  II.1 A multiphase transport model

The string melting version of AMPT model consists of four main stages of
heavy-ion collisions, i.e. initial state, parton cascade, hadronization, and
hadronic rescatterings. The initial state with fluctuating initial conditions
is generated by the heavy ion jet interaction generator (HIJING) model
Gyulassy and Wang (1994). In HIJING model, minijet partons and excited strings
are produced by hard processes and soft processes, respectively. In the string
melting mechanism, all excited hadronic strings in the overlap volume are
converted to partons according to the flavor and spin structures of their
valence quarks Lin and Ko (2002). The partons are generated by string melting
after a formation time:

$t_{f}=E_{H}/m^{2}_{T,H},$ (1)

where $E_{H}$ and $m_{T,H}$ represent the energy and transverse mass of the
parent hadron. The initial positions of partons originating from melted
strings are determined by tracing their parent hadrons along straight-line
trajectories. The interactions among partons are described by the Zhang’s
parton cascade (ZPC) model Zhang (1998), which includes only two-body elastic
scatterings with a g+g $\rightarrow$ g+g cross section, i.e.,

$\frac{d\sigma}{d\hat{t}}=\frac{9\pi\alpha^{2}_{s}}{2}(1+\frac{\mu^{2}}{\hat{s}})\frac{1}{(\hat{t}-\mu^{2})^{2}},$
(2)

where $\alpha_{s}$ is the strong coupling constant (taken as 0.33), while
$\hat{s}$ and $\hat{t}$ are the usual Mandelstam variables. The effective
screening mass $\mu$ is taken as a parameter in ZPC for adjusting the parton
scattering cross section. Note that previous AMPT model studies have shown
that a parton scattering cross section of 3 mb can well describe both large
and small colliding systems at RHIC and the LHC energies Lin (2014); Orjuela
Koop et al. (2015); Ma and Bzdak (2016); Ma and Lin (2016); He and Lin (2017);
Lin and Zheng (2021). A quark coalescence model is used for hadronization at
the freeze out of parton system. The hadronic scatterings in hadronic phase
are simulated by a relativistic transport (ART) model Li and Ko (1995).

In this study, we simulated 2 million events each for peripheral Pb$+$Pb
collisions and $p$$+$Pb minimum bias collisions at $\sqrt{s_{NN}}$ = 5.02 TeV,
using a string-melting version of the AMPT model with parton scattering cross
sections of 0 mb, 3 mb, and 10 mb, respectively.

###  II.2 Anisotropic flow

To quantify the azimuthal anisotropy $v_{n}$ in final momentum space, the
azimuthal angle distribution of measured particles can be decomposed into a
Fourier series,

$\displaystyle\frac{{\rm d}N}{{\rm
d}\phi}=\frac{1}{2\pi}\Bigg{[}1+\sum_{n=1}^{\infty}2v_{n}\cos{\left(n(\phi-\psi_{n})\right)}\Bigg{]},$
(3)

where $\phi$ is the azimuthal angle, $\psi_{n}$ is the $n^{\rm th}$ harmonic
event plane angle, $v_{n}=\langle\cos{[n(\phi-\psi_{n})]}\rangle$ is the
$n^{\rm th}$ order of harmonic flow coefficient Voloshin and Zhang (1996). The
second ($v_{2}$) and third ($v_{3}$) Fourier coefficients represent the
amplitude of elliptic and triangular flow, respectively. The linearized
hydrodynamic response shows that anisotropy flow $v_{n}$ is correlated to the
geometry asymmetry of energy density profile in spatial space of initial
state, namely the initial eccentricity $\varepsilon_{n}$,

$\varepsilon_{n}=\frac{\sqrt{\langle r^{n}\cos(n\phi_{r})\rangle^{2}+\langle
r^{n}\sin(n\phi_{r})\rangle^{2}}}{\langle r^{n}\rangle},$ (4)

where $r$ and $\phi_{r}$ are the polar coordinates of participating nucleons
Alver and Roland (2010). In hydrodynamics, harmonic flows are responding to
eccentricities,

$v_{n}=v_{n}(\varepsilon_{n},k),$ (5)

where the constant $k$ is sensitive to the properties of the QGP, such as the
transport coefficient $\eta/s$ Yan (2018).

Input PointCloudN$\times$d$\begin{matrix}\text{input}\\\
\text{transform}\end{matrix}$N$\times$dMLP(32)sharedN$\times$32$\begin{matrix}\text{feature}\\\
\text{transform}\end{matrix}$N$\times$32MLP(32,64,512)sharedN$\times$512$\begin{matrix}\text{max}\\\
\text{pool}\end{matrix}$512$\begin{matrix}\text{MLP}\\\
\text{(256,128,2)}\end{matrix}$global feature2output scores Figure 1: The
architecture of the devised point cloud network (PCN). In this network, it
takes final state $N$ points as input, applies input and feature
transformations, and then aggregates point features by max-pooling. The output
is classification scores for 2 classes, i.e., the probability of Pb$+$Pb or
$p$$+$Pb colliding systems.

On the other hand, the EbyE distribution of $v_{n}$, P($v_{n}$), can be
obtained by an unfolding method, which suppresses the nonflow contribution.
However, there is some problem with the unfolding process for small systems.
For $p+A$ systems, the response matrix in the unfolding method can not be
obtained reliably. For example, there is no convergent result in p+p system.
Therefore, the unfolding method is not used in this paper.

##  III Training the PCN for classifying two systems

In this section, we introduce the detailed analysis procedures, including the
input and output, the PCN architecture, training, and evaluation of the PCN.

###  III.1 Input to the machine

The AMPT events for both $p$$+$Pb and Pb$+$Pb collisions are grouped into the
data sets according to the number of final charged particles $N_{ch}$ measured
in the kinetic window of $|\eta|<$ 2.4 and $p_{\rm T}>$ 0.4 GeV/c. These data
are divided into various centrality classes for the colliding systems, which
are the $N_{ch}$ bins of 60-90, 90-120, 120-150, 150-185, 185-230. Each data
set has about 0.2 million events which are separated into training events and
validation events by a ratio of 70% to 30%. All the information about the
final state particles within $|\eta|<$ 2.4 in each event are the input to the
PCN as a sample, which consists of a list of particles with their information
on ($p_{\rm x}$, $p_{\rm y}$, $p_{z}$, $E$) or ($p_{\rm x}$, $p_{\rm y}$).
Simultaneously, the two true labels, $p$$+$Pb and Pb$+$Pb, are marked on each
event to perform the supervised training.

###  III.2 Network architecture

The architecture of PCN is shown in Fig.1. It begins with an input alignment
network. The following are a shared pointwise multilayer perceptron (MLP)
implemented by 1D-convolution neural network (CNN) to extract 32 feature maps,
a feature alignment network, a shared MLP to extract 32, 64 and 512 feature
maps, respectively. A global max pooling then gets the maximum values of each
feature among all particles as one global feature of the particle cloud.
Finally, a shared MLP implemented by 3 layers fully connected deep neural
network (DNN) with 256, 128, and 2 neurons tags each event as $p$$+$Pb or
Pb$+$Pb collision. Batch normalization layers are present between every
convolution layer. The LeakyReLU($\alpha$=0.01) activation function is used
for all layers except the final layer. The sigmoid activation is used on the
final layer for binary classification. The models use the Adam optimizer with
a learning rate of $10^{-4}$ with total decay $10^{-4}$ and categorical cross
entropy as the loss function. In addition, a dropout layers (with drop out
probability 0.3) and L2 regularization are present to tackle the overfitting
issue. We use a maximum of 50 epochs with 32 batch size to train the data set.

##  IV Training results and discussion

In our study, two cases of input data have been investigated. In the case 1,
the input for training is an EbyE list of four-momentum ($p_{\rm x}$, $p_{\rm
y}$, $p_{z}$, $E$) of the selected final state hadrons from the AMPT model.
The input data have a dimension of $N\times 4$, where $N$ is the maximum
number of selected particles in an event. Events with fewer particles are
filled with zeros to maintain the same input dimension. The case 2 is the same
as the case 1, but the input is a list of two-momentum ($p_{\rm x}$, $p_{\rm
y}$) of selected final state hadrons.

###  IV.1 Case 1: training with four-momentum of final hadrons

Figure 2: (Color online) Validation accuracy as a function of the number of
charged particles $N_{ch}$ by training with four-momentum ($p_{\rm x}$,
$p_{\rm y}$, $p_{z}$, $E$) of final particles in pseudorapidity ranges
-2.4$<\eta<$2.4 (solid curve), -2.4$<\eta<$0 (dashed curve), and 0$<\eta<$2.4
(dotted curve).

Figure 2 shows validation accuracy as a function of the number of charged
particles $N_{ch}$ by learning four-momentum ($p_{\rm x}$, $p_{\rm y}$,
$p_{z}$, $E$) of final particles in three different pseudorapidity ranges. All
the validation accuracies increase as the multiplicity increases, which
indicates that the two colliding systems are more distinguishable for the
events with higher multiplicity. The validation accuracies are very close for
the pseudorapidity ranges of -2.4$<\eta<$2.4 and 0$<\eta<$2.4. However, there
are about 10% drops in validation accuracy for the pseudorapidity range of
-2.4$<\eta<$0 relative to the above two cases. This indicates the discrepancy
between the two colliding systems mainly comes from the difference in the
forward pseudorapidity range, which is the proton-going direction in $p+$Pb
systems.

Figure 3: (Color online) The pseudorapidity distributions of final hadrons for
different $N_{ch}$ classes in Pb$+$Pb collisions (solid curve) and $p$ $+$ Pb
collisions (dashed curve).

To find why the validation accuracy is sensitive to the pseudorapidity range,
Fig. 3 compares the pseudorapidity distributions of final hadrons in Pb$+$Pb
collisions and $p$$+$Pb collisions for different $N_{ch}$ classes. Between
Pb$+$Pb and $p$ $+$ Pb systems, there is a significant difference in the
pseudorapidity distribution due to the lower hadron yield in $p$ beam
(forward) direction for $p$$+$Pb collisions. This leads to the discrepancy in
the above validation accuracies when considering different pseudorapidity
ranges. Since the difference in the pseudorapidity distribution in the forward
direction becomes more significant at larger $N_{ch}$, it explains why the
validation accuracy increases with multiplicity when we train PCN by learning
the four-momentum of final particles in Fig.2.

###  IV.2 Case 2: training with two-momentum of final hadrons

To investigate whether the PCN can learn the difference of collective flow
between $p+$Pb collisions and peripheral Pb$+$Pb collisions, three scenarios
of lists are used as training inputs, i.e., two-momentum ($p_{\rm x}$, $p_{\rm
y}$), normalized two-momentum ($p_{\rm x}^{\rm norm}$, $p_{\rm y}^{\rm
norm}$), and normalized meanwhile randomly rotated two-momentum ($p_{\rm
x}^{\rm norm,rand}$, $p_{\rm y}^{\rm norm,rand}$), respectively. The
normalized two-momentum $p_{\rm x}^{\rm norm}$ and $p_{\rm y}^{\rm norm}$ are
defined as,

$\displaystyle p_{\rm x}^{\rm norm}$ $\displaystyle=$
$\displaystyle\frac{p_{\rm x}}{p_{\rm T}},$ $\displaystyle p_{\rm y}^{\rm
norm}$ $\displaystyle=$ $\displaystyle\frac{p_{\rm y}}{p_{\rm T}},$ (6)

where $p_{\rm T}$ is the transverse momentum of each particle. The randomly
rotated two-momentum $p_{\rm x}^{\rm rand}$ and $p_{\rm y}^{\rm rand}$ are,

$\displaystyle p_{\rm x}^{\rm rand}$ $\displaystyle=$ $\displaystyle p_{\rm
x}\times\cos{\phi_{\rm rand}}-p_{\rm y}\times\sin{\phi_{\rm rand}},$
$\displaystyle p_{\rm y}^{\rm rand}$ $\displaystyle=$ $\displaystyle p_{\rm
x}\times\sin{\phi_{\rm rand}}-p_{\rm y}\times\cos{\phi_{\rm rand}},$ (7)

where $\phi_{\rm rand}$ is a random angle between 0 and 2$\pi$. While, the
normalized and randomly rotated two-momentum $p_{\rm x}^{\rm norm,rand}$ and
$p_{\rm y}^{\rm norm,rand}$ are,

$\displaystyle p_{\rm x}^{\rm norm,rand}$ $\displaystyle=$
$\displaystyle\frac{p_{\rm x}\times\cos{\phi_{\rm rand}}-p_{\rm
y}\times\sin{\phi_{\rm rand}}}{p_{\rm T}},$ $\displaystyle p_{\rm y}^{\rm
norm,rand}$ $\displaystyle=$ $\displaystyle\frac{p_{\rm x}\times\sin{\phi_{\rm
rand}}-p_{\rm y}\times\cos{\phi_{\rm rand}}}{p_{\rm T}}.$ (8)

Figure 4: (Color online) Validation accuracy as a function of the number of
charged particles $N_{ch}$ by learning with input to be two-momentum ($p_{\rm
x}$, $p_{\rm y}$) (a) and the normalized two-momentum ($p_{\rm x}^{\rm norm}$,
$p_{\rm y}^{\rm norm}$) (b), each without (solid curve) or with (dashed curve)
random rotations.

Figure 4(a) shows validation accuracy as a function of the number of final
charged particles $N_{ch}$ with input to be two-momentum ($p_{\rm x}$, $p_{\rm
y}$) and its random rotation ($p_{\rm x}^{\rm rand}$ and $p_{\rm y}^{\rm
rand}$) of final particles. Validation accuracy of two-momentum input is over
$60\%$, and it will drop less than $2\%$, if two-momentum are randomly
rotated. It indicates that beside anisotropic flow there is other discrepancy
between the two colliding systems, like in $p_{\rm T}$ distributions. To
eliminate the effect from different $p_{\rm T}$ distribution of hadron, we
train the PCN with normalized two-momentum ($p_{\rm x}^{\rm norm}$, $p_{\rm
y}^{\rm norm}$) and its random rotation ($p_{\rm x}^{\rm norm,rand}$ and
$p_{\rm y}^{\rm norm,rand}$) of final particles. The validation accuracy as a
function of the number of final charged particles is shown in Fig.4(b).
Compared to the two-momentum training, there is more than $10\%$ decrease of
validation accuracy by learning two-momentum after normalization, which
indicates that $p_{\rm T}$ distribution of hadrons is the main feature that
distinguishes the two colliding systems, if we train the PCN by two-momentum.
Furthermore, the PCN can not distinguish the two colliding systems by learning
normalized two-momentum with random rotations of final hadrons, because these
operations eliminate the information about the magnitude of $p_{\rm T}$ and
anisotropic flow. Compared to the training by the normalized two-momentum with
random rotations, the validation accuracy is improved by a few percentages by
learning two-momentum with the normalization only. This discrepancy comes from
the difference of anisotropic flow between two colliding systems. However, the
anisotropic flow between two colliding systems is so similar that the PCN can
not distinguish between different colliding systems very well, even though it
is better than the case with random rotations.

Figure 5: (Color online) Validation accuracy as a function of the number of
charged particles $N_{ch}$ by learning two-momentum ($p_{\rm x}$, $p_{\rm y}$)
of final hadrons in the AMPT model with different parton cross sections.

Figure 5 shows validation accuracy as a function of the number of final
charged particles $N_{ch}$ by learning two-momentum ($p_{\rm x}$, $p_{\rm y}$)
of final particles with different parton cross sections. It can be seen that
all validation accuracies increase with the increase of $N_{ch}$ and parton
cross section. In addition, the PCN can not distinguish the two colliding
systems with 0 mb parton cross section at low $N_{ch}$. These results can
basically be explained by Fig. 6, since the average transverse momentum
$\left\langle p_{\rm T}\right\rangle$ can quantify the feature of $p_{\rm T}$
spectra shape.

Figure 6: (Color online) Average transverse momenta $\left\langle p_{\rm
T}\right\rangle$ in Pb $+$ Pb collisions (solid curve) and p $+$ Pb collisions
(dashed curve) as a function of the number of charged particles $N_{ch}$ with
parton cross section of 0 mb (a), 3 mb (b), and 10 mb (c).

Figure 6 shows average transverse momenta $\left\langle p_{\rm
T}\right\rangle$ as a function of the number of final charged particles
$N_{ch}$ with different parton cross sections. We can see that the
$\left\langle p_{\rm T}\right\rangle$ in p $+$ Pb collisions is larger than
that in Pb $+$ Pb collisions, and the difference is more significant with a
larger parton cross section and higher multiplicity. Although we observe that
the difference of $\left\langle p_{\rm T}\right\rangle$ is almost zero for 0
mb parton cross section, some validation accuracies are larger than 0.5 for 0
mb in Fig. 5. It indicates that there is a difference of non-flow effect
between small and large colliding systems. It could be attributed to jets,
because the effects of the jet transverse momentum broadening and multiple
scatterings Leonidov et al. (1997) have been found to be stronger in p $+$ Pb
collisions than in peripheral Pb $+$ Pb collisions Abelev et al. (2009).

Figure 7: (Color online) Validation accuracy as a function of the number of
charged particles $N_{ch}$ by learning normalized two-momentum ($p_{\rm
x}^{\rm norm}$, $p_{\rm y}^{\rm norm}$) of final hadrons in the AMPT model
with different parton cross sections.

Figure 7 shows validation accuracy as a function of the number of final
charged particles $N_{ch}$ by learning normalized two-momentum ($p_{\rm
x}^{\rm norm}$, $p_{\rm y}^{\rm norm}$) of final particle with different
parton cross sections. It can be seen that the validation accuracy with 0 mb
is close to the corresponding result in Fig. 5, and both of them increase with
$N_{ch}$, which indicates that the discrepancy of non-flow effect between the
two systems is larger for higher multiplicity. However, the discrepancy
between 0 mb and nonzero parton cross sections must come from anisotropic
flow. The discrepancy increases with parton cross section, which indicates
that there is larger discrepancy of anisotropic flow between two colliding
systems due to more parton collisions.

Figure 8: (Color online) Averaged output
$\left\langle\widehat{y}\right\rangle$ of the ensembles of two systems
identified by the model trained by normalized two-momentum ($p_{\rm x}^{\rm
norm}$, $p_{\rm y}^{\rm norm}$) of final hadrons in the AMPT model with
different parton cross sections.

If an event is identified as $p$$+$Pb or Pb $+$ Pb, it is marked by the output
of $\widehat{y}$=0 or 1, respectively, in our analysis. Thus the averaged
output $\left\langle\widehat{y}\right\rangle$ of our model can represent the
probability of Pb $+$ Pb of a event. In other words, The closer the
$\left\langle\widehat{y}\right\rangle$ is to 0, the more likely the events are
$p$$+$Pb collisions, and vice versa for Pb $+$ Pb collisions. Figure 8 shows
the averaged output $\left\langle\widehat{y}\right\rangle$ of the ensembles of
each system identified by the model trained by normalized two-momentum
($p_{\rm x}^{\rm norm}$, $p_{\rm y}^{\rm norm}$) of final hadrons in the AMPT
model with different parton cross sections. It can been seen that
$\left\langle\widehat{y}\right\rangle$ of the ensembles of Pb $+$ Pb events is
larger than that for $p$$+$Pb events in every situation, which indicates that
the PCN is able to distinguish two different ensembles of systems, according
to the difference of averaged anisotropic flow between two systems.

To investigate the relationship between anisotropic flow and parton cross
section, we further calculate the harmonic flow coefficients, e.g., elliptic
flow $v_{2}$, by fitting the long range part of the two-particle azimuthal
correlation function $C(\Delta\phi)$, which defined as

$C(\Delta\phi)=\frac{Y_{same}(\Delta\phi)}{Y_{mixed}(\Delta\phi)}\times\frac{\int
Y_{mixed}(\Delta\phi)d\Delta\phi}{\int Y_{same}(\Delta\phi)d\Delta\phi},$ (9)

where $Y_{same}(\Delta\phi=\phi_{2}-\phi_{1})$ and $Y_{mixed}(\Delta\phi)$ are
the number of particle pairs at a given $\Delta\phi$ and within a given
$p_{\rm T}$ range for the same and mixed events. This definition of
$C(\Delta\phi)$ removes a trivial dependence on the number of produced
particles Adler et al. (2006); Adare et al. (2015); Bzdak and Ma (2014).

Figure 9: (Color online) Integrated elliptic $v_{2}$ and triangular $v_{3}$
flow coefficient in Pb $+$ Pb collisions (solid curve) and p $+$ Pb collisions
(dashed curve) as a function of the number of charged particles $N_{ch}$ with
parton cross section of 0 mb (a), 3 mb (b), and 10 mb (c).

Figure 9 shows $v_{2}$ and $v_{3}$ as a function of the number of final
charged particles $N_{ch}$ with different parton cross sections. It can be
observed that $v_{2}$ and $v_{3}$ increases with parton cross section for
Pb$+$Pb collisions. But the dependence on parton cross section is non-
monotonous for $p$$+$Pb collisions, which has already been found in Ref. Zhao
et al. (2023). It indicates that collective flow can be not only built up but
also damaged by the parton collisions in small colliding systems Wang and Ma
(2022). The most significant difference of $v_{2}$ between the two systems
appears when parton cross section is taken as 10 mb for high-multiplicity
events, because there are the largest $v_{2}$ in Pb$+$Pb collisions and
relatively small $v_{2}$ in $p$$+$Pb collisions. On the other hand, $v_{3}$ of
the two systems are similar for the two parton cross sections of 0 mb and 10
mb. However, there is more obvious difference of $v_{3}$ between the two
systems with parton cross section of 3 mb, because there are the largest
$v_{3}$ in $p$$+$Pb collisions and relatively small $v_{3}$ in Pb$+$Pb
collisions. Surprisingly, the discrepancy of $v_{2}$ and $v_{3}$ between two
ensembles of each system can be captured by the PCN, even the similar $N_{ch}$
and parton cross section dependences are obtained in Figs. 7 and 8, although
the validation accuracies are not high enough for the PCN to distinguish two
systems due to similar EbyE flow distributions which will be shown next.

Figure 10: (Color online) The EbyE 2-D distributions of $v_{2}$ vs $v_{3}$
[P($v_{2}$, $v_{3}$)] of Pb$+$Pb (red) and $p$$+$Pb (blue) collisions in the
AMPT model with the parton cross section of 3 mb (first row) and 10 mb (second
row) for the $N_{ch}$ bin of 90$<N_{ch}<$120 (first column), 120$<N_{ch}<$150
(second column), 150$<N_{ch}<$185 (third column), and 185$<N_{ch}<$230 (fourth
column). Figure 11: (Color online) The percentage of the 2-D overlapping area
between the distribution P($v_{2}$) and P($v_{3}$) for Pb$+$Pb collisions and
that for $p$$+$Pb collisions as a function of the number of final charged
particles $N_{ch}$ in the AMPT model with the parton cross section of 0 mb
(solid curve), 3 mb (dashed curve), and 10 mb (dotted curve).

Figure 10 shows the EbyE 2-D distributions of $v_{2}$ vs $v_{3}$ [P($v_{2}$,
$v_{3}$)] in Pb$+$Pb collisions and $p$$+$Pb collisions with parton cross
sections of 3 and 10 mb for different $N_{ch}$ classes. Based on P($v_{2}$,
$v_{3}$), the percentage of the overlapping area of P($v_{2}$, $v_{3}$)
between Pb$+$Pb collisions and $p$$+$Pb collisions as a function of the number
of final charged particles $N_{ch}$ can be calculated, which is shown in Fig.
11. Note that the result for 0 mb parton cross section is not shown, since
$v_{n}$ come from non-flow for this case. We can see that the overlapping
percentage decreases with parton cross section, which indicates that $v_{2}$
and $v_{3}$ between two systems are more different with a larger parton cross
section. This is also consistent with the result of validation accuracy in
Fig. 7. It suggests that more pronounced difference of collective flow between
the two colliding systems is produced by more parton collisions. On the other
hand, the large overlapping volume percentage (over $80\%$) indicates that the
P($v_{2}$, $v_{3}$) distributions between the two systems are so similar that
they are hard to be identified. This is also in line with the observation that
PCN can not distinguish the two colliding systems very well in EbyE manner,
when we train PCN with input to be the normalized two-momentum of final
particles, as shown in Fig. 7.

##  V Summary and outlook

In summary, we employ the point cloud network to identify the events of
$p$$+$Pb and peripheral Pb$+$Pb collisions from a multiphase transport model.
Many different features between the two systems are learned and captured by
the point cloud network, such as pseudorapidity distribution, $p_{\rm T}$
spectra, anisotropic flow. In the transverse momentum plane, the point cloud
network can learn the different features of $p_{\rm T}$ spectra that can
classify two different colliding systems. After normalizing transverse
momentum of final hadrons, the point cloud network finally distinguish two
different colliding systems according to their collective flow. It is
difficult for the PCN to distinguish two systems by EbyE collective flow only,
since the EbyE distributions of collective flow P($v_{2}$, $v_{3}$) are very
close between $p$$+$Pb and Pb$+$Pb collisions. However, the PCN is able to
distinguish between two ensembles of each system though the feature of $v_{2}$
and $v_{3}$, and obtain the $N_{ch}$ and parton cross section dependence of
the discrepancy between two systems.

In this big data and ML approach, by changing the different input types, we
confirm that the discrepancy between two systems is more reflected in the
pseudorapidity distribution and the $p_{\rm T}$ spectra than in the
anisotropic flow. We also demonstrate that the collective flow between
$p$$+$Pb and Pb$+$Pb collisions are more different with a larger parton
scattering cross section, which is the consistent with the characteristics of
escape mechanism for collective flow in transport model He et al. (2016); Lin
et al. (2016).

Admittedly, the ability of our PCN to distinguish $p$$+$Pb from peripheral
Pb$+$Pb collisions in EbyE manner by collective flow is limited. However, some
meaningful results have been obtained by using the useful tool of machine
learning. In future, the method in this study could be applied to investigate
other rich physics topics, such as identifying initial flow and final flow in
small colliding systems Nie et al. (2019).

###### Acknowledgements.

This work is supported in part by the National Natural Science Foundation of
China under Grants No.12147101, No. 11890714, No. 11835002, No. 11961131011,
No. 11421505, the National Key Research and Development Program of China under
Grant No. 2022YFA1604900, the Strategic Priority Research Program of Chinese
Academy of Sciences under Grant No. XDB34030000, and the Guangdong Major
Project of Basic and Applied Basic Research under Grant No. 2020B0301030008
(H.-S.W. S. G. and G.-L.M.), the Germany BMBF under the ErUM-Data project (K.
Z.).

## References

  * Adams et al. (2005) J. Adams et al. (STAR), Nucl. Phys. A 757, 102 (2005), eprint nucl-ex/0501009.
  * Adcox et al. (2005) K. Adcox et al. (PHENIX), Nucl. Phys. A 757, 184 (2005), eprint nucl-ex/0410003.
  * Aamodt et al. (2008) K. Aamodt et al. (ALICE), JINST 3, S08002 (2008).
  * Bzdak et al. (2020) A. Bzdak, S. Esumi, V. Koch, J. Liao, M. Stephanov, and N. Xu, Phys. Rept. 853, 1 (2020), eprint 1906.00936.
  * Luo and Xu (2017) X. Luo and N. Xu, Nucl. Sci. Tech. 28, 112 (2017), eprint 1701.02105.
  * Heinz and Snellings (2013) U. Heinz and R. Snellings, Ann. Rev. Nucl. Part. Sci. 63, 123 (2013), eprint 1301.2826.
  * Song and Heinz (2008) H. Song and U. W. Heinz, Phys. Rev. C 77, 064901 (2008), eprint 0712.3715.
  * Jeon and Heinz (2015) S. Jeon and U. Heinz, Int. J. Mod. Phys. E 24, 1530010 (2015), eprint 1503.03931.
  * Shen and Yan (2020) C. Shen and L. Yan, Nucl. Sci. Tech. 31, 122 (2020), eprint 2010.12377.
  * Gale et al. (2013) C. Gale, S. Jeon, and B. Schenke, Int. J. Mod. Phys. A 28, 1340011 (2013), eprint 1301.5893.
  * Yan (2018) L. Yan, Chin. Phys. C 42, 042001 (2018), eprint 1712.04580.
  * Alver and Roland (2010) B. Alver and G. Roland, Phys. Rev. C 81, 054905 (2010), [Erratum: Phys.Rev.C 82, 039903 (2010)], eprint 1003.0194.
  * Ma and Wang (2011) G.-L. Ma and X.-N. Wang, Phys. Rev. Lett. 106, 162301 (2011), eprint 1011.5249.
  * Yan et al. (2014) L. Yan, J.-Y. Ollitrault, and A. M. Poskanzer, Phys. Rev. C 90, 024903 (2014), eprint 1405.6595.
  * Qin et al. (2010) G.-Y. Qin, H. Petersen, S. A. Bass, and B. Muller, Phys. Rev. C 82, 064903 (2010), eprint 1009.1847.
  * Ollitrault (1992) J.-Y. Ollitrault, Phys. Rev. D 46, 229 (1992).
  * Stoecker (2005) H. Stoecker, Nucl. Phys. A 750, 121 (2005), eprint nucl-th/0406018.
  * Khachatryan et al. (2010) V. Khachatryan et al. (CMS), JHEP 09, 091 (2010), eprint 1009.4122.
  * Chatrchyan et al. (2013) S. Chatrchyan et al. (CMS), Phys. Lett. B 718, 795 (2013), eprint 1210.5482.
  * Abelev et al. (2013a) B. Abelev et al. (ALICE), Phys. Lett. B 719, 29 (2013a), eprint 1212.2001.
  * Aad et al. (2013a) G. Aad et al. (ATLAS), Phys. Rev. Lett. 110, 182302 (2013a), eprint 1212.5198.
  * Sirunyan et al. (2018) A. M. Sirunyan et al. (CMS), Phys. Rev. C 98, 044902 (2018), eprint 1710.07864.
  * Acharya et al. (2019) S. Acharya et al. (ALICE), Phys. Rev. Lett. 123, 142301 (2019), eprint 1903.01790.
  * Adare et al. (2015) A. Adare et al. (PHENIX), Phys. Rev. Lett. 114, 192301 (2015), eprint 1404.7461.
  * Adamczyk et al. (2015) L. Adamczyk et al. (STAR), Phys. Lett. B 747, 265 (2015), eprint 1502.07652.
  * Aidala et al. (2019) C. Aidala et al. (PHENIX), Nature Phys. 15, 214 (2019), eprint 1805.02973.
  * Acharya et al. (2022) U. A. Acharya et al. (PHENIX), Phys. Rev. C 105, 024901 (2022), eprint 2107.06634.
  * Abdulameer et al. (2023) N. J. Abdulameer et al. (PHENIX), Phys. Rev. C 107, 024907 (2023), eprint 2203.09894.
  * Dusling et al. (2016) K. Dusling, W. Li, and B. Schenke, Int. J. Mod. Phys. E 25, 1630002 (2016), eprint 1509.07939.
  * Loizides (2016) C. Loizides, Nucl. Phys. A 956, 200 (2016), eprint 1602.09138.
  * Nagle and Zajc (2018) J. L. Nagle and W. A. Zajc, Ann. Rev. Nucl. Part. Sci. 68, 211 (2018), eprint 1801.03477.
  * Dumitru et al. (2011) A. Dumitru, K. Dusling, F. Gelis, J. Jalilian-Marian, T. Lappi, and R. Venugopalan, Phys. Lett. B 697, 21 (2011), eprint 1009.5295.
  * Dusling and Venugopalan (2013) K. Dusling and R. Venugopalan, Phys. Rev. D 87, 094034 (2013), eprint 1302.7018.
  * Skokov (2015) V. Skokov, Phys. Rev. D 91, 054014 (2015), eprint 1412.5191.
  * Schenke et al. (2015) B. Schenke, S. Schlichting, and R. Venugopalan, Phys. Lett. B 747, 76 (2015), eprint 1502.01331.
  * Schlichting and Tribedy (2016) S. Schlichting and P. Tribedy, Adv. High Energy Phys. 2016, 8460349 (2016), eprint 1611.00329.
  * Kovner et al. (2017) A. Kovner, M. Lublinsky, and V. Skokov, Phys. Rev. D 96, 016010 (2017), eprint 1612.07790.
  * Iancu and Rezaeian (2017) E. Iancu and A. H. Rezaeian, Phys. Rev. D 95, 094003 (2017), eprint 1702.03943.
  * Mace et al. (2018) M. Mace, V. V. Skokov, P. Tribedy, and R. Venugopalan, Phys. Rev. Lett. 121, 052301 (2018), [Erratum: Phys.Rev.Lett. 123, 039901 (2019)], eprint 1805.09342.
  * Nie et al. (2019) M. Nie, L. Yi, G. Ma, and J. Jia, Phys. Rev. C 100, 064905 (2019), eprint 1906.01422.
  * Shi et al. (2021) Y. Shi, L. Wang, S.-Y. Wei, B.-W. Xiao, and L. Zheng, Phys. Rev. D 103, 054017 (2021), eprint 2008.03569.
  * Bozek (2012) P. Bozek, Phys. Rev. C 85, 014911 (2012), eprint 1112.0915.
  * Bzdak et al. (2013) A. Bzdak, B. Schenke, P. Tribedy, and R. Venugopalan, Phys. Rev. C 87, 064906 (2013), eprint 1304.3403.
  * Shuryak and Zahed (2013) E. Shuryak and I. Zahed, Phys. Rev. C 88, 044915 (2013), eprint 1301.4470.
  * Qin and Müller (2014) G.-Y. Qin and B. Müller, Phys. Rev. C 89, 044902 (2014), eprint 1306.3439.
  * Bozek and Broniowski (2013) P. Bozek and W. Broniowski, Phys. Rev. C 88, 014903 (2013), eprint 1304.3044.
  * Bozek et al. (2015) P. Bozek, A. Bzdak, and G.-L. Ma, Phys. Lett. B 748, 301 (2015), eprint 1503.03655.
  * Song et al. (2017) H. Song, Y. Zhou, and K. Gajdosova, Nucl. Sci. Tech. 28, 99 (2017), eprint 1703.00670.
  * Lin et al. (2005) Z.-W. Lin, C. M. Ko, B.-A. Li, B. Zhang, and S. Pal, Phys. Rev. C 72, 064901 (2005), eprint nucl-th/0411110.
  * Bzdak and Ma (2014) A. Bzdak and G.-L. Ma, Phys. Rev. Lett. 113, 252301 (2014), eprint 1406.2804.
  * Orjuela Koop et al. (2015) J. D. Orjuela Koop, A. Adare, D. McGlinchey, and J. L. Nagle, Phys. Rev. C 92, 054903 (2015), eprint 1501.06880.
  * Ma and Bzdak (2016) G.-L. Ma and A. Bzdak, Nucl. Phys. A 956, 745 (2016).
  * He et al. (2016) L. He, T. Edmonds, Z.-W. Lin, F. Liu, D. Molnar, and F. Wang, Phys. Lett. B 753, 506 (2016), eprint 1502.05572.
  * Lin et al. (2016) Z.-W. Lin, L. He, T. Edmonds, F. Liu, D. Molnar, and F. Wang, Nucl. Phys. A 956, 316 (2016), eprint 1512.06465.
  * Ma et al. (2021) L. Ma, G.-L. Ma, and Y.-G. Ma, Phys. Rev. C 103, 014908 (2021), eprint 2102.01872.
  * Wang and Ma (2022) H.-S. Wang and G.-L. Ma, Phys. Rev. C 106, 064907 (2022), eprint 2208.06854.
  * Aad et al. (2013b) G. Aad et al. (ATLAS), JHEP 11, 183 (2013b), eprint 1305.2942.
  * Adams et al. (2006) J. Adams et al. (STAR), Phys. Rev. D 74, 032006 (2006), eprint nucl-ex/0606028.
  * Abelev et al. (2014) B. B. Abelev et al. (ALICE), Phys. Lett. B 728, 25 (2014), eprint 1307.6796.
  * Adam et al. (2015) J. Adam et al. (ALICE), Phys. Rev. C 91, 064905 (2015), eprint 1412.6828.
  * Acharya et al. (2020) S. Acharya et al. (ALICE), Eur. Phys. J. C 80, 693 (2020), eprint 2003.02394.
  * Abelev et al. (2009) B. I. Abelev et al. (STAR), Phys. Rev. C 79, 034909 (2009), eprint 0808.2041.
  * Abelev et al. (2013b) B. B. Abelev et al. (ALICE), Phys. Lett. B 727, 371 (2013b), eprint 1307.1094.
  * Adcox et al. (2004) K. Adcox et al. (PHENIX), Phys. Rev. C 69, 024904 (2004), eprint nucl-ex/0307010.
  * Pang et al. (2018) L.-G. Pang, K. Zhou, N. Su, H. Petersen, H. Stöcker, and X.-N. Wang, Nature Commun. 9, 210 (2018), eprint 1612.04262.
  * Zhou et al. (2019) K. Zhou, G. Endrődi, L.-G. Pang, and H. Stöcker, Phys. Rev. D 100, 011501 (2019), eprint 1810.12879.
  * Graczykowski et al. (2022) L. K. Graczykowski, M. Jakubowska, K. R. Deja, and M. Kabus (ALICE), JINST 17, C07016 (2022), eprint 2204.06900.
  * ZHOU et al. (2022) M. ZHOU, Y. LUO, and H. SONG, Sci. China Phys. Mech. Astron. 52, 252002 (2022).
  * WANG et al. (2022) L. WANG, L. PANG, and K. ZHOU, Sci. China Phys. Mech. Astron. 52, 252003 (2022).
  * HE et al. (2022) W. HE, J. HE, R. Wang, and Y. MA, Sci. China Phys. Mech. Astron. 52, 252004 (2022).
  * DU et al. (2022) Y.-L. DU, D. PABLOS, and K. Tywoniuk, Sci. China Phys. Mech. Astron. 52, 252017 (2022).
  * LI et al. (2023) F. LI, L. PANG, and X. WANG, Nucl. Sci. Tech. 46 (2023).
  * Qi et al. (2016) C. R. Qi, H. Su, K. Mo, and L. J. Guibas, arXiv e-prints arXiv:1612.00593 (2016), eprint 1612.00593.
  * Steinheimer et al. (2019) J. Steinheimer, L. Pang, K. Zhou, V. Koch, J. Randrup, and H. Stoecker, JHEP 12, 122 (2019), eprint 1906.06562.
  * Omana Kuttan et al. (2020a) M. Omana Kuttan, J. Steinheimer, K. Zhou, A. Redelbach, and H. Stoecker, Phys. Lett. B 811, 135872 (2020a), eprint 2009.01584.
  * Omana Kuttan et al. (2021) M. Omana Kuttan, J. Steinheimer, K. Zhou, A. Redelbach, and H. Stoecker, Particles 4, 47 (2021).
  * Omana Kuttan et al. (2020b) M. Omana Kuttan, K. Zhou, J. Steinheimer, A. Redelbach, and H. Stoecker, JHEP 21, 184 (2020b), eprint 2107.05590.
  * Huang et al. (2022) Y. Huang, L.-G. Pang, X. Luo, and X.-N. Wang, Phys. Lett. B 827, 137001 (2022), eprint 2107.11828.
  * Gyulassy and Wang (1994) M. Gyulassy and X.-N. Wang, Comput. Phys. Commun. 83, 307 (1994), eprint nucl-th/9502021.
  * Lin and Ko (2002) Z.-w. Lin and C. M. Ko, Phys. Rev. C 65, 034904 (2002), eprint nucl-th/0108039.
  * Zhang (1998) B. Zhang, Comput. Phys. Commun. 109, 193 (1998), eprint nucl-th/9709009.
  * Lin (2014) Z.-W. Lin, Phys. Rev. C 90, 014904 (2014), eprint 1403.6321.
  * Ma and Lin (2016) G.-L. Ma and Z.-W. Lin, Phys. Rev. C 93, 054911 (2016), eprint 1601.08160.
  * He and Lin (2017) Y. He and Z.-W. Lin, Phys. Rev. C 96, 014910 (2017), eprint 1703.02673.
  * Lin and Zheng (2021) Z.-W. Lin and L. Zheng, Nucl. Sci. Tech. 32, 113 (2021), eprint 2110.02989.
  * Li and Ko (1995) B.-A. Li and C. M. Ko, Phys. Rev. C 52, 2037 (1995), eprint nucl-th/9505016.
  * Voloshin and Zhang (1996) S. Voloshin and Y. Zhang, Z. Phys. C 70, 665 (1996), eprint hep-ph/9407282.
  * Leonidov et al. (1997) A. Leonidov, M. Nardi, and H. Satz, Z. Phys. C 74, 535 (1997).
  * Adler et al. (2006) S. S. Adler et al. (PHENIX), Phys. Rev. Lett. 97, 052301 (2006), eprint nucl-ex/0507004.
  * Zhao et al. (2023) X.-L. Zhao, Z.-W. Lin, L. Zheng, and G.-L. Ma, Phys. Lett. B 839, 137799 (2023), eprint 2112.01232.


# Tensors Fitting Perfectly

Adam Paszke
Google Research
<EMAIL_ADDRESS>Brennan Saeta
Google Research
<EMAIL_ADDRESS>
###### Abstract

Multidimensional arrays (NDArrays) are a central abstraction in modern
scientific computing environments. Unfortunately, they can make reasoning
about programs harder as the number of different array shapes used in an
execution of a program is usually very large, and they rarely appear
explicitly in program text. To make things worse, many operators make implicit
assumptions about the shapes of their inputs: array addition is commonly
enriched with broadcasting semantics, while matrix multiplication assumes that
the lengths of contracted dimensions are equal. Because precise reasoning
about shapes is crucial to write correct programs using NDArrays, and because
shapes are often hard to infer from a quick glance at the program, we
developed _Tensors Fitting Perfectly_ , a static analysis tool that reasons
about NDArray shapes in Swift for TensorFlow programs by synthesizing a set of
shape constraints from an abstract interpretation of the program. It can both
(1) check for possible inconsistencies, and (2) provide direct insights about
the shapes of intermediate values appearing in the program, including via a
mechanism called _shape holes_. The static analysis works in concert with
optional runtime assertions to improve the productivity of program authors.

##  1 Introduction

Numerical computing has been completely transformed by the concept of
multidimensional arrays and a significant fraction of modern scientific
computing workloads are written as compositions of multidimensional array
operations, instead of explicit, nested loops operating on scalars. There are
two key reasons for why this abstraction ended up being so useful for
efficient execution of scientific simulations. First, even though most of
those programs are written in an imperative language with sequential
semantics, most array operators are deeply parallel in nature, and so such
programs easily map to modern vectorized hardware which is the main driving
force of performance improvements since the end of Dennard scaling. Second,
because the number of array operators invoked is usually orders of magnitude
smaller than the number of corresponding scalar operations that need to be
performed, the whole computation can be expressed in a relatively inefficient
high-level language, enabling rapid development and iteration. This paradigm,
pioneered by APL [13], has fuelled major breakthroughs in science, and led to
the development of numerous software packages with similar functionality, such
as Matlab or NumPy [17]. More recently libraries like Theano [21], TensorFlow
[1], PyTorch [18] and JAX [4] have extended the NDArray abstraction to allow
seamless execution of array operators on accelerators like GPUs, and
efficiently perform automatic differentiation.

While this paradigm has been very fruitful, one of its biggest unsolved pain
points is shape mismatch errors. Every multidimensional array is associated
with a finite sequence of natural numbers which represent the domain of
indices valid for accessing its elements. For example, $\mathbb{R}^{n\times
m}$ matrices would usually be represented by arrays of shape [n, m] (although
the ordering of those dimensions depends on the convention of the libraries
and user code). The length of the shape sequence is often called the rank. If
one wants to take a matrix product A.dot(B), then they need to ensure that
both A and B are rank 2, and that the equality A.shape[1] == B.shape[0] holds.
In most widely used systems this condition is checked solely at run time,
meaning that such errors often go unnoticed until such an operation is
invoked. Because many of these NDArray-based programs run for hours or days
before completing (e.g. evaluating model performance on a validation or test
dataset, or pre-training on a large unsupervised dataset before fine-tuning
for a supervised task), iteratively running the program to debug shape errors
can be costly and unproductive.

Scientific computing has greatly benefited from the rise of the open source
culture, but unfortunately the difficulty of reasoning about shapes manifests
itself especially clearly when analyzing published code. Most projects treat
shapes as irrelevant metadata, and shape contracts of individual user or
library functions are left completely implicit. However, because the semantics
of many operations depend on the shapes, having no annotations makes those
programs difficult to understand, modify, and maintain, reducing the benefits
of free code sharing. Additionally, when shape mismatches are encountered at
runtime, the errors usually originate from the lowest levels of abstraction,
often with little context. In practice, NDArray program authors inevitably
need to understand the implementations of the abstractions they consume,
reducing the value of the abstraction itself.

One partial solution to that issue, taken by some authors, is to defensively
annotate the programs with comments signifying the shape of intermediate
variables. While this makes it much easier for readers to understand what the
program is actually doing, such documentation is not checked and can easily
get out of date. Because the shape is just a list of numbers, it seems natural
to ask to have it checked by a computer. Unfortunately, no such systems are
used today (although some early prototypes are under development [16]).

Most NDArray libraries support implicit broadcasting, or dynamically
increasing the rank of an NDArray by repeating its contents in order to
satisfy the shape constraints of a given operation. For example, adding a
vector v of size $\mathbb{R}^{n}$ to a matrix m of size $\mathbb{R}^{n\times
m}$ results in v being implicitly tiled $m$ times to form a
$\mathbb{R}^{n\times m}$ matrix. Broadcasting improves performance by avoiding
materializing copies of v in memory, and implicit broadcasting allows users to
write rank-polymorphic code. Unfortunately, implicit broadcasting makes
reasoning about shapes and their resulting errors substantially more difficult
and sometimes can even mask them, resulting in a silent divergence from the
desired program semantics.

Another possible solution to the shape mismatch problem would be to encode all
of the shape constraints in a type system, such that the type checking
procedure would additionally prove that no shape errors can occur at run time.
Unfortunately, while type systems that are expressive enough for that exist
(e.g. in Haskell [8]), their corresponding type checking problems are
undecidable and still require a fair amount of user supervision to have their
program accepted, even if it is correct. On the other hand, type systems used
in the most popular imperative languages are not up to the task, and cannot
deal with typing of common operations with sophisticated constraints (e.g.
functions that would accept inputs of arbitrary ranks, but with some
constraints on the trailing dimensions).

⬇

// A sample program entry point to check.

let m = Model()

let batchSize = 12

let input = Tensor(zeros: (12, 32, 32, 3))

let output = m(input)

|-> [batchSize, 10]

\par// A simple model used in the program.

//

// A couple shape assertions are

// sprinkled throughout the model.

struct Model {

let conv = Conv2D(

filterShape: (2, 2, 3, 5),

strides: (2, 2))

let dense = Dense(

inputSize: ____,

outputSize: 10)

\parfunc callAsFunction(

_ input: Tensor

) -> Tensor {

let batchSize = input.shape[0]

let c = relu(conv(input))

let m = maxPool2D(c, stride: (2, 2))

|-> [batchSize, 16, 16, 5] // !!!!

let dIn = flatten(m)

let dOut = dense(dIn)

|-> [batchSize, 10] // Optional.

return dOut

}

}

\parinfix operator |-> // Shape assert operator

func |-> (

_ a: Tensor,

_ b: TensorShape

) -> Tensor {

assert(a.shape == b)

return a

}

\parextension Tensor {

var shape4d: (Int, Int, Int, Int) {

assert(self.rank == 4)

let shape = self.shape

return (

shape[0], shape[1],

shape[2], shape[3])

}

}

⬇

// Select operator implementations,

// including shape assertions.

\par// Computes the valid shape for the output

// of the maxPool2D and conv2D operations.

func validWindowShape(

_ input: Tensor,

kernelSize: (Int, Int),

stride: (Int, Int),

output: Int

) -> TensorShape {

let (iN, iH, iW, _) = input.shape4d

return [

iN,

(iH - kernelSize.0) / stride.0 + 1,

(iW - kernelSize.1) / stride.1 + 1,

outputs

]

}

\parfunc maxPool2D(

_ input: Tensor,

kernelSize: (Int, Int),

stride: (Int, Int) = (2, 2)

) -> Tensor {

let result = TF.maxPool2D(

input, kernelSize, stride)

// Shape assertion without shape

// assertion operator.

assert(result.shape == validWindowShape(

input,

kernelSize: kernelSize,

stride: stride,

output: input.shape[3]) // Channels.

return result

}

\parfunc conv2D(

_ input: Tensor,

_ weight: Tensor

) -> Tensor {

let (_, iH, iW, iC) = input.shape4d

let (kH, kW, iF, oF) = weight.shape4d

assert(bF == oF); assert(iC == iF)

assert(iH >= kH); assert(iW >= kW)

let result = TF.conv2D(input, weight)

|-> validWindowShape(

input, kernelSize: (kH, kW),

stride: (1, 1), output: oF)

return result

}

Figure 1:  A sample program instantiating a neural network and evaluating the
network at a given input, along with select implementations of library
operations and their shape assertions. The shape assertion operator (|->) is
defined in terms of the assert primitive. maxPool2D and conv2D are examples of
sophisticated shape constraints inherent in neural networks. The static
analysis tool highlights a shape error on the line annotated with !!!!; the
correct shape for m is: [batchSize, 8, 8, 5]. Finally, the static analysis
tool informs the user that the shape hole (____) should be replaced with 320.

In this article, we propose to address the problem with a combination of a
static analysis tool coupled with more precise optional run-time shape
checking. Instead of trying to prove that the user program is correct, our
static analysis tool searches for provable failures, and will raise an error
if any of those are found. In this way, while the checking procedure is not
complete and can miss some issues, we can provide a low-noise signal. Run-time
assertions complement the static analysis to reason about shapes more
accurately, but in case their costs are prohibitive for any application, most
programming languages (Swift included) allow skipping their evaluation
entirely through a compiler or interpreter flag.

Apart from checking, our tool is also able to provide the user with some hints
about the shape semantics of the program through functionality like shape
holes that will be discussed in Section 3.8. Future work could integrate
functionality into the popular editors to interactively assist the researchers
and programmers in their daily tasks.

##  2 High level overview

This article is a description of Tensors Fitting Perfectly, an open-source
tool111https://github.com/google-research/swift-tfp for finding shape bugs in
Swift programs developed as part of the Swift for TensorFlow project [19]. It
takes a list of Swift files as input and passes those on to the Swift
compiler, which is asked to lower it to the Swift Intermediate Language (SIL)
instead of producing an executable or a library. SIL is a static single
assignment (SSA) based intermediate representation (IR) of the user program.
One interesting note is that SIL does not make use of phi-instructions, but
instead allows each block to take a list of arguments. All jumps to that block
have to specify the values that those arguments are supposed to take. SIL
preserves the full power of the Swift type system, and allows for storing most
(e.g. compound) values in virtual registers. Additionally, while the supported
instruction set is quite complex, it already desugars and normalizes most of
the numerous language features of Swift into much simpler constructs; for
example all closures are already converted into top-level functions at this
point. This step also includes type checking of the program, so it lets us
ensure that we only process well-formed code, but unfortunately limits the use
in the context of incomplete programs which would be necessary to provide live
hints in text editors.

Each top level function in the analyzed SIL file is preprocessed to eliminate
loops (Section 3.2) from its control flow graph (CFG) and to remove the need
to analyze memory loads and stores (e.g. stack allocations are hoisted into
virtual registers). Each block appearing in the SIL representation is then
subjected to symbolic execution (Section 3.3), which attempts to recover the
assertions specified by the user, and lift them into the raw form of the
constraint language used later. Initially the representation includes function
calls which have to be resolved to produce the canonical form on which
checking can be performed. After processing each function, we produce a
function summary which includes the set of constraints (including calls to
other functions) as well as the expressions corresponding to the arguments and
returned values. Note that the produced summary is not of bounded size. In
general, the inputs to the verification procedure may grow even exponentially
together with the size of the program (due to inlining), but the pathological
examples generally do not resemble real world programs.

Once an entry point and a path condition is chosen by the verification
procedure (Section 3.5), function calls are resolved by recursively
instantiating the summaries of called functions in a single constraint system
(with appropriate substitutions applied). The system is then simplified (some
operations are expensive to represent in the language of the solver),
translated into the solver’s intermediate representation (effectively a
logical formula), and checked. If the solver reports that the system is
unsatisfiable, a failure is reported to the user.

##  3 Implementation details

In this section, we discuss the inner workings of the tool in more detail. We
start with an overview of all the steps that prepare the program for symbolic
execution, and then discuss the algorithm that decides whether it should be
accepted or not. Finally, since the goal is to provide useful insights to our
users, we comment on how to present verification failures (i.e. shape errors)
in a way that makes them more approachable.

###  3.1 Shape Assertions

The static analysis tool is configured with special knowledge of two symbols:
the NDArray type (e.g. Swift for TensorFlow’s Tensor type), and the
corresponding shape type (e.g. TensorShape). Standard assertions are sprinkled
throughout the implementation of Tensor methods which form the basis for how
the tool reasons about the user’s code. Importantly, the tool is not given any
special knowledge of the semantics of the operations on NDArray; instead
everything is built up from these assertions within the NDArray
implementation.

###  3.2 Making the control flow graph loop-free

One of the important decisions that have to be made when designing a
verification tool based on SMT solvers is how to handle loops or recursion.
Although loops do not play a crucial role in typical NDArray programs, their
presence should not prevent the tool from applying the facts derived about the
program before the loop to the continuation after the loop finishes.

Because there is no way to express fixed points in standard propositional and
first-order logic [2], there are a number of techniques to model back edges in
the control flow graph. One common approach is to rewrite each loop appearing
in the program as a deeply nested tower of conditional statements, effectively
assuming that all loops have their trip counts bounded by the unrolling
factor. However, if one considers all trip counts up to the unrolling factor
as possible program paths, then this approach would additionally assume that
the loops might execute as many times as this threshold, which might cause
e.g. some shapes to appear to become negative on those paths if the loop
shrinks them. Because we have tried to err on the side of soundness (i.e.
trying to not report errors for programs that can execute successfully) when
designing the tool, we did not take that approach.

Instead, we replace each loop with a single conditional statement,
representing a choice between running the loop at least once or skipping it
entirely. If the loop executes at least once then, assuming that the loop
terminates, we know that there exists the first and the last iteration, and so
we repeat the loop body twice. The only question is what happens with loop
carried variables, and the answer is that the first instance uses the regular
loop inputs that would get passed in to the first iteration, while the second
instance gets “fresh” inputs in the sense that they get materialized out of
thin air and we do not know of any dependencies between them and any other
program variable. Note that after this transformation, although the program
can no longer be executed (as we would need a way to actually supply the
values for the second iteration), it can still be analyzed. What this means is
that we effectively reject to analyze how the shapes change throughout the
loop, but only want to use the assertions from inside of the loop body to be
cross-checked with the ones that appear both after and before the loop. It
also has the benefit of reducing the number of possible program paths that
have to be considered.

    // Swift source.

    func loopingFn(
        k: Int,
        input: Tensor
    ) -> Tensor {
      var x = input
      x = preOp(x)
      for i in 0..<k {
        x = loopOp(x)
      }
      return postOp(x)
    }

    // SIL-like SSA repr.

    func loopingFn {
    bb0(%
      br bb1(loopCount, %
    bb1(%
      cond_br %
    bb2:
      br bb1(%
    bb3:
      return %
    }

    // Function summary.

    loopingFn(s1, s7):
      // Before loop
      preOp(s1, s2)
      s4a == s2
      loopOp(s4a, s6a)
      // After loop
      loopOp(s6b, s4b)
      postOp(s4b, s7)

A simple function containing a loop and tensor operations, a SIL-like SSA
representation of the function that is used as input to the TFP tool, and a
representation of the resulting shape constraint summary inferred from the
function’s control flow graph. In the function summary, the calls to sub-
procedures should be understood to calls to their summaries, which constrain
the shapes of input and output tensors.

###  3.3 Symbolic execution

Now that the control flow graph is loop-free the program is executed
symbolically to extract the assertions appearing inside it. We initialize
arguments of every block to a variable of their corresponding type, if it is a
supported type, and an additional variable is created to represent the value
returned from the function. Note that there are no variables for compound
types, so compound arguments are associated with a compound expression
containing variables. SIL does not have a concept of methods (they are lowered
to curried top-level functions), so we treat compound data types like structs
as isomorphic to tuples containing their members.

$\displaystyle\text{Integer expressions}\quad n::$ $\displaystyle=l\mid v$
Literals / variables
$\displaystyle\mid\texttt{\\_\\_\\_\\_(}\texttt{loc}\texttt{)}$ Shape hole
(loc is a source location) $\displaystyle\mid\texttt{rank(}s\texttt{)}$ Rank
(number of dimensions) $\displaystyle\mid s[c]$ Dimension size
$\displaystyle\mid n+n\mid n-n\mid n\cdot n\mid n/n$ Arithmetic (division
rounds down) $\displaystyle\text{Boolean expressions}\quad b::$
$\displaystyle=\texttt{true}\mid\texttt{false}\mid v$ Literals / variables
$\displaystyle\mid\neg b$ Negation $\displaystyle\mid b\land\ldots\land b\mid
b\lor\ldots\lor b$ Conjunction / disjunction $\displaystyle\mid n=n\mid
s=s\mid b=b$ Primitive type equality $\displaystyle\mid n>n\mid n\geq n\mid
n<n\mid n\leq n$ Relations between naturals $\displaystyle\text{Shape
expressions}\quad s::$ $\displaystyle=v$ Variables
$\displaystyle\mid[n,\ldots,n]$ Literals
$\displaystyle\mid\texttt{broadcast(}s,s\texttt{)}$ Broadcasting operation
$\displaystyle\text{Compound expressions}\quad c::$
$\displaystyle=(e,\ldots,e)$ Tuple $\displaystyle\text{Expressions}\quad e::$
$\displaystyle=\texttt{integer(}n\texttt{)}\mid\texttt{boolean(}b\texttt{)}$
$\displaystyle\mid\texttt{shape(}s\texttt{)}\mid\texttt{compound(}c\texttt{)}$
Figure 2: Language of shape constraints

In the following, a path condition for a given simple block should be
understood as a logical formula that has to be satisfiable if the block is to
be reachable. Intuitively, when we encounter a program that corresponds to an
if statement, the block corresponding to one of its branches will have the
checked expression in its path condition, while the other one will have its
negation. This way we ensure that we never attempt to cross-verify assertions
made in one branch with those made in the other one, since those can easily be
contradictory, while the actual program will only ever be subject to one of
them.

We begin the execution by setting the path condition of the function entry
block to true, and conditions of all other blocks to false. When processing a
block, all instructions in its body are executed symbolically over a limited
set of abstract values which is a superset of the constraint language
supported by TFP (Figure 2) that additionally includes e.g. function pointers
and their partial applications. Any calls to the assert built-in function are
added to the set of constraints, guarded by the block’s path condition. Once a
block-terminating instruction is reached, the set of constraints is extended
with equations between the symbolic values of the terminator operands and the
arguments of the successor blocks (if the terminator is a jump), or with the
variable representing the function result. Finally, the successor blocks have
their path condition extended by taking a disjunction of their current
condition and a conjunction of the current block’s path condition with the
condition derived from the terminator if possible (e.g. if it is a conditional
jump) or true otherwise (e.g. if it is an unconditional jump).

###  3.4 Representing shape formulas in UFNIA

To verify specifications represented in the constraint language (Figure 2), we
translate those into formulas over the UFNIA logic (uninterpreted functions
and non-linear integer arithmetic). All shape variables are represented as
uninterpreted functions from integers to integers and a single integer
encoding its rank. This representation is motivated by the fact that we cannot
constrain the domain to a finite subset of natural numbers a priori, because
the rank does not have to be known statically. However, at each point where
this function is evaluated, an assertion ensuring that its argument falls
between 0 and the rank is added, guaranteeing that the formulas only reason
about the fragment of the domain bounded by the corresponding rank variable.
The full translation procedure is described in Figure 3.

Despite the fact that there is no decision procedure for the UFNIA logic and
that undecidable problems can easily be encoded in the language of TFP, the Z3
solver we have used in our implementation has not failed to resolve any of the
representative examples we have tried so far. One practical note is that
representing shape equality is quite expensive, as it requires the use of
universal quantification to express function equality. It is however one of
the most fundamental operations in our constraint language, and because of
that we preprocess the constraints before the UFNIA translation by trying to
eliminate as many equalities as we can (by eliminating variables that can be
substituted without affecting correctness). With this simple preprocessing
step, the verification process is very quick, but one can easily get solver
timeouts if it is skipped.

$\displaystyle\mathcal{I}\llbracket l\rrbracket$ $\displaystyle=l$
$\displaystyle\mathcal{I}\llbracket v\rrbracket$ $\displaystyle=v$
$\displaystyle\mathcal{I}\llbracket\texttt{\\_\\_\\_\\_(}\texttt{loc}\texttt{)}\rrbracket$
$\displaystyle=v_{\texttt{loc}}~{}\text{(a variable name unique to program
location $\texttt{loc}$ where the hole appeared)}$
$\displaystyle\mathcal{I}\llbracket\texttt{rank(}s\texttt{)}\rrbracket$
$\displaystyle=\mathcal{S}\llbracket s\rrbracket\texttt{.rank}$
$\displaystyle\mathcal{I}\llbracket s[c]\rrbracket$
$\displaystyle=\begin{cases}\mathcal{S}\llbracket
s\rrbracket\texttt{.shape(}c\texttt{)}~{}\textbf{assuming}~{}0<-c\leq\mathcal{S}\llbracket
s\rrbracket\texttt{.rank}&\text{if}~{}c<0\\\ \mathcal{S}\llbracket
s\rrbracket\texttt{.shape(}\mathcal{S}\llbracket
s\rrbracket\texttt{.rank}-c\texttt{)}~{}\textbf{assuming}~{}0\leq
c<\mathcal{S}\llbracket s\rrbracket\texttt{.rank}&\text{if}~{}c\geq 0\\\
\end{cases}$ $\displaystyle\mathcal{I}\llbracket n_{1}\odot n_{2}\rrbracket$
$\displaystyle=\mathcal{I}\llbracket n_{1}\rrbracket\odot\mathcal{I}\llbracket
n_{2}\rrbracket~{}\text{(where $\odot\in\\{+,-,\cdot,/\\}$)}$
$\displaystyle\mathcal{B}\llbracket\texttt{true}\rrbracket$
$\displaystyle=\top$
$\displaystyle\mathcal{B}\llbracket\texttt{false}\rrbracket$
$\displaystyle=\bot$ $\displaystyle\mathcal{B}\llbracket v\rrbracket$
$\displaystyle=v$ $\displaystyle\mathcal{B}\llbracket\neg b\rrbracket$
$\displaystyle=\neg\mathcal{B}\llbracket b\rrbracket$
$\displaystyle\mathcal{B}\llbracket b_{1}\diamond\ldots\diamond
b_{k}\rrbracket$ $\displaystyle=\mathcal{B}\llbracket
b_{1}\rrbracket\diamond\ldots\diamond\mathcal{B}\llbracket
b_{k}\rrbracket~{}\text{(where $\diamond\in\\{\land,\lor\\}$)}$
$\displaystyle\mathcal{B}\llbracket n_{1}=n_{2}\rrbracket$
$\displaystyle=\mathcal{I}\llbracket n_{1}\rrbracket=\mathcal{I}\llbracket
n_{2}\rrbracket$ $\displaystyle\mathcal{B}\llbracket b_{1}=b_{2}\rrbracket$
$\displaystyle=\mathcal{B}\llbracket b_{1}\rrbracket=\mathcal{B}\llbracket
b_{2}\rrbracket$ $\displaystyle\mathcal{B}\llbracket n_{1}\diamond
n_{2}\rrbracket$ $\displaystyle=\mathcal{I}\llbracket
n_{1}\rrbracket\diamond\mathcal{I}\llbracket n_{2}\rrbracket~{}\text{(where
$\diamond\in\\{>,\geq,<,\leq\\}$)}$ $\displaystyle\mathcal{B}\llbracket
s_{1}=s2\rrbracket$ $\displaystyle=\left(\mathcal{S}\llbracket
s_{1}\rrbracket\texttt{.rank}=\mathcal{S}\llbracket
s_{2}\rrbracket\texttt{.rank}\right)\land\left(\forall_{i}~{}\mathcal{S}\llbracket
s_{1}\rrbracket\texttt{.shape(}i\texttt{)}=\mathcal{S}\llbracket
s_{2}\rrbracket\texttt{.shape(}i\texttt{)}\right)$
$\displaystyle\mathcal{S}\llbracket v\rrbracket$
$\displaystyle=(v,v_{\text{rank}})$
$\displaystyle\mathcal{S}\llbracket[n_{1},\ldots,n_{k}]\rrbracket$
$\displaystyle=(v_{f},v_{f_{\text{rank}}})~{}\textbf{assuming}~{}\mathcal{S}\llbracket
v_{f}\rrbracket\texttt{.rank}=k\land\bigwedge\mathcal{S}\llbracket
v_{f}\rrbracket\texttt{.shape(}i\texttt{)}=\mathcal{I}\llbracket
n_{i}\rrbracket$
$\displaystyle\mathcal{S}\llbracket\texttt{broadcast(}s_{1},s_{2}\texttt{)}\rrbracket$
$\displaystyle=(v_{f},v_{f_{\text{rank}}})~{}\textbf{assuming}~{}\mathcal{S}\llbracket
v_{f}\rrbracket\texttt{.rank}=\textbf{max}\left(\mathcal{S}\llbracket
s_{1}\rrbracket\texttt{.rank},~{}\mathcal{S}\llbracket
s_{2}\rrbracket\texttt{.rank}\right)$ $\displaystyle\hskip
92.70018pt\forall_{i}~{}\mathcal{S}\llbracket
v_{f}\rrbracket\texttt{.shape(}i\texttt{)}=\textbf{max}\left(\mathcal{S}\llbracket
s_{1}\rrbracket\texttt{.shape(}i\texttt{)},~{}\mathcal{S}\llbracket
s_{2}\rrbracket\texttt{.shape(}i\texttt{)}\right)$ $\displaystyle\hskip
92.70018pt\forall_{i}~{}\mathcal{S}\llbracket
s_{1}\rrbracket\texttt{.shape(}i\texttt{)}=1\lor\mathcal{S}\llbracket
s_{2}\rrbracket\texttt{.shape(}i\texttt{)}=1$ $\displaystyle\hskip
189.45035pt\lor\mathcal{S}\llbracket
s_{1}\rrbracket\texttt{.shape(}i\texttt{)}=\mathcal{S}\llbracket
s_{2}\rrbracket\texttt{.shape(}i\texttt{)}$ Figure 3: Translation of shape
constraints into the UFNIA logic. $v_{f}$ denotes a fresh variable that is
created each time a rule is used. Additionally, each time a rule containing
$~{}\textbf{assuming}~{}\phi$ is evaluated, the formula $\phi$ is added to the
set of assumptions. The result is a conjunction of the output formula and all
assumptions made during the translation. The result of
$\mathcal{S}\llbracket\cdot\rrbracket$ is assumed to be a named pair, with
first component named dims and second component named rank. Also,
$\mathcal{S}\llbracket x\rrbracket\texttt{.shape(}i\texttt{)}$ is syntax sugar
for $\mathcal{S}\llbracket x\rrbracket.\texttt{dims}(\mathcal{S}\llbracket
x\rrbracket\texttt{.rank}-i-1)$.

###  3.5 Checking procedure

To verify the program, the algorithm first gathers all possible path
conditions appearing in the program constraints. For each such condition, we
first check that this path is actually feasible, i.e. that the blocks labeled
by this condition are actually reachable. This is performed by asserting all
constraints except those that are guarded by the current condition along with
the condition. The intuition here is that we want to verify whether all the
constraints that are in blocks which are less constrained than this one
actually allow for its execution. If the path feasibility check succeeds, we
additionally assert the conditions that we have omitted in the check, and run
the analysis.

To see why verifying satisfiability of the path condition alone is not
sufficient consider the following program:

⬇

func g(_ y: Int) {

if (y == 2) {

…

}

}

\parfunc f(_ x: Int) {

if (x == 1) {

g(x)

}

}

After processing, the system of constraints derived from f would look
approximately like this: $(x=1\Rightarrow y=x)\land(x=1\land y=2\Rightarrow
true)$. While $x=1\land y=2$ is certainly a satisfiable formula on its own, it
is not a good assumption to make in this program given that it implies the
path condition of the equality $y=x$.

A careful reader will notice that the above algorithm has a failure mode which
is that if blocks that have weaker path conditions contain a shape
contradiction then it will be considered as a reason to reject the strongly
constrained path as unfeasible. This however should not be an issue, as it
means that the program does in fact contain an error. Because we iterate over
all possible path conditions it will be reported when we reach the weakest
possible assumption that is necessary to prove it.

Instead of verifying all possible conditions one could also refine the
algorithm such that it first considers the most restricted conditions, and
only descends to those implied by it if a failure is found or the path is
found to be infeasible. This however is left for a future improvement.

###  3.6 Error reporting

One of the most significant challenges facing verification tools that depend
on SMT solvers is providing good error messages. In our case, it is very hard
to analyze the exact cause of failure, and the proofs produced by Z3 on small
programs can easily grow to thousands of operations. Providing those
explanations is hard, but making it easy to understand the failures for users
is crucial for success of a static analysis tool. Hence, to give users some
insights, we extract an unsat core from the solver (i.e. a subset of the input
assertions that lead to a contradiction), and apply a simplification procedure
to those. It looks for expressions of the form v = ..., and tries to eliminate
the variable v from the constraint system by inlining the other side of the
equality into other constraints and simplifying those. As it turns out, this
method works quite well in practice.

[h!]

⬇

func matmul(_ x: Tensor<Float>, _ y: Tensor<Float>) -> Tensor<Float> {

assert(x.shape[1] == y.shape[0])

let r = TensorFlow.matmul(x, y)

assert(r.shape == [x.shape[0], y.shape[1]])

return r

}

\parlet x = randn([20, 10])

let y = randn([30, 10])

let z = matmul(x, y)

An incorrect use of the matrix multiplication function.

[h!]

    In main():
    Something doesn’t fit!
      - 10 = 30
          Asserted at tmp.swift:2
                | func matmul(_ x: Tensor<Float>, _ y: Tensor<Float>) -> Tensor<Float> {
             2  |   assert(x.shape[1] == y.shape[0])
                |   let r = TensorFlow.matmul(x, y)

An error reported by the TFP tool when checking the program shown in Listing
3.6.

To see what kind of output can be provided with this method, consider a
program shown in Listing 3.6. Running TFP over it, would produce the error in
Listing 3.6, pointing to a single assertion that (after simplification) is
sufficient to prove the error. Note that it is also printed as 10 = 30, which
shows what kind of abstract values were deduced for both sides of this
equality.

In the future there are two potential improvements that we could consider.
Firstly, instead of using Z3 for all the solving, there’s an option of writing
a custom model checker that only defers to an SMT solver for arithmetic
reasoning, which would allow it to give more context when a failure occurs.
Secondly, one could try to synthesize summaries from Z3 unsat proofs (e.g.
match some patterns of most common shape errors), but that might not be robust
in practice.

###  3.7 Higher level abstractions

Although our system supports analysing un-annotated user programs, the more
assertions appear in the program, the better the error messages produced by
TFP become. Additionally, the tool can know whether the shape error is within
the implementation of a function or in the use of the function. As a result,
developing ergonomic syntax for expression assertions is an integral part of a
successful tool.

Because the tool derives shape analysis from standard assertions and not from
the semantics of the NDArray operations themselves, the entire system
gracefully extends to arbitrary user abstractions. For example, deep learning
libraries often vend a variety of Layer abstractions (e.g. a convolutional
layer). Library authors simply add standard asserts within their code.

In _Tensors Fitting Perfectly_ , we leverage Swift’s support for defining
custom operators to define the |-> operator which is syntactic sugar for shape
assertions. (See Listing 1 for the implementation.) It asserts the Tensor has
a given (symbolic) shape, and returns the Tensor itself, allowing for
ergonomic use in variable assignment. The following is an example in practice:

⬇

func runDigitDetection(

_ images: Tensor<Float>

) -> Tensor<Float> {

let batchSize = input.shape[0]

let output = network(input)

|-> [batchSize, 10] // 10 digits

let predictions = softmax(output)

|-> [batchSize]

return predictions

}

In addition to providing information about shapes statically, these assertions
are validated at runtime to help debug shape errors static analysis did not
catch. Note: the runtime checks are disabled when compiling with optimizations
to avoid introducing any overheads for high performance applications.

###  3.8 Shape holes

Finding shape errors can definitely make writing correct programs easier, but
that functionality alone does not really improve one’s understanding of the
shape semantics of the code. Ideally, one could imagine having an editing
environment in which e.g. hovering over a tensor expression or variable would
reveal its shape in relation to other variables present in scope, or some
globally defined constants (like batch size for example). Such a tight
integration is not yet supported by our tool, but we do provide a way to
synthesize constants in the user program based on the shape constraints which
we call shape holes.

More concretely, whenever an identifier ____ (of type Int) is present in the
user program, _Tensors Fitting Perfectly_ will treat it as a shape hole. As
shown in the constraint language specification (Figure 2) each hole is
associated with its program location. Later, during the translation to logical
formulas (Figure 3) each hole effectively gets assigned a unique variable
labelled by its program location. Assuming that there are no contradictions in
the extracted shape specification, we can retrieve the model from the SMT
solver, and see what the valuation of each such special variable is. Moreover,
once we see an example valuation we can query the solver for alternative
solutions to determine whether each such value is in fact unique, or to
generate a number of examples. Based on that, if the checking procedure
succeeds, TFP will display a message showcasing a few possible values for each
program hole, making it very easy to fill e.g. incomplete machine learning
model specifications automatically.

For example, consider this program where randn creates a new array with a
given size and all entries sampled i.i.d. from the unit Normal distribution
and matmul being the matrix multiplication function:

⬇

let x = randn([20, 10])

let y = randn([____, ____])

let z = matmul(x, y)

Matrix multiplication requires that the contracted dimensions (i.e. second
dimension of x and first dimension of y) be equal, while the second dimension
of y is completely unconstrained in this case. Running TFP over it would
produce the output in Listing 3.8.

    In main() -> ():
      - The hole at tmp.swift:2:19 has to be exactly 10
            |   let x = randn([20, 10])
         2  |   let y = randn([____, ____])
            |   let z = matmul(x, y)
      - Some example values that the hole at tmp.swift:2:25 might take are: 1, 2, 3
            |   let x = randn([20, 10])
         2  |   let y = randn([____, ____])
            |   let z = matmul(x, y)

Tool output guiding the user to values for shape holes that allow the program
to execute correctly. There are 2 shape holes—both on line 2—in this program.
The tool prints the valid value when only a single shape value satisfies all
the program’s constraints (the first shape hole), while suggesting a range of
valid values if the constraints do not force a single solution (the second
shape hole).  This confirms the insights stated in the previous paragraph and
indicates that the first dimension of y necessarily has to be of size 10, or a
shape error will occur otherwise. It is clear that this approach can be
extended further and used for providing functionality like tooltips for code
editors or synthesis of program constants that might be hard to compute by
hand (a good example here is e.g. the number of features after flattening of a
sequence of convolutional layers).

##  4 Related work

While the approach to shape checking taken in this paper is novel, the
technique of applying model checkers for bug finding has been known for a long
time. In fact, our tool works similarly to the sketch described in the seminal
work of Engler et al. [9]: we infer the specification from the user program,
and report errors if it is contradictory. Another tool that has adopted a
similar approach is PREfix [5]. However, both of those are mostly focused on
finding e.g. memory safety issues instead of reasoning about numerical
metadata like shapes. A good introduction to those methods can be found in an
excellent review by the authors of the Z3 SMT solver [3].

Much more academic work has been devoted to the task of proving the absence of
errors instead. If we do not limit ourselves to the shape checking problem
stated previously, techniques like abstract interpretation [7, 6] have been
quite successful. However, their reliance on the over-approximation of the set
of reachable states means that they can often report errors which are not
actually present in the user code. This does not work very well in the context
of shape verification, because one does not want to force all users into a
defensive mode where they have to write a lot of code devoted to checking
shape conditions. It has also been known to have trouble scaling to larger
programs, but successful reasoning about shapes without annotations crucially
depends on extracting knowledge from large fragments of the program via inter-
procedural analysis.

Annotation checkers [14, 10] also deserve a mention here, but those require
numerous extra directives from the user, which is very time-consuming in large
programs. Such a productivity hit certainly outweighs the benefits in many
research applications where most programs are run once and then modified or
discarded.

When one turns to shape verification, most tools built specifically for that
purpose are based on an attempt to design a decidable type system that can
effectively prove that no shape errors can occur at run time. Those are
usually explored in the context of functional programming, and some good
examples include languages like Remora [20], Dex [15] or Futhark [11].
Unfortunately, many operators used commonly in practice today require features
that would make the type system undecidable, meaning that all those approaches
cannot possibly scale to the way those programs are expressed today. This is
not to say that the current interfaces are perfect and a decidable dialect
cannot be made useful. It is only that those attempts cannot help the
thousands of people writing e.g. NumPy programs today.

Projects utilizing undecidable systems aided by solvers like Hasktorch [12]
deserve a mention too. They can usually assign types to most of the commonly
used operators, and the modern solvers are actually powerful enough to
discharge most constraints emitted by the type checking procedure. However,
their success critically depends on the interesting programs actually forming
a decidable subset of the type checking problem, which can be handled with
available solvers.

The approaches discussed above provide a lot of safety and can certainly be
beneficial in mission critical applications. However, many numerical programs
are written purely for research purposes, and do not need the full guarantee
of correctness, especially at a cost of the verification process restricting
the expressiveness of the language or significantly slowing down the
development speed. Additionally, all of the typing approaches have been
presented in functional programming languages which are not the tool of choice
for the vast majority of users today. Hence, applying those techniques without
significantly disrupting existing workflows seems very hard. This is why we
believe that bug finding approaches like the one proposed in this article are
an interesting alternative, as they can provide insights into the programs,
while avoiding the need for the user to spend too much time wrangling with the
tool.

##  5 Conclusion and future work

This article introduces _Tensors Fitting Perfectly_ , a static analysis tool
that attempts to infer tensor shape specification from optionally annotated
user programs, and find shape errors based on those. The error checking is
based on encoding the specification in the UFNIA logic, and using an SMT
solver to answer satisfiability queries.

This is unlike any prior work, where all shape checks are either solely
carried out at run time or are encoded as part of the type-checking problem,
leading to certain limitations and forcing users to change their workflows
significantly. This work presents a practical alternative that gracefully
combines a novel static representation of programs with optional dynamic
checks to form a pragmatic and practical system to improve user efficiency.

Apart from that, TFP provides the means to infer values of different shape
constants, making it possible for programmers to get interactive feedback
about the shapes without executing the code. All of this has been designed to
potentially work with partially incomplete or incorrect code (containing e.g.
syntax errors), such that in the future it can be used to power high-level
code intelligence tools and provide direct insights from within the editors.
Some other future plans for our work include: (a) Expanding support for more
language constructs (e.g. var declarations); (b) Adding support for
verification across Swift modules; and (c) Writing front-ends for other
languages (e.g. Python).

## References

  * [1] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng.  TensorFlow: Large-scale machine learning on heterogeneous systems, 2015\.  Software available from tensorflow.org.
  * [2] Alfred V Aho and Jeffrey D Ullman.  Universality of data retrieval languages.  In Proceedings of the 6th ACM SIGACT-SIGPLAN symposium on Principles of programming languages, pages 110–119. ACM, 1979.
  * [3] Nikolaj Bjørner and Leonardo de Moura.  Applications of SMT solvers to program verification.  2014\.
  * [4] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman-Milne.  JAX: composable transformations of Python+NumPy programs.  http://github.com/google/jax, 2018.
  * [5] William R Bush, Jonathan D Pincus, and David J Sielaff.  A static analyzer for finding dynamic programming errors.  Software: Practice and Experience, 30(7):775–802, 2000.
  * [6] Patrick Cousot and Radhia Cousot.  Introduction to abstract interpretation.  1998\.
  * [7] Patrick Cousot, Radhia Cousot, Jérôme Feret, Laurent Mauborgne, Antoine Miné, David Monniaux, and Xavier Rival.  The astrée analyzer.  In European Symposium on Programming, pages 21–30. Springer, 2005\.
  * [8] Richard A Eisenberg.  Dependent types in haskell: Theory and practice.  PhD thesis, University of Pennsylvania, 2016.
  * [9] Dawson Engler, David Yu Chen, Seth Hallem, Andy Chou, and Benjamin Chelf.  Bugs as deviant behavior: A general approach to inferring errors in systems code.  In ACM SIGOPS Operating Systems Review, volume 35, pages 57–72. ACM, 2001.
  * [10] Cormac Flanagan, Cormac Flanagan, K Rustan M Leino, Mark Lillibridge, Greg Nelson, James B Saxe, and Raymie Stata.  Extended static checking for java.  In ACM Sigplan Notices, volume 37, pages 234–245. ACM, 2002.
  * [11] Troels Henriksen, Niels GW Serup, Martin Elsman, Fritz Henglein, and Cosmin E Oancea.  Futhark: purely functional gpu-programming with nested parallelism and in-place array updates.  In ACM SIGPLAN Notices, volume 52, pages 556–571. ACM, 2017.
  * [12] Austin Huang, Junjie Hashimoto, Adam Paszke, Sam Stites, and Torsten Scholak.  Hasktorch.  https://github.com/hasktorch/hasktorch, 2019.
  * [13] Kenneth E. Iverson.  A Programming Language.  John Wiley & Sons, Inc., New York, NY, USA, 1962.
  * [14] Daniel Jackson.  Aspect: Detecting bugs with abstract dependences.  ACM Transactions on Software Engineering and Methodology (TOSEM), 4(2):109–145, 1995.
  * [15] Dougal Maclaurin, Radulm Alexey, Matthew J. Johnson, and Dimitrios Vytiniotis.  Dex: array programming with typed indices.  In NeurIPS Program Transformations Workshop, 2019.
  * [16] Nishant Sinha.  tsanley.  https://github.com/ofnote/tsanley, 2019.
  * [17] Travis Oliphant.  NumPy: A guide to NumPy.  USA: Trelgol Publishing, 2006.
  * [18] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In Advances in Neural Information Processing Systems, 2019.
  * [19] Brennan Saeta, Denys Shabalin, Marc Rasi, Brad Larson, Xihui Wu, Parker Schuh, Michelle Casbon, Daniel Zheng, Saleem Abdulrasool, Aleksandr Efremov, Dave Abrahams, Chris Lattner, and Richard Wei.  Swift for TensorFlow: A portable, flexible platform for deep learning.  MLSys, 2021.
  * [20] Justin Slepak, Olin Shivers, and Panagiotis Manolios.  An array-oriented language with static rank polymorphism.  In European Symposium on Programming Languages and Systems, pages 27–46. Springer, 2014.
  * [21] Theano Development Team.  Theano: A Python framework for fast computation of mathematical expressions.  arXiv e-prints, abs/1605.02688, May 2016.