Learn Julia For Beginners – The Future Programming Language of Data Science and Machine Learning Explained

By: n Logan Kilpatrick n

Re-posted from: https://www.freecodecamp.org/news/learn-julia-programming-language/

Learn Julia For Beginners – The Future Programming Language of  Data Science and Machine Learning Explained

Julia is a high-level, dynamic programming language, designed to give users the speed of C/C++ while remaining as easy to use as Python. This means that developers can solve problems faster and more effectively.

Julia is great for computational complex problems. Many early adopters of Julia were concentrated in scientific domains like Chemistry, Biology, and Machine Learning.

This said, Julia is general-purpose language and can be used for tasks like Web Development, Game Development, and more. Many view Julia as the next-generation language for Machine Learning and Data Science, including the CEO of Shopify (among many others):

Praise for Julia by Shopify’s CEO Tobi

How to Download the Julia Programming Language ⤵️

There are two main ways to run Julia: via a .jl file in an IDE like VS Code or command by command in the Julia REPL (Read Evaluate Print Loop). In this guide, we will mainly use the Julia REPL. Before you can use either, you will need to download Julia:

or just head to: https://julialang.org/downloads/

After you have Julia installed, you should be able to launch it and see:

Screen-Shot-2021-12-24-at-5.56.14-AM
Julia 1.7 REPL after install

Julia Programming Language Basics for Beginners

Before we can use Julia for all of the exciting things it was built for like Machine Learning or Data Science, we first need to get familiar with the basics of the language.

We will start by going over variables, types, and conditionals. Then, we will talk about loops, functions, and packages. Last, we’ll touch on more advanced concepts like structs and talk about additional learning resources.

This is going to be a whirlwind tour so strap in and get ready! It is also worth noting that this tutorial assumes you have some basic familiarity with programming. If you don’t, check out this course on an Intro to Julia for Nervous Beginners.

An Introduction to Julia Variables and Types ⌨️

In Julia, variables are dynamically typed, meaning that you do not need to specify the variable’s type when you create it.

julia> a = 10 # Create the variable "a" and assign it the number 10
10

julia> a + 10 # Do a basic math operation using "a"
20

(Note that in code snippets, when you see julia> it means the code is being run in the REPL)

Just like we defined a variable above and assigned it an integer (whole number), we can also do something similar for strings and other variable types:

julia> my_string = "Hello freeCodeCamp" # Define a string variable
"Hello freeCodeCamp"

julia> balance = 238.19 # Define a float variable 
238.19

When creating variables in Julia, the variable name will always go on the left-hand side, and the value will always go on the right-hand side after the equals sign. We can also create new variables based on the values of other variables:

julia> new_balance = balance + a
248.19

Here we can see that the new_balance is now the sum (total) of 238.19 and 10. Note further that the type of new_balance is a float (number with decimal place precision) because when we add a float and int together, we automatically get the type with higher precision, which in this case is a float. We can confirm this by doing:

julia> typeof(new_balance)
Float64

Due to the nature of dynamic typing, variables in Julia can also change type. This means that at one point, holder_balance could be a float, and then later on it could be a string:

julia> holder_balance = 100.34
100.34

julia> holder_balance = "The Type has changed"
"The Type has changed"

julia> typeof(holder_balance)
String

You may also be excited to know that variable names in Julia are very flexible, in fact, you can do something like:

julia> ? = 10
10

julia> ? = -10
-10

julia> ? + ?
0

On top of emoji variable names, you can also use any other Unicode variable name which is very helpful when you are trying to represent mathematical ideas. You can access these Unicode variables by doing a \ and then typing the name, followed by pressing tab:

julia> \sigma # press tab and it will render the symbol

julia> σ = 10 # set sigma equal to 10

Overall, the variable system in Julia is flexible and provides a huge set of features that make writing Julia code easy while still being expressive. If you want to learn more about variables in Julia, check out the Julia documentation: https://docs.julialang.org/en/v1/manual/variables/

How to Write Conditional Statements in Julia  ?

In programming, you often need to check certain conditions in order to make sure that specific lines of code run. For example, if you write a banking program, you might only want to let someone withdraw money if the amount they are trying to withdraw is less than the amount they have present in their account.

