This section presents experiments which show that the graph simplifications we propose to collapse standard Taylor mode are currently not applied by the \texttt{jit} compiler in JAX.

\begin{figure*}[!t]
  \centering

  % From https://tex.stackexchange.com/a/7318
  \newcolumntype{C}{ >{\centering\arraybackslash} m{0.11\textwidth} }
  \newcolumntype{D}{ >{\centering\arraybackslash} m{0.4\textwidth} }
  \begin{tabular}{CDD}
    & \textbf{Laplacian $(D=50)$}
    & \makecell{\textbf{Biharmonic $(D=5)$} \\ \textbf{(via nested Laplacians)}}
    \\
    \textbf{Exact}
    & \includegraphics{jet/exp/exp04_jax_benchmark/figures/architecture_tanh_mlp_768_768_512_512_1_device_cuda_dim_50_name_jax_laplacian_vary_batch_size.pdf}
    & \includegraphics{jet/exp/exp04_jax_benchmark/figures/architecture_tanh_mlp_768_768_512_512_1_device_cuda_dim_5_name_jax_bilaplacian_vary_batch_size.pdf}
    \\
    & $(N=2048)$
    & $(N=256)$
    \\
    \textbf{Stochastic}
    & \includegraphics{jet/exp/exp04_jax_benchmark/figures/architecture_tanh_mlp_768_768_512_512_1_batch_size_2048_device_cuda_dim_50_distribution_normal_name_jax_laplacian_vary_num_samples.pdf}
    & \includegraphics{jet/exp/exp04_jax_benchmark/figures/architecture_tanh_mlp_768_768_512_512_1_batch_size_256_device_cuda_dim_5_distribution_normal_name_jax_bilaplacian_vary_num_samples.pdf}
    \\
  \end{tabular}
  \captionof{figure}{\textbf{JAX's \texttt{jit} compiler does not apply our graph simplifications to standard Taylor mode.} Colors: \textcolor{tab-green}{Collapsed Taylor mode}, \textcolor{tab-orange}{standard Taylor mode}, and \textcolor{tab-blue}{nested first-order automatic differentiaion}, \textcolor{black!50!white}{opaque} memory consumptions are for non-differentiable computations.
    Results are on GPU and we use a $D \to 768 \to 768 \to 512 \to 512 \to 1$ MLP with tanh activations, varying the batch size.
    For each approach, we fit a line to the data and report the slope in \cref{tab:jax-benchmark} to quantify the relative speedup and memory reduction.
  }
  \label{fig:jax-benchmark}
\end{figure*}

\begin{figure*}[!t]
  \centering

  \captionof{table}{\textbf{JAX Benchmark from \cref{fig:jax-benchmark} in numbers.}
    We fit linear functions and report their slopes, \ie, how much runtime and memory increase when incrementing the batch size.
    % Our collapsed Taylor mode is up to two times faster than nested first-order autodiff, while using 80\% of memory in the differentiable, and 70\% in the non-differentiable, setting.
    All numbers are shown with two significant digits and bold values are best according to parenthesized values.}
  \label{tab:jax-benchmark}
  \vspace{1.5ex}
  % paths where the performances are stored
  \def\datapathJAXLaplacianExact{jet/exp/exp04_jax_benchmark/performance/architecture_tanh_mlp_768_768_512_512_1_device_cuda_dim_50_name_jax_laplacian_vary_batch_size}
  \def\datapathJAXBilaplacianExact{jet/exp/exp04_jax_benchmark/performance/architecture_tanh_mlp_768_768_512_512_1_device_cuda_dim_5_name_jax_bilaplacian_vary_batch_size}
  \def\datapathJAXLaplacianStochastic{jet/exp/exp04_jax_benchmark/performance/architecture_tanh_mlp_768_768_512_512_1_batch_size_2048_device_cuda_dim_50_distribution_normal_name_jax_laplacian_vary_num_samples}
  \def\datapathJAXBilaplacianStochastic{jet/exp/exp04_jax_benchmark/performance/architecture_tanh_mlp_768_768_512_512_1_batch_size_256_device_cuda_dim_5_distribution_normal_name_jax_bilaplacian_vary_num_samples}
  % configuration options for the \num command
  \sisetup{%
    % scientific-notation=true,%
    round-mode=figures,%
    round-precision=2,%
    detect-weight, % for bolding to work
    tight-spacing=true, % less space around \cdot
  }
  % temporarily overwrite the \num command to process nan's
  \let\origsiunitxnum\num
  % Redefine \num to check for non-numeric strings
  \renewcommand{\num}[1]{\IfStrEq{#1}{nan}{\text{n/a} }{\origsiunitxnum{#1}}}
  \begin{tabular}{ccc|ccc}
    \toprule
    \textbf{Mode}
    & \makecell{\textbf{Per-datum or} \\ \textbf{sample cost}}
    & \textbf{Implementation}
    & \textbf{Laplacian}
    & \makecell{\textbf{Biharmonic} \\ \!\!\textbf{(via nested Laplacians)}\!\!}
    \\
    \midrule
    \multirow{9}{*}{\textbf{Exact}}
    & \multirow{3}{*}{Time [ms]}
    & \textcolor{tab-blue}{Nested 1\textsuperscript{st}-order}
    & \input{\datapathJAXLaplacianExact/hessian_trace_best.txt}
    & \input{\datapathJAXBilaplacianExact/hessian_trace_best.txt}
    \\
    &
    & \textcolor{tab-orange}{Standard Taylor}
    & \input{\datapathJAXLaplacianExact/jet_naive_best.txt}
    & \input{\datapathJAXBilaplacianExact/jet_naive_best.txt}
    \\
    &
    & \textcolor{tab-green}{Collapsed (ours)}
    & \textbf{\input{\datapathJAXLaplacianExact/jet_simplified_best.txt}}
    & \textbf{\input{\datapathJAXBilaplacianExact/jet_simplified_best.txt}}
    \\ \cmidrule{2-5}
    & \multirow{3}{*}{\makecell{Mem.\,[MiB] \\ (differentiable)}}
    & \textcolor{tab-blue}{Nested 1\textsuperscript{st}-order}
    & \input{\datapathJAXLaplacianExact/hessian_trace_peakmem.txt}
    & \input{\datapathJAXBilaplacianExact/hessian_trace_peakmem.txt}
    \\
    &
    & \textcolor{tab-orange}{Standard Taylor}
    & \input{\datapathJAXLaplacianExact/jet_naive_peakmem.txt}
    & \input{\datapathJAXBilaplacianExact/jet_naive_peakmem.txt}
    \\
    &
    & \textcolor{tab-green}{Collapsed (ours)}
    & \textbf{\input{\datapathJAXLaplacianExact/jet_simplified_peakmem.txt}}
    & \textbf{\input{\datapathJAXBilaplacianExact/jet_simplified_peakmem.txt}}
    \\ \cmidrule{2-5}
    & \multirow{3}{*}{\makecell{Mem.\,[MiB] \\ (non-diff.)}}
    & \textcolor{tab-blue}{Nested 1\textsuperscript{st}-order}
    & \input{\datapathJAXLaplacianExact/hessian_trace_peakmem_nondifferentiable.txt}
    & \input{\datapathJAXBilaplacianExact/hessian_trace_peakmem_nondifferentiable.txt}
    \\
    &
    & \textcolor{tab-orange}{Standard Taylor}
    & \input{\datapathJAXLaplacianExact/jet_naive_peakmem_nondifferentiable.txt}
    & \input{\datapathJAXBilaplacianExact/jet_naive_peakmem_nondifferentiable.txt}
    \\
    &
    & \textcolor{tab-green}{Collapsed (ours)}
    & \textbf{\input{\datapathJAXLaplacianExact/jet_simplified_peakmem_nondifferentiable.txt}}
    & \textbf{\input{\datapathJAXBilaplacianExact/jet_simplified_peakmem_nondifferentiable.txt}}
    \\
    \midrule
    \multirow{9}{*}{\textbf{Stochastic}}
    & \multirow{3}{*}{Time [ms]}
    & \textcolor{tab-blue}{Nested 1\textsuperscript{st}-order}
    & \textbf{\input{\datapathJAXLaplacianStochastic/hessian_trace_best.txt}}
    & \input{\datapathJAXBilaplacianStochastic/hessian_trace_best.txt}
    \\
    &
    & \textcolor{tab-orange}{Standard Taylor}
    & \input{\datapathJAXLaplacianStochastic/jet_naive_best.txt}
    & \textbf{\input{\datapathJAXBilaplacianStochastic/jet_naive_best.txt}}
    \\
    &
    & \textcolor{tab-green}{Collapsed (ours)}
    & Not implemented %\textbf{\input{\datapathJAXLaplacianStochastic/jet_simplified_best.txt}}
    & Not implemented %\textbf{\input{\datapathJAXBilaplacianStochastic/jet_simplified_best.txt}}
    \\ \cmidrule{2-5}
    & \multirow{3}{*}{\makecell{Mem.\,[MiB] \\ (differentiable)}}
    & \textcolor{tab-blue}{Nested 1\textsuperscript{st}-order}
    & \input{\datapathJAXLaplacianStochastic/hessian_trace_peakmem.txt}
    & \input{\datapathJAXBilaplacianStochastic/hessian_trace_peakmem.txt}
    \\
    &
    & \textcolor{tab-orange}{Standard Taylor}
    & \textbf{\input{\datapathJAXLaplacianStochastic/jet_naive_peakmem.txt}}
    & \textbf{\input{\datapathJAXBilaplacianStochastic/jet_naive_peakmem.txt}}
    \\
    &
    & \textcolor{tab-green}{Collapsed (ours)}
    & Not implemented %\textbf{\input{\datapathJAXLaplacianStochastic/jet_simplified_peakmem.txt}}
    & Not implemented %\textbf{\input{\datapathJAXBilaplacianStochastic/jet_simplified_peakmem.txt}}
    \\ \cmidrule{2-5}
    & \multirow{3}{*}{\makecell{Mem.\,[MiB] \\ (non-diff.)}}
    & \textcolor{tab-blue}{Nested 1\textsuperscript{st}-order}
    & \textbf{\input{\datapathJAXLaplacianStochastic/hessian_trace_peakmem_nondifferentiable.txt}}
    & \input{\datapathJAXBilaplacianStochastic/hessian_trace_peakmem_nondifferentiable.txt}
    \\
    &
    & \textcolor{tab-orange}{Standard Taylor}
    & \input{\datapathJAXLaplacianStochastic/jet_naive_peakmem_nondifferentiable.txt}
    & \textbf{\input{\datapathJAXBilaplacianStochastic/jet_naive_peakmem_nondifferentiable.txt}}
    \\
    &
    & \textcolor{tab-green}{Collapsed (ours)}
    & Not implemented %\textbf{\input{\datapathJAXLaplacianStochastic/jet_simplified_peakmem_nondifferentiable.txt}}
    & Not implemented %\textbf{\input{\datapathJAXBilaplacianStochastic/jet_simplified_peakmem_nondifferentiable.txt}}
    \\
    \bottomrule
  \end{tabular}
  % Re-set the \num command to the original one
  \let\num\origsiunitxnum
\end{figure*}

\paragraph{Comparing JAX implementations.} Similar to our PyTorch experiment in \Cref{sec:experiments}, we compare three implementations of the Laplacian in JAX (all compiled with \texttt{jax.jit}):

\begin{enumerate}
\item \textbf{\textcolor{tab-blue}{Nested 1\textsuperscript{st}-order AD}} computes the Hessian using \texttt{jax.hessian}, which relies on forward-over-reverse mode, then traces it.

\item \textbf{\textcolor{tab-orange}{Standard Taylor mode}} propagates multiple univariate Taylor polynomials, each of which computes one element of the Hessian diagonal, then sums them to obtain the Laplacian.
  This is implemented with \texttt{jax.experimental.jet.jet} and \texttt{jax.vmap}.

\item \textbf{\textcolor{tab-green}{Collapsed Taylor mode}} relies on the forward Laplacian implementation in JAX provided by the \texttt{folx} library \cite{gao2023folx} and implements our proposed collapsed Taylor mode for the specific case of the Laplacian.
  \texttt{folx} also enables leveraging sparsity in the tensors, which is beneficial for architectures in VMC.
  To disentangle runtime improvements from sparsity detection versus collapsing Taylor coefficient, we disable \texttt{folx}'s sparsity detection.
\end{enumerate}
For the biharmonic operator, we simply nest the Laplacian implementations.

We only investigate computing the exact Laplacian, as the forward Laplacian in \texttt{folx} currently does not support stochastic computation.
We use the same neural network architecture as for our PyTorch experiments, fix the input dimension to $D=50$ and vary the batch size, recording the runtime and peak memory with the same protocol as described in the main text.
JAX is purely functional and therefore does not have a mechanism to build up a differentiable computational graph similar to evaluating a function in PyTorch where some leafs have \texttt{requires\_grad=True}.
To approximate the peak memory of computing a differentiable Laplacian in JAX, we measure the peak memory of first computing the Laplacian, then evaluating the gradient \wrt the neural network's parameters which backpropagates through the same computation graph built by PyTorch.

\paragraph{Results (Laplacian).} The left column of \cref{fig:jax-benchmark} visualizes the performance of the three implementations.
We fit linear functions to each of them and report the cost incurred by adding one more datum to the batch in \cref{tab:jax-benchmark}.
From them, we draw the following conclusions:

\begin{enumerate}
\item \textbf{Performance is consistent between PyTorch and JAX.} Although our PyTorch implementation does not leverage compilation, the values reported in \cref{tab:benchmark,tab:jax-benchmark} are consistent and only in rare cases differ by a factor of more than two.
  This confirms that our PyTorch-based implementation of Taylor mode is reasonably efficient, and that the presented performance results in the main text are transferable to other frameworks like JAX.

\item \textbf{Our implementation of collapsed Taylor mode based on graph rewrites in PyTorch achieves consistent speed-up with the Laplacian-specific implementation in JAX.}
  Specifically, we observe that \textcolor{tab-green}{collapsed Taylor mode/forward Laplacian} use roughly half the runtime of \textcolor{tab-blue}{nested 1\textsuperscript{st}-order AD} (compare \cref{tab:benchmark,tab:jax-benchmark}).
  This supports our argument that our collapsed Taylor is indeed a generalization of the forward Laplacian, \ie, the latter does not employ additional tricks (leveraging sparsity could also be applied to our approach but we are not aware of a drop-in implementation).
  It also illustrates that the savings we report in PyTorch carry over to other frameworks like JAX.

\item \textbf{JAX's \texttt{jit} compiler is unable to apply the graph rewrites we propose in this work.}
  If the JAX compiler was able to perform our proposed graph rewrites, then the \texttt{jit}-compiled \textcolor{tab-orange}{standard Taylor mode} should yield similar performance than the \textcolor{tab-green}{forward Laplacian}.
  However, we observe a clear performance gap in runtime and memory, from which we conclude that the compilation did not collapse the Taylor coefficients.
  Our contribution is to point out that such rewrites could easily be added to the compiler's ability to unlock these performance gains at zero user overhead.
\end{enumerate}

\paragraph{Results (biharmonic operator).}
For the biharmonic operator (right column of \cref{fig:jax-benchmark} and \cref{tab:jax-benchmark}), we conclude that (i) the most efficient way to compute biharmonics is by nesting Laplacians (compare with \cref{tab:benchmark} where Taylor mode uses the approach for general linear differential operators) and (ii) that nesting Taylor mode Laplacians is more efficient than nesting 1\textsuperscript{st}-order AD Laplacians, while also allowing to apply our collapsing technique.

% \begin{figure*}[!t]
%   \centering
%   %   From https://tex.stackexchange.com/a/7318
%   \newcolumntype{C}{ >{\centering\arraybackslash} m{0.11\textwidth} }
%   \newcolumntype{D}{ >{\centering\arraybackslash} m{0.4\textwidth} }
%   \begin{tabular}{CDD}
    %     & \textbf{PyTorch}
    %     & \textbf{JAX}
%     \\
%     \textbf{Exact}
%     & \includegraphics{jet/exp/exp01_benchmark_laplacian/figures/architecture_tanh_mlp_768_768_512_512_1_batch_size_256_device_cuda_name_bilaplacian_vary_dim.pdf}
%     & \includegraphics{jet/exp/exp04_jax_benchmark/figures/architecture_tanh_mlp_768_768_512_512_1_batch_size_256_device_cuda_name_jax_bilaplacian_vary_dim.pdf}
%   \end{tabular}
%   \caption{\textbf{Comparison of our proposed approach for computing the Laplacian in PyTorch (left) and by nesting Laplacians in JAX (right).}
%     The net is the same as before, but we fix the batch size to be $256$ and vary the net's input dimension.
%   }
%   \label{fig:benchmark-bilaplacians}
% \end{figure*}

%%% Local Variables:
%%% mode: LaTeX
%%% TeX-master: "../main"
%%% End:
