All right, good. Everybody has a seat. Yeah, there's. Is it the cable or is it. It could be the
cable.
So, yep, hello everyone. I'm Matt Graham. I am a research data scientist at the UCL Centre
for Advanced Research Computing, and today I'm going to be giving you a whirlwind to sort
of accelerate it in differentiable scientific computing with Jacks. So, just to begin with
an overview of what Jacks is, so it's a Python library which is perhaps best well known
for its use in kind of logic air machine learning and artificial intelligence context, but it's
actually also kind of using a broader scientific computing high-performance numerical computing context,
and that's what we're considering on here. It started out as a research project in Google
Brain by a group including some of the key contributors to the Python autographed package and the
XLA kind of compiler project. The initial version of Jacks was described in a 2018
system L conference paper, and there was an associated open source release under it
and I'm actually 2.0 license on getter been october 2018. Since that point, it's had seen
widespread adoption in both machine learning and scientific computing context, and there's a sort of large
and vibrant open source community that's developed around it, and in particular quite a kind of
very huge downstream package ecosystem building various functionality on top of it. So, just to kind of
give some of the key features that Jacks offers. So, I would say one of the sort of standout
features is if you're already familiar with the NumPy, kind of way of coding up problems
and with the scientific Python ecosystem, there's a very kind of shallow learning curve to get
them up and running with Jacks. It offers a almost drop in NumPy API replacement that you can
directly use in place of NumPy, but the real power of Jacks comes from the use
composable function transformations that are offered for doing things like just in time
compilation, automatic differentiation, automated vectorization, and parallelization. Under the
platform, it allows you running across both CPUs, but also accelerated devices, it's just
graphics processing units, and Google's 10s processing units. So, just to kind of give a bit of context
while I'm talking about Jacks, so my main experience is from developing downstream packages that
build on top of it, and so I'm not going to have much time to talk about these today, but if you're
interested and come and want to chat to me afterwards, I work on a package for approximate
which builds upon Jacks' differentiability and another one for doing spherical monotransforms
in sort of a space in F sciences context. So, how is Jacks actually structured? At the
top, it's the actual kind of Jacks' Python package, and it's public API, which exposes both
these functional transforms. I'll talk a little bit about various kind of numerical modules
for kind of building your functions on top of, and there's also a relatively well-developed
kind of series of APIs for providing extensions on top of Jacks.
These are the internals. Underneath that, I'm one of the key components of that is Jack's
loop, which is the sort of Python wrappers around the binary components of Jacks, which
is separated from the sort of pure Python Jacks package. Alongside that, there are some
vendor-specific plugins for some of the multi-device support. Underneath all of that, Jacks
basically builds along on top of the kind of open-axle compiler infrastructure, uses the
stable-hitch-low MLI-elect for the kind of communication between the Python package and the
xeno-compiler tool chain, and that's what gives you this kind of device-agnostic behavior.
So, just to kind of flag that from our recent versions of Jacks and, and, indeed,
just numpy, there's support for what's called the Python API, standard, which allows you to write code,
which is possible across different array back-ends. So, this would look something like rather than
importing numpy as NP, you would take your array namespace from one of your array arguments,
using this special, and the array namespace method, and then define your numerical operations
in terms of that. And this gives you the advantage of if you pass in a Jacks array, you will
dispatch to the Jacks operations, but if you pass in a numpy array, you will dispatch to the numpy operations.
So, I've mentioned Jacks has these functional transformations, just to dig into a little bit how they actually
work, and in particular, the just-in-time compilation transformation. So, Jacks transforms work by
twisting a function with these kind of abstract arguments that represent the specific shapes and types of the
arguments and trace all of the different numerical operations that apply to those.
For particularly, the just-in-time compilation, the output of that tracing is an internal
intermediate representation called a Jacksper, short-for-a-Jacks expression.
There are then a series of lowering rules defined to translate those numerical
primitives that you've traced down the Jacksper to operations in the stable HTML,
which is then what has passed to the XLE compilers, executable compiled,
wrapped as a new Python function, and executed in the functions with outputs returned,
and the executable cached, so that we don't hit that kind of compiler overhead on subsequent
function implications. So, in code, what does this look like? If we have some Python function,
we apply our Jacks.ditch transform, we get a new function out, and then we would expect
if we evaluate that on some real argument, we will get numerically equivalent outputs,
but for larger functions, we would have a performance gain from the sort of jittered version once
we amortized a cost of compilation. So, as well as just-in-time compilation,
one of the other key features that Jacksper provides is automatic differentiation.
So, one of the most basic instances of this is the Jacks.grad,
transform, so this takes a skill of valued function, which is the function which computed the
gradient, we would respect any real arguments. So, if we had applied that to the simple example,
we quoted up earlier, rather than getting scalar out now, we're getting a ray that represents the
derivatives of respect to each of the elements. And just to kind of flag, because this is using a sort of
efficient reverse mode, ultimately a definition implementation, the gradients are actually
similar kind of constant overhead factor from just evaluating the function itself.
The AD transform uses the same tracing approach as just-in-time compilation,
and indeed the other transforms that Jacks provides, also built on top of this tracing architecture,
and can be composed with each other, so you can do just-in-time compilation of your derivatives,
and so on. So, just to kind of touch on some of their sharp edges,
it can be a little bit challenging to work around, maybe when you're first using Jacks.
So, the eagle-eyed menu may have noticed, and some of those examples there that we were
getting kind of single-position, floating-point outputs. So, this is kind of aligned with
you know, Jacks is typically used often on accelerators, so it's default
semantics around it, it types a more detailed for that. Use case,
if you want to get closer behavior to what you would expect from kind of base number pie,
you will need to enable this Jacks. Use this Jacks enable X64 configuration option,
Jacks arrays are immutable, and in general it kind of assumes a sort of function
program style, whether the function is being transformed appear,
and to do kind of indexed updates, you need to use these kind of functional
app methods for other than the standard numpy in place update syntax.
So, just-in-time compilation in particular, put some requirements on the trace control flow,
so this can only depend on if using base python control flow on the shapes and
data types of non-stack array arguments, and you also need to be able to
infer all of the intermediate and output arrays, statically from the
input arguments. There are structured control primitives that
Jacks provides that do allow data dependent control flow within a
jit compilation context, but then you're having kind of step away from
just using basic python syntax. So just to conclude, in summary,
so I would say the sort of major selling point of Jacks is it's a very
low barrier entry to exploiting GPU acceleration,
and indeed doing that in a multi-banger way, it does support
Nvidia and AMD GPUs, and I think there's some limited support for
Intel as well, and on top of that, getting automatic differentiation
and the other trans-films at Jacks offers, with Jacks and NumPy
and a number of other libraries now supporting this array API.
You can actually write portable code across these different
ray back ends with relatively minimal effort, just compared to talking
just one of them, and the requirement, though, to be able to trace
functions as part of Jacks' transform infrastructure does put
some limitations on, for example, control flow, and that can somewhat
run counter-to-differ example that device portable, so back end portability,
so if you're having to lean on these structure control
permutives, they're Jacks specific, and so you do end up being a
little bit more tied in some cases. Okay, that's everything.
Next for listening, if you are interested more learning more about Jacks,
there's a longer tutorial I give to the archer to training series,
back in December, and there's all material available for it here.
Thank you.