Let us look at a basic example of a conditional statement in Julia:

julia> bank_balance = 4583.11
4583.11

julia> withdraw_amount = 250
250

julia> if withdraw_amount <= bank_balance
           bank_balance -= withdraw_amount
           print("Withdrew ", withdraw_amount, " from your account")
       end
Withdrew 250 from your account

Let us take a closer look here at some parts of the if statement that might differ from other code you have seen: First, we use no : to denote the end of the line and we also are not required to use () around the statement (though it is encouraged). Next, we don’t use {} or the like to denote the end of the conditional, instead, we use the end keyword.

Just like we used the if statement, we can chain it with an else or an elseif:

julia> withdraw_amount = 4600
4600

julia> if withdraw_amount <= bank_balance
           bank_balance -= withdraw_amount
           print("Withdrew ", withdraw_amount, " from your account")
       else
           print("Insufficent balance")
       end
Insufficent balance

You can read more about control flow and conditional expressions in the Julia documentation: https://docs.julialang.org/en/v1/manual/control-flow/#man-conditional-evaluation

How to use Loops in Julia ?

There are two main types of loops in Julia: a for loop and a while loop. As is the same with other languages, the biggest difference is that in a for loop, you are going through a pre-defined number of items whereas, in a while loop, you are iterating until some condition is changed.

Syntactically, the loops in Julia look very similar in structure to the if conditionals we just looked at:

julia> greeting = ["Hello", "world", "and", "welcome", "to", "freeCodeCamp"] # define greeting, an array of strings
6-element Vector{String}:
 "Hello"
 "world"
 "and"
 "welcome"
 "to"
 "freeCodeCamp"

julia> for word in greeting
           print(word, " ")
       end
Hello world and welcome to freeCodeCamp 

In this example, we first defined a new type: a vector (also called an array). This array is holding a bunch of strings we defined. The behavior is very similar to that of arrays in other languages but it is worth noting that arrays are mutable (meaning you can change the number of items in the array after you create it).

Again, when we look at the structure of the for loop, you can see that we are iterating through the greeting variable. Each time through, we get a new word (in this case) from the array and assign it to a temporary variable word which we then print out. You will notice that the structure of this loop looks similar to the if statement and again uses the end keyword.

Now that we explored for loops, let us switch gears and take a look at a while loop in Julia:

julia> while user_input != "End"
           print("Enter some input, or End to quit: ")
           user_input = readline() # Prompt the user for input
       end
Enter some input, or End to quit: hi
Enter some input, or End to quit: test
Enter some input, or End to quit: no
Enter some input, or End to quit: End

For this while loop, we set it up so that it will run indefinitely until the user typed the word “End”. As you have now seen it a few times, the structure of the loop should start to look familiar.

If you want to see some more examples of loops in Julia, you can check out the Julia Documentation’s section on loops: https://docs.julialang.org/en/v1/manual/control-flow/#man-loops

How to use Functions in Julia

Functions are used to create multiple lines of code, chained together, and accessible when you reference a function name. First, let us look at an example of a basic function:

julia> function greet()
           print("Hello new Julia user!")
       end
greet (generic function with 1 method)

julia> greet()
Hello new Julia user!

Functions can also take arguments, just like in other languages:

julia> function greetuser(user_name)
           print("Hello ", user_name, ", welcome to the Julia Community")
       end
greetuser (generic function with 1 method)

julia> greetuser("Logan")
Hello Logan, welcome to the Julia Community

In this example, we take in one argument, and then add its value to the print out. But what if we don’t get a string?

julia> greetuser(true)
Hello true, welcome to the Julia Community

In this case, since we are just printing, the function continues to work despite not taking in a string anymore and instead of taking a boolean value (true or false). To prevent this from occurring, we can explicitly type the input arguments as follows:

julia> function greetuser(user_name::String)
           print("Hello ", user_name, ", welcome to the Julia Community")
       end
greetuser (generic function with 2 methods)

julia> greetuser("Logan")
Hello Logan, welcome to the Julia Community

So now the function is defined to take in only a string. Let us test this out to make sure we can only call the function with a string value:

julia> greetuser(true)
Hello true, welcome to the Julia Community

Wait a second, why is this happening? We re-defined the greetuser function, it should not take true anymore.

What we are experiencing here is one of the most powerful underlying features of Julia: Multiple Dispatch. Julia allows us to define functions with the same name and number of arguments but that accept different types. This means we can build either generic or type specific versions of functions which helps immensely with code readability since you don’t need to handle every scenario in one function.

We should quickly confirm that we actually defined both functions:

julia> methods(greetuser)
# 2 methods for generic function "greetuser":
[1] greetuser(user_name::String) in Main at REPL[34]:1
[2] greetuser(user_name) in Main at REPL[30]:1

The built-in methods function is perfect for this and it tells us we have two functions defined, with the only difference being one takes in any type, and the other takes in just a string.

It is worth noting that since we defined a specialized version that accepts just a string, anytime we call the function with a string it will call the specialized version. The more generic function will not be called when a string is passed in.

Next, let us talk about returning values from a function. In Julia, you have two options, you can use the explicit return keyword, or you can opt to do it implicitly by having the last expression in the function serve as the return value like so:

julia> function sayhi()
           "This is a test"
           "hi"
       end
sayhi (generic function with 1 method)

julia> sayhi()
"hi"

In the above example, the string value “hi” is returned from the function since it is the last expression and there is no explicit return statement. You could also define the function like:

julia> function sayhi()
           "This is a test"
          return "hi"
       end
sayhi (generic function with 1 method)

julia> sayhi()
"hi"

In general, from a readability standpoint, it makes sense to use the explicit return statement in case someone reading your code does not know about the implicit return behavior in Julia functions.

Another useful functions feature is the ability to provide optional arguments:


julia> function sayhello(response="hello")
          return response
       end
sayhello (generic function with 2 methods)

julia> sayhello()
"hello"

julia> sayhello("hi")
"hi"

In this example, we define response as an optional argument so that we can either allow it to use the default behavior we defined or we can manually override it when necessary. These examples just scratch the surface on what is possible with functions in Julia. If you want to read more about all the cool things you can do, check out: https://docs.julialang.org/en/v1/manual/functions/

How to use Packages in Julia ?

The Julia package manager and package ecosystem are some of the most important features of the language. I actually wrote an entire article on why it is one of the most underrate features of the language.

With that said, there are two ways to interact with packages in Julia: via the REPL or using the Pkg package. We will mostly focus on the REPL in this post since it is much easier to use in my experience.

After you have Julia installed, you can enter the package manager from the REPL by typing ].

Screen-Shot-2021-12-26-at-9.50.05-AM
Pkg mode in the Julia REPL

Now that we are in the package manager, there are a few things we commonly want to do:

  • Add a package
  • Remove a package
  • Check what is already installed

If you want to see all the possible commands in the REPL, simply enter Pkg mode by typing ] and then type ?  followed by the enter / return key.

How to Add Julia Packages ➕

Let’s add our first package, Example.jl . To do so, we can run:

(@v1.7) pkg> add Example

which should provide output that looks something like:

(@v1.7) pkg> add Example
Updating registry at `~/.julia/registries/General`
Updating git-repo `https://github.com/JuliaRegistries/General.git`
Updating registry at `~/.julia/registries/JuliaPOMDP`
Updating git-repo `https://github.com/JuliaPOMDP/Registry`
Resolving package versions...
Installed Example ─ v0.5.3
Updating `~/.julia/environments/v1.7/Project.toml`
[7876af07] + Example v0.5.3
Updating `~/.julia/environments/v1.7/Manifest.toml`
[7876af07] + Example v0.5.3
Precompiling project...
1 dependency successfully precompiled in 1 seconds (69 already precompiled)
(@v1.7) pkg>

For space reasons, I will skip further outputs under the assumption that you are following along with me.

How to Check the Package Status in Julia ?

Now that we think we have a package installed, let’s doublecheck if it is really there by typing status (or st for shorthand) into the package manager:

(@v1.7) pkg> st
Status `~/.julia/environments/v1.7/Project.toml`
[7876af07] Example v0.5.3
[587475ba] Flux v0.12.8

Here we can see I have two packages installed, Flux and Example. It also gives me the path to the file which manages my current environment (in this case, global Julia v1.7) along with the package versions I have installed.

How to Remove a Julia package ?

If I wanted to remove a package from my active environment, like Flux, I can simply type remove Flux (or rm as the shorthand):

(@v1.7) pkg> rm Flux
Updating `~/.julia/environments/v1.7/Project.toml`
[587475ba] - Flux v0.12.8

A quick status afterward shows this was successful:

(@v1.7) pkg> st
Status `~/.julia/environments/v1.7/Project.toml`
[7876af07] Example v0.5.3

We now know the very basics of working with packages. But we have committed a major programming crime, using our global package environment.

How to Use Julia Packages ?

Now that we have gone over how to manage packages, let’s explore how to use them. Quite simply, you just need to type using packageName to use a specific package you want. One of my favorite new features in Julia 1.7 (highlighted in this blog post) is shown below:

1-jI58_UDd87Q4fQ326r6E6Q
Image captured by Author

If you recall, we removed the Flux package, and of course, I forgot this so I went to use it and load it in by typing using Flux. The REPL automatically prompts me to install it via a simple “y/n” prompt. This is a small feature but saves a tremendous amount of time and potential confusion.

It is worth noting that there are two ways to access a package’s exported functions: via the using keyword and the import keyword. The big difference is that using automatically brings all of the functions into the current namespace (for which you can think about as a big list of functions which Julia knows the definitions) whereas import gives you access to all of the functions but you have to prefix the function with the package name like: Flux.gradient() where Flux is the name of the package and gradient() is the name of a function.


How to use Structs in Julia?

Julia does not have Object Orientated Programming (OOP) paradigms built into the language like classes. However, structs in Julia can be used similar to classes to create custom objects and types. Below, we will show a basic example:

julia> mutable struct dog
           breed::String
           paws::Int
           name::String
           weight::Float64
       end

julia> my_dog = dog("Australian Shepard", 4, "Indy", 34.0)
dog("Australian Shepard", 4, "Indy", 34.0)

julia> my_dog.name
"Indy"

In this example, we define a struct to represent a dog. In the struct, we define four attributes which make up the dog object. In the lines after that, we show the code to actually create a dog object and then access some of its attributes. Note that you need not specify the types of the attributes, you could leave it more open. For this example, we defined explicit types to highlight that feature.

You will notice that similar to classes in Python (and other languages), we did not define an explicit constructor to create the dog object. We can, however, define one if that would be useful to use:

julia> mutable struct dog
           breed::String
           paws::Int
           name::String
           weight::Float64
           
           function dog(breed, name, weight, paws=4)
               new(breed, paws, name, weight)
           end
       end

julia> new_dog = dog("German Shepard", "Champ", 46)
dog("German Shepard", 4, "Champ", 46.0)

Here we defined a constructor and used the special keyword new in order to create the object at the end of the function. You can also create getters and setters specifically for the dog object by doing the following:

julia> function get_name(dog_obj::dog)
           print("The dogs's name is: ", dog_obj.name)
       end
get_name (generic function with 1 method)

julia> get_name(new_dog)
The dogs's name is: Champ

In this example, the get_name function only takes an object of type dog. If you try to pass in something else, it will error out:

julia> get_name("test")
ERROR: MethodError: no method matching get_name(::String)
Closest candidates are:
  get_name(::dog) at REPL[61]:1
Stacktrace:
 [1] top-level scope
   @ REPL[63]:1

It is worth noting that we also defined the struct to be mutable initially so that we could change the field values after we created the object. You omit the keyword mutable if you want the objects initial state to persist.

Structs in Julia not only allow us to create object’s, we also are defining a custom type in the process:

julia> typeof(new_dog)
dog

In general, structs are used heavily across the Julia ecosystem and you can learn more about them in the docs: https://docs.julialang.org/en/v1/base/base/#struct

Additional Julia Programming Learning Resources ?

I hope that this tutorial helped get you up to speed on many of the core ideas of the Julia language. With that said, I know that there are still gaps as this is an extended but non-comprehensive guide. To learn more about Julia, you can check out the learning tab on the Julia website: https://julialang.org/learning/ which has guided courses, YouTube videos, and mentored practice problems.

If you have other questions or need help getting started with Julia, please feel free to get in touch with me: https://twitter.com/OfficialLoganK

Optimization with JuMP in Julia

By: Jaafar Ballout

Re-posted from: https://www.supplychaindataanalytics.com/optimization-with-jump-in-julia/

As a result of the emergence and progress of open-source languages SCM and OR analysts have access to a growing number of constantly improving tools and solvers. Julia, for example, is a flexible dynamic open-source programming language that is appropriate for scientific and numerical computing. Julia’s source code is available on GitHub and the programming language is supported by a growing number of frameworks and applications for optimization and scheduling.

I introduced Julia in one of my previous posts already, covering a linear programming example. In this article I will give a more comprehensive introduction to Julia, comparing it with other programming languages, and specifically focusing on optimization.

JuMP: A package for mathematical programming in Julia

JuMP (Julia for Mathematical Programming) is a package available in Julia. This modelling framework allows users to formulate an optimization problem (linear, mixed-integer, quadratic, conic quadratic, semidefinite, and nonlinear) with easy-to-read code. The problem can then be solved by optimizers (solvers) written in low-level languages. The selected solver for this demonstration is GLPK which is a state-of-art open-source solver for linear programming (LP) and mixed-integer programming (MIP) problems. It incorporates algorithms such as the simplex method and the interior-point method.

Comparing Julia vs. Matlab for optimization problems

At first, a quick comparison between Julia and Matlab is necessary to appreciate the power of open-source programming languages.

Criteria Julia ? Matlab ?
Cost Free Hundreds of dollars per year
License Open-Source One year user license
Source Non-profitable foundation and community MathWorks company
Scope Mainly numeric Numeric only
Performance Excellent Faster than Python but slower than Julia
Editor/IDE Jupyter is recommended IDE already included
Packaging Pkg manager included Toolboxes included
Usage Generic Industry and research in academia
Fame Young but starts to be known Old and well-known
Support Community MathWorks
Documentation Growing Okay

JuMP and GLPK packages for optimization with Julia

JuMP:

  • Modeling language and collection of supporting packages for mathematical optimization in Julia
  • Supported solvers: GLPK, CLP, CBC, CPLEX, ECOS, GUROBI, HiGHS, MOSEK, XPRESS, OSQP, PATH, ProxSDP SCIP, SCS, SDPA.

The first step is to add the required packages: JuMP and GLPK. This is done using Pkg which is a built-in package manager in Julia. I can add the requried or desired packages using the following commands in Julia REPL. If Julia is installed, typing Julia in the command line is enough to open REPL (Julia run command).

using Pkg
Pkg.add("JuMP")
Pkg.add("GLPK")

Now, I should import the necessary packages in the IDE which is a Jupyter notebook in our case.

# Import pacakges: JuMP and GLPK 
using JuMP # Julia package - high level programming language
using GLPK # open-source solver - low level language

A test problem

Let us consider the following optimization model:

The feasible region (as x and y are continuous decision variables) is shown below:

According to the gradient of the objective (the red arrow) and direction of the optimization (maximization), the green point is the optimal solution, in which x=3.75, y=1.25, and the optimal value of the objective function is 23.75. I will solve this simple continuous linear programming model with the JuMP package in Julia and comment on its syntax.

Solving the problem

The following commented code aims to solve the proposed linear programming model with JuMP (the name of the package) in Julia:

# Define the model
m = Model(); # m stands for model

# Set the optimizer
set_optimizer(m,GLPK.Optimizer); # calling the GLPK solver

# Define the variables
@variable(m, x>=0);
@variable(m, y>=0);

# Define the constraints 
@constraint(m, x+y<=5);
@constraint(m, 10x+6y<=45);

# Define the objective function
@objective(m, Max, 5x+4y);

# Run the solver
optimize!(m);

# Output
println(objective_value(m)) # optimal value z
println("x = ", value.(x), "\n","y = ",value.(y)) # optimal solution x & y

In future work I will explore the power of Julia in solving supply chain management problems, specifically mixed integer programming (MIP).

The post Optimization with JuMP in Julia appeared first on Supply Chain Data Analytics.

Engineering Trade-Offs in Automatic Differentiation: from TensorFlow and PyTorch to Jax and Julia

By: Christopher Rackauckas

Re-posted from: http://www.stochasticlifestyle.com/engineering-trade-offs-in-automatic-differentiation-from-tensorflow-and-pytorch-to-jax-and-julia/

To understand the differences between automatic differentiation libraries, let’s talk about the engineering trade-offs that were made. I would personally say that none of these libraries are “better” than another, they simply all make engineering trade-offs based on the domains and use cases they were aiming to satisfy. The easiest way to describe these trade-offs is to follow the evolution and see how each new library tweaked the trade-offs made of the previous.

Early TensorFlow used a graph building system, i.e. it required users to essentially define variables in a specific graph language separate from the host language. You had to define “TensorFlow variables” and “TensorFlow ops”, and the AD would then be performed on this static graph. Control flow constructs were limited to the constructs that could be represented statically. For example, an `ifelse` function statement is very different from a conditional `if` then `else` of code because `ifelse` would semantically be the same as always calling both branches and then choosing the result, thus only having a single code path (though I say semantically because further compiler optimizations may and usually do reduce that). This static sublanguage is then represented in an intermediate representation (IR) known as XLA which then performed a lot of simplification of linear algebra, and AD was done using the simple graph representation algorithms given that there was no true control flow at this representation. While this gives a lot of efficiency (XLA is great for simplification because it can easily see the whole world), it of course had some major downsides in terms of flexibility and convenience.

Thus you can almost think of this as a source code transformation because all of the autodiff is done on essentially an IR for a language which is not the same as the host language, but for the most part it was requiring the user does the translation to the new language for the AD system which is… rather inconvenient.

PyTorch came along to solve the flexibility and convenience issues by instead using a tape-based method. It generates the code to autodiff every time you run the forward pass by simply storing the operation that it sees in a given forward pass, and then differentiates that set of operations in reverse. This “building of the tape” is done by operator overloading as part of the Tensor type PyTorch says you need to use. How it works is easy to see. For example, f(2.0) would take the first branch of the if statement and then run the while loop 5 times. So then the AD pass would take that set of operations and start running backpropagation through 5 passes of the while loop and back through the first branch. Notice that by using this form, the AD does not “see” any dynamic control flow: that was all in Python, but not in the tape. Thus the AD does not have to handle dynamic control flow, and this makes it very easy to handle a lot of odd cases of the language. The downside to this approach though is that the AD is “per value”, i.e. you cannot do a lot of optimizations on the backwards passes because you will not necessarily ever see the same backwards pass again, and this allows for a lot less optimization.

Does this harm PyTorch’s efficiency beyond repair? Well, no and yes. No it does not harm efficiency in the sense of, most machine learning algorithms are so heavily reliant on expensive kernels, such as matrix multiplication (`A*x`), `conv`, etc., so the amount of work per operation is extremely high in most ML applications that it hides the overhead of this approach. This allows the PyTorch team to spend most of its time optimizing the 2,000+ operators that it provides, and so most people in ML see PyTorch as fast because it comes with fast kernels (fast conv calls, fast GPU linear algebra) despite the AD overhead. That said, you can very easily run into cases where AD and Python interpreter overhead are not washed out. Cases of that are where your arrays are small or where a lot of scalar operations are happening, for example the Julia vs PyTorch Neural ODE benchmarks on cases matching scientific model discovery workflows you see a 100x performance improvement in Julia (even major differences without AD in the ODE and SDE solvers), and can mostly be attributed to language and AD overhead due to the small kernels used in these cases. For this reason the PyTorch team has been working on things like `torch.@jit` as a separate sublanguage that can compile and optimize differently from the rest of the code, specific to handling these cases, though there’s a lot of discussion of the long-term viability of that approach. But anyways, PyTorch has done really well because it made good choices for its domain of use.

So then TensorFlow Eager (2.0) comes around as adds dynamic control flow support in a manner similar to PyTorch as a sad attempt to get everyone back, but of course then it doesn’t play nicely will all of the XLA tooling (because it cannot see the whole graph of all possible operations for all input values to optimize it well) so it didn’t hit the TensorFlow speeds everyone was expecting, so it was kind of the worst of all worlds.

Subsequent tools then all sought ways to either expand the domains of these ideas or try to mix some of the advantages of the two sides. Jax is one of those. Jax uses non-standard interpretation to build a copy of the full code in its own IR to then perform AD on, finally lowering it to TensorFlow’s XLA for optimizations. Jax’s non-standard interpretation is kind of like operator overloading in that it has special objects walk through code in order to build out the exprs (this is called the “tracing” step). But wait, how is it able to trace the full code if there’s dynamic control flow, won’t it have the same issues as PyTorch that it only sees parts of the full code’s potential paths? Indeed that is true, and that’s why it doesn’t want you to use full dynamic control flow and instead use Jax primitives like lax.while, which are function calls that can be caught during tracing to avoid the code having true dynamic behavior at trace time. Also, for this to be true you need that what your function does can be completely determined by its inputs, i.e. the functions must be “pure”. For this reason Jax requires programming in a functional style with pure functions rather than the object-oriented standard of Python, thus a notable trade-off of the abstract interpretation approach. But what you essentially get is a more natural graph builder for TensorFlow, because at the end of the day it ends up in TensorFlow’s XLA IR, and so you get the same efficiency there but in a form that can look and feel a lot more natural. The downside of course is that you still don’t have true dynamism which is why those linked primitives exist, and why they are not well optimized as described in Jax – The Sharp Bits. However, “most” ML algorithms don’t use very much dynamism (example: recurrent neural networks know how many layers they have, they don’t have a while loop iterate to tolerance), and so “most” algorithms tend to do well in this sublanguage. In that sense, it can optimize a lot of codes rather naturally.

What about keeping dynamism in the AD?

This of course then begs the question, is it possible to keep the full dynamism of the host language in the AD system? It is possible, but it is hard. This is what a lot of the Julia AD tools have focused on with source code transformations (along with Swift for TensorFlow). However, since source code is “for humans”, it can be a rather difficult level to algorithmically work on. Thus instead these tools work on lowered IR, where these lowered representations remove a lot of the “cruft” of syntax to give a much smaller support surface. This was the core of Zygote.jl’s approach where it saw that by acting on the SSA IR it could directly support control flow like while loops without unrolling them into sets of operations (like PyTorch or TensorFlow Eager) or only supporting a sublanguage of control flow (like Jax). This is essentially done by converting while loops and other dynamic constructs into static (source code) representations that have new lines of code in there for things like stacks that keep information about the forward pass (like which branch is taken), and then these stacks are accessed and used in the generated backwards pass. Thus what code is generated is not dynamic (a while loop forward gives a for loop in reverse), but the generated backwords pass is dynamic (because it uses the stack to tell it how many times to walk the for loop). This allows AD to have a single code for all branches (unlike the tape-building forms) and thus it can optimize more like TensorFlow but in a world where the dynamic control flow is not eliminated.

Well that sounds like the best of both worlds, so why isn’t everyone using it? There’s two factors involved in that. First, accepting that your AD will have to deal with the full dynamic nature of an entire programming language means accepting a much more difficult job. The whole purpose of the AD approaches in TensorFlow/PyTorch/Jax is for these constructs to be eliminated before the AD, so they have a much smaller surface of language support required. Because of this added complexity, this pretty much guarantees you cannot use Python because it’s such a crazy language in terms of what it allows with dynamism (fun fact, the Jax folks at Google Brain did have a Python source code transform AD at one point but it was scrapped essentially because of these difficulties), and so people working on these solutions flocked to languages with clear syntax that is easy for compilers to optimize, i.e. Julia and Swift. Python has most of the ML crowd, so that creates a barrier to entry.

But even then, the problem is still very hard. In Julia it was found that Zygote acts on too high of an IR, i.e. before compiler optimizations, which then requires you do AD on unoptimized code only delete most of the work later, and so it would be better for it to go even lower. This is the reason why the Diffractor.jl project started. But there’s even a reason to act lower, since some optimization only occurs at the LLVM level, which is why Julia developers started directly building an AD system that acts on LLVM’s IR itself known as Enzyme (note that while this project included members of the Julia Lab like Valentin, because it acts at the LLVM level it is applicable to any LLVM compiled language, such as C/C++ (Clang) or Rust). There is then a trade-off that occurs with source code transform methods as you go lower and lower in the IRs which I describe in a separate post. tl;dr there: Enzyme can act after compiler optimizations so much of the higher level information might be deleted (at least, without completion of dialects like MLIR which aren’t quite ready). Enzyme only sees the barest of low level code so it may not have the high level linear algebra definitions to do all of the linear algebra simplifications, like how XLA will fuse many matrix-vector multiplications into a matrix-matrix multiplication, since some of the function calls may have been inlined and deleted. Optimizing this remining loopy code to reach BLAS speeds is thus as hard as generating looping code that reaches BLAS speeds, and history shows this is hard but not impossible. Additionally, function calls to a nonlinear solver may have already been deleted, so optimized adjoints which outperform the direct differentiation of code, like in the case of Deep Equilibrium Models (DEQs), may end up less optimized. But that lowest level allows for very efficient scalar code differentiation and mutation support. On the other hand, Diffractor uses Julia’s typed IR so it can apply higher level rules easily and consistently, and in theory it can do transformations similar to XLA (i.e. keeping BLAS calls intact and fusing them). But writing such analyses on a fully dynamic compute IR is difficult enough that it has not been done. Tooling around escape analysis and shape propagation are being built to try and enable such optimizations, but the fact remains that it’s a lot more work to do it on a language IR instead of a sublanguage graph like XLA. In theory you could have compiler passes prove that a function is semi-static in the sense of XLA and get the same optimizations as Jax or TensorFlow, but that doesn’t happen today and it’s not easy to do. The future of Julia AD systems will likely mix the Enzyme and Diffractor approaches to tackle this issue, but the clear trade-off being made here is generality at the cost of implementation complexity.

The second factor, and probably the more damning one, is that most ML codes don’t actually use that much dynamism. Recurrent neural networks, transformers, convolutions, etc. all have simpler forms of dynamism which in some sense is quite static. That’s an important trade-off most people don’t always consider: why solve problems your users don’t have? The number of layers you have do not depend on the values coming out of the layers. Support for dynamism for ML workflows is thus mostly about convenience, not necessity. When algorithms do have dynamism, in most cases you can get away with wrapping it as an operation in the language, i.e. defining a function and defining the adjoint derivative for that function. This for example is how Jax supports ODEs even though adaptive ODE solvers require knowing the calculated values in order to determine the number of steps. You cannot differentiate an ODE code with Jax, but if you use an ODE solver with a defined adjoint you are okay. While this does mean that some algorithms are not possible with Jax (at least without forgoing a lot of optimizations), and algorithms where differntiating solvers is fundamentally different from adjoint definitions can limit which performance/stability trade-offs can be made (see the supplemental section 8 for details in the case of ODEs about stability of “discrete adjoints”]), these factors seem to be rather rare in standard ML use cases which is why most people haven’t bothered to learn a new programming language to get around these issues.

That leaves us where we are today. Are more ML algorithms of the future going to require handling more dynamic structures? Is optimizing scalar and mutating code going to be important for people using AD systems? The reason why I know this story so well is because the answer for my domain, scientific machine learning (SciML), is yes. Climate models use mutation because reallocating huge buffers would greatly effect performance. Adaptive solvers on stiff equations are a fact of life, so simple adjoints used in PyTorch and Jax are unstable and simply give Inf as the gradients in these cases. Time will tell whether this physics-informed, expert-guided, science-guided, scientific machine learning domain becomes standard, but hopefully this describes how all of the choices made here were not “better” or “worse”, but instead it’s all about domain-specific engineering trade-offs.

The post Engineering Trade-Offs in Automatic Differentiation: from TensorFlow and PyTorch to Jax and Julia appeared first on Stochastic Lifestyle.