If you have any experience with Async Rust, you must have come across Pin and pinning. Pinning is a complex topic and many Rust programmers find it confusing. If you feel uncomfortable around pinning, or like me, have had too many “why” questions about it, this article is for you!

All the articles (e.g., [1], [2], [3]) I have read on the topic of pinning state that the need for pinning is due to the existence of self-referential objects. However, that cannot be the only reason, as we see self-referential objects in other programming languages too, but nothing similar to Rust’s pinning.

Some other factor, unique to Rust, must be at play. In this article, we are aiming to understand this other factor. Along the way, we will answer the following questions:

1. How are mutability, move semantics, and movability related?
2. How does Pin work, and what memory safety guarantees does it provide?
3. When do you need to use Pin in your API?
4. How do other languages deal with self-referential objects?

Pin in a nutshell

Pin is a smart pointer. It is often said that Pin “prevents an object from being moved in memory, unless it is address-insensitive”. This article is all about deciphering this sentence, as it is not quite accurate, or rather, does not give the right impression when you first read it. Pin, alone, does not have a mechanism to keep an object in its place in memory, but is the way to express a guarantee about the object’s immovability. You, as the developer, may need to use other language constructs to achieve the stated guarantee.

The following diagram shows the structure of a Pin. We will come back to this diagram after clarifying some concepts, and will have a closer look at the three players in it when discussing Pin’s guarantees.

Movability

If I ask you what is movability, chances are, you’ll think of Rust’s move semantics and transfer of ownership. Here is a brief summary:

Question: Is ownership transfer the type of move that Pin expresses a guarantee about?

Not really. In the context of pinning, “move” and “movability” mean that the value has moved in the mechanical sense of being located at a new place in memory. In the following, I am going to refer to this as a mechanical move.

Two types of move

Question: Does a “move”, or transfer of ownership of a value, involve a mechanical move?

The short answer is “yes”. The longer answer is “yes, but sometimes only partially”.

Question: Can the location of a value in memory change without a transfer of ownership?

This one is trickier to answer. This is where the boundaries between transfer of ownership, mechanical move, and even mutability get blurred. To answer this question, it is best to look at some examples.

In the first example (playground), we create a vector with capacity 2, then push 4 elements to it. The elements are of type SelfPoint:

// Example 1:

struct SelfPoint {
    data: u32,
    ptr: *const u32,
}

impl SelfPoint {
    fn new(data: u32) -> Self {
        SelfPoint {
            data,
            ptr: std::ptr::null(),
        }
    }

    fn init(&mut self) {
        let ptr: *const u32 = &self.data;
        self.ptr = ptr;
    }
}

fn main() {
    let mut v = Vec::with_capacity(2);

    for i in 0..4 {
        let x = SelfPoint::new(i);
        v.push(x);
        v.last_mut().expect("vec is empty").init();
        println!("Vector addr & cap: {:p}, {}; heap location: {:p}", &v, v.capacity(), &v[0]);
    }
    ...
}

Each time we push an element to the vector, a move happens. This move involves both a transfer of ownership and a mechanical move (from the stack, to the call stack of push, and finally to the heap). After pushing each element to the vector, we call init on it to set the value of ptr to the address of data. Since we initialized the vector with capacity 2, pushing the third element will result in allocating a larger buffer, and moving the elements to the new buffer. The following diagram shows the memory layout of the vector, before and after pushing the third element to it.

Notice how after the third push, the ptr fields of the first two elements of the vector point to memory locations they no longer own. This is exactly the kind of problem that Pin is designed to prevent.

In this example, the memory allocated to the vector on the heap is moved, mechanically, but no explicit transfer of ownership has happened. Deep in the implementation of Vec, the vector’s ptr is updated and mutated, but even this is not something that I can call a transfer of ownership.

In the next example (playground), we initialize the vector with a larger capacity to prevent memory reallocation. We also pass the vector to a function causing a transfer of ownership.

// Example 2:

fn validate(v: Vec<SelfPoint>) {
    for x in &v {
        let data = &x.data as *const u32;
        assert!(data == x.ptr);
    }
    println!("Validated vector at {:p}; buffer at: {:p}", &v, v.as_ptr());
}

fn main() {
    let mut v = Vec::with_capacity(4);

    for i in 0..4 {
        let x = SelfPoint::new(i);
        v.push(x);
        v.last_mut().expect("vec empty").init();
    }
    
    println!("Created vector\t at {:p}; buffer at: {:p}", &v, v.as_ptr());
    validate(v);
}

The following diagram shows the memory layout in this case.

In both of the examples, vector v is partially moved to a new location in memory. In the first example, only the heap-allocated memory is moved, without any transfer of ownership. In the second example, only the stack-allocated memory is moved, involving a transfer of ownership.

With these examples in mind, we can now answer the question “Can the location of a value in memory change without a transfer of ownership?”

The answer is yes. In other words, mechanical move is a broader concept than ownership transfer.

Movability and mutability

The mechanical move in Example 1 was caused by a mutation. In fact, there is a very close connection between mutability and movability. More specifically, any mutation entails a mechanical move (see the classic example with mem::swap).

Due to this close connection between mutability and movability, in the description of Pin, we see many references to mutability, despite Pin being all about movability!

Note that mutating a value does not necessarily entail a transfer of ownership. Transfer of ownership involves marking a memory location as invalid, but in the case of mem::swap for instance, both memory locations remain valid after the operation.

Leaky Abstractions

While ownership and transfer of ownership are powerful abstractions, they are somewhat leaky, particularly when it comes to capturing the notion of mechanical moves. This is hinted at in the explanation of ownership in the Rust book. Moreover, as we saw in the use of Vec in Example 1, we may need to study the internal details of a type when deciding whether its manipulation can result in a mechanical move or not.

Having a leaky abstraction or only partially hiding details in a programming language is nothing new. Niklaus Wirth, in his 1974 paper “On the Design of Programming Languages”, writes:

I found a large number of programs perform poorly because of the language’s tendency to hide “what is going on” with the misguided intention of “not bothering the programmer with details”.

In my view, Rust programs are anything but “programs that perform poorly”, but knowing that an abstraction is leaky allows us to use it more effectively.

Address sensitivity

So far, we’ve learned that objects in Rust are movable in a mechanical sense, and that this kind of move is different from ownership transfer.

In some cases, the validity of an object’s invariants and the soundness of its behavior depend on its location in memory. We say that such an object is in an address-sensitive state. Self-referential objects (i.e., objects with pointers to themselves, or to fields in themselves) are the archetype of address-sensitivity. In the examples above, calling init on a SelfPoint instance made it self-referential and address-sensitive.

Strictly speaking, address-sensitivity is a state of an object, not necessarily a property that can be specified via its type (e.g., instances of SelfPoint are not address-sensitive before calling init on them). However, it is often a good practice to associate a different type to each state of a stateful entity. Following this practice, a self-referential object becomes address-sensitive early at the beginning of its lifetime, and remains address-sensitive until the end of its lifetime. In other words, once your object is no longer address-sensitive, it is a good practice to consume it and convert it into an instance of another type.

As a result, in practice, address-sensitivity can be specified via types. This allows the compiler to reason about the address-sensitivity of objects at compile time.

Why is pinning needed?

Once an object is in an address-sensitive state, the compiler can guarantee memory-safe interaction with it, only if the object stays in its location in memory. If an object is not address-sensitive, it can freely move around in memory, without violating any of the compiler’s safety guarantees.

To tell the compiler that instances of a type will not be address-sensitive, we implement the Unpin auto trait for that type. By default, all types are Unpin. To announce potential address-sensitivity of instances of a type, we implement !Unpin for it. We will learn more about Unpin below.

To promise to the compiler that an object has a fixed location in memory, or otherwise is address-insensitive, we wrap a pointer (i.e., pinning pointer) to that object (i.e., pointee) in a Pin. As we will see below, Pin restricts mutable access to a pointee that is !Unpin.

In short, pinning (i.e., the use of Pin and Unpin), are the means that you as a developer use to help the compiler provide its memory-safety guarantees, and avoid undefined behavior when it comes to address-sensitive types and objects.

The Unpin Trait

Unpin is a marker trait, and is automatically implemented for almost all types. It is implemented even for type SelfPoint in Example 1. This is because the compiler has no reason to assume that at runtime you are going to set SelfPoint.ptr to point to SelfPoint.data.

To tell the compiler that instances of SelfPoint will be self-referential and address-sensitive, you have to explicitly implement !Unpin for it. The conventional way of doing this is to add a field of type PhantomPinned to your type. PhantomPinned is !Unpin, and a type that contains a PhantomPinned does not get a default Unpin implementation. We change SelfPoint as follows:

// Example 3:

use std::marker::PhantomPinned;

struct SelfPoint {
    data: u32,
    ptr: *const u32,
    _pinned: PhantomPinned,
}

If your type has a field that is !Unpin (e.g., a tokio::time::Sleep field), and you want to make your type Unpin, you have to explicitly implement Unpin for it. Doing so is as simple as writing a single line of code.

impl Unpin for MyType {}

But you have to be very careful, because when you implement Unpin, you are promising to the compiler that instances of your type will be address-insensitive and freely movable.

Pin and its guarantees

Now that we know why pinning is needed, we can take a closer look at Pin, its API, and its guarantees.

The structure of a Pin

As mentioned above, to promise to the compiler that an object either has a fixed location in memory (i.e, is immovable) or is address-insensitive (i.e., is Unpin), we wrap a pointer to it in a Pin. This is shown in the diagram below (repeated from above).

• The object that we want to express the promise about is called the pointee.
• The pointer to pointee is called the pinning pointer.
• The pin itself is a smart pointer that wraps around the pinning pointer, and takes ownership of it. We say that pin pins the pointee.

Mutability, Movability, and Pin’s guarantees

Once you express your promise about the immovability or address-insensitivity of an object, by pinning it, Pin guarantees that it will maintain your promise. To achieve this, Pin won’t easily give mutable access to the pointee. This is important, because as we saw above, mutable access to a value can result in moving it in memory. For address-sensitive objects, this may violate the invariant of the object, or cause memory-safety issues.

APIs utilizing Pin, with Future::poll being a prominent example, are concerned with integrity and validity of the pointee’s invariants rather than its specific memory location. These APIs need to mutate the pointee, and need to know that doing so won’t result in undefined behaviors or a violation of the pointee’s invariants. To establish this proposition, immovability is used as an overapproximation – a sufficient and easier-to-verify condition that ensures the object’s integrity as it is being mutated.

By relying on immovability, Pin provides a simple and generic API that allows safe interaction with address-sensitive objects in arbitrary contexts without having to make any additional assumptions about the internal invariants of those objects.

Pinning address-insensitive objects

Instances of an Unpin type are promised to be address-insensitive, and therefore allowed to move freely in memory. Pin has a safe API for working with Unpin objects. The following table lists the methods in Pin’s safe API. Those in green require the pointee to be Unpin. The table shows the API from stable Rust 1.77.0.

If the pointee is Unpin, you can easily wrap it in a Pin, using the new method, without any additional restrictions. Moreover, if the pointee is Unpin, you can use the get_mut method to get mutable access to it. Mutating the pointee is safe because it is address-insensitive, and can move and mutate freely.

The methods in blue do not give direct mutable access to pointee, and cannot be used to mutate or move the object, without additional, possibly unsafe, code. Therefore they are safe.

Pinning address-sensitive objects

If instances of a type can enter an address-sensitive state, then that type must be declared !Unpin.

To soundly Pin instances of an !Unpin type, you have to promise that the pointee’s data will not be moved nor have its storage invalidated until it gets dropped. The compiler cannot check these promises, so when working with !Unpin types, you have to write unsafe code. Pin offers a number of unsafe methods particularly for working with !Unpin types.

To create an instance of Pin in this case, you have to use the unsafe new_unchecked method. The documentation for this method states a few additional promises that you have to make. These promises allow Pin to guarantee that your use of its safe API (colored in blue in the previous table) will be sound.

These additional promises are about the pinning pointer, which is the argument to new_unchecked. Needless to say, the pinning pointer must be Deref. You have to promise that in your implementation of Deref::deref, you do not “move [any data] out of self” (even a partial move violates this promise). You have to make a similar promise about your implementation of DerefMut if you have it. This is because Pin relies on Deref::deref and DerefMut::deref_mut in the implementation of some of its safe methods (e.g., as_ref and as_mut), and can only maintain the immovability promise, if Deref::deref and DerefMut::deref_mut guarantee it.

Note that inside DerefMut::deref_mut, you have mutable access to pointee. A malicious implementation could simply use swap to move a value out of self.

The methods map_unchecked and map_unchecked_mut are consuming methods that, among other things, are particularly useful for converting the pointee to an Unpin type, signaling its transition to an address-insensitive state. As an example, we could map instances of SelfPoint to address-insensitive u32 values holding only the content of the field data.

Putting it all together

We can use Pin and Unpin to fix the problem that we faced in Example 1. Here are the changes we have to make for the fix:

1. Change SelfPoint to be !Unpin as in Example 3.
2. Change init to take Pin<&mut Self> as its receiver. It only makes sense to make an instance of SelfPoint self-referential if we can guarantee that it is immovable. One way to do this is to move instances of SelfPoint to the heap.
3. Change the vector to store pinned references to SelfPoint instances.

Here is the final code, with these changes applied (playground).

// Example 4:

use std::marker::PhantomPinned;
use std::pin::Pin;

struct SelfPoint {
    data: u32,
    ptr: *const u32,
    _pinned: PhantomPinned,
}

impl SelfPoint {
    fn new(data: u32) -> Self {
        SelfPoint {
            data,
            ptr: std::ptr::null(),
            _pinned: PhantomPinned,
        }
    }

    fn init(self: Pin<&mut Self>) {
        let ptr: *const u32 = &self.data;
        let this = unsafe { self.get_unchecked_mut() };
        this.ptr = ptr;
    }
}

fn validate(v: Vec<Pin<Box<SelfPoint>>>) {
    for x in &v {
        let data = &x.data as *const u32;
        assert!(data == x.ptr);
    }
    println!("Validated vec:\t{:p}; buffer addr: {:p}; v[0] addr: {:p}", &v, v.as_ptr(), v[0]);
}

fn main() {
    let mut v = Vec::with_capacity(2);

    for i in 0..4 {
        let mut x = Box::pin(SelfPoint::new(i));
        x.as_mut().init();
        v.push(x);
        println!("{i} Vector addr:\t{:p}; buffer addr: {:p}; v[0] addr: {:p}", &v, &v[0], v[0]);
    }
    
    validate(v);
}

Do I need to use Pin in my API?

When I started digging into pinning, I wanted to be able to comfortably decide whether to use Pin in my APIs or not. This decision becomes relevant when designing generic APIs or traits, where we don’t know the exact types of the objects that will interact with the API.

In these cases, you might want to use Pin, if:

1. Your API is likely to work with address-sensitive types and objects.
   1. For instance, types that implement the Future trait are likely to be self-referential, and therefore address-sensitive. APIs that work with a Future need to Pin it.
   2. On the other hand, elements stored in a Vec don’t have that tendency, and more often than not are address-insensitive.
2. Your API needs to mutate objects.
   1. For instance, Future::poll is called to make progress on a Future object. This requires changing, or mutating, the state of the Future object.
   2. On the other hand, in its implementation, Vec does not mutate the elements in its buffer.
3. The safety of your API depends on the object’s integrity and the validity of its invariants.
   1. For instance, as a safe method, calls to Future::poll “must never cause undefined behavior (memory corruption, incorrect use of unsafe functions, or the like), regardless of the future’s state.”
   2. On the other hand, while most of the Vec’s methods are safe, its contract is not affected by the invariants of the elements in its buffer. The presence of address-sensitive objects in a Vec does not cause undefined behavior when invoking any of its methods.

If your API passes all these checks, you’ll need Pin in your API. Passing these checks implies that your API depends tightly on the inner working of the types that it interacts with. Using Pin hides these details behind an immovability promise. This promise, although possibly more conservative than what you want, is easier to reason about.

Otherwise, even if you have a general purpose API like Vec, you won’t need Pin in your API; and your API will work just fine with pinned objects too.

Before wrapping up this section, I want to emphasize again that Pin does not force the pointee to stay in its location. It merely expresses an immovability promise about it. While all objects in Rust are movable, most objects never leave their original locations in memory throughout their lifetimes. For instance, most heap-allocated objects never move, as we saw in Example 4.

How do other languages deal with self-referential objects?

To answer this question, let’s start with Java, as a representative of garbage-collected languages.

In Java, all objects are heap-allocated. In your programs, you work with references to these objects, and can easily make self-referential objects. Consider the following example, which declares class SelfRef.

// Example 5:
class SelfRef {
   SelfRef self;

   void init() {
       this.self = this;
   }
}

You can make instances of SelfRef self-referential by calling init on them. But how about moving instances of SelfRef in memory? Well, the language does not provide you with a mechanism to do that. The notion of movability, in the sense of manually relocating objects in memory, does not exist in Java or similar garbage-collected languages. In these languages, objects reside safely in the heap, each hiding behind a reference. However, this does not mean that objects in these languages never relocate in the memory. The runtime, more specifically the garbage collector, does move objects around. This happens specifically when defragmenting or compacting the heap after removing the garbage collected objects. The garbage collector manages this process by tracking references to ensure their validity.

The same is true for other garbage-collected languages, such as Python. An exception worth mentioning is Golang, which (as opposed to Java and Python) allows pointer types. In Golang, you can declare a type similar to SelfPoint in Example 1, create self-referential instances of it, and move them around in memory. I am not aware of any mechanism in Golang for preventing undefined behaviors in these situations.

In C++, as far as I understand, you can move objects around. However, the language allows you to custom-implement move constructors. I suppose, move constructors are where you can ensure that invariants of your type are maintained as your objects move around in memory.

Conclusion

In my quest for understanding Pin, I came to the conclusion that it is a construct that sits at the intersection of a few leaky abstractions, and provides a simple and generic mechanism for expressing invariants about address-sensitive objects without depending on their internal details.

The main leaky abstraction is the concept of moving a value. You want to think of it as a transfer of ownership, but it is more complex than that. Sometimes you don’t have an explicit transfer of ownership, but parts of the value move to another location in memory. Sometimes you have an explicit transfer of ownership, but parts of the value (e.g., the heap-allocated parts) stay in their memory location.

These subtleties become important when working with address-sensitive objects, whose invariants may be violated if their location in memory changes.

It is important to note that it is not the location of these objects that we care about, but their integrity and the validity of their invariants. However, it is generally easier to reason about the movability of an object rather than its invariants. Pin and Unpin are the constructs that we use for expressing movability guarantees.

Despite some unsoundness issues in earlier versions of Pin, in my assessment, Pin and Unpin form a neat and effective abstraction that conceals the complexities of memory management that are leaked via the concepts of ownership and move semantics. I am curious about their origin and evolution. The story is certainly somewhere in the discussions or on Github, if you go looking!

This article is a follow up on my earlier article titled AsyncWrite and a Tale of Four Implementations, which goes into an in-depth discussion of the decisions you have to make when working with async and Futures.

References

1. Pinning in the Async Rust book
2. Pin and Suffering
3. Pin, Unpin, and why Rust needs them
4. The vision for Rust specification from the Inside Rust Blog
5. Nightly Rust documentation of pinning
6. Rust internals

What is fascinating is that Helen, being deafblind, communicated primarily through touch. She could not remember reading the original story herself, in braille. So somebody must have “read” it to her, by spelling the words on her fingers.

Following the incident, Helen’s teacher, Miss Sullivan, with the assistance of Alexander Graham Bell, finds out that the original story must have been “read” to Helen, by a friend, 3-4 years earlier, in 1888, during a vacation. At that time, Helen had been “learning the spelling of strange words”, and stories did not have much meaning to her.

Later, Helen’s own analysis of her letters and writings revealed to her that her mind was saturated with the story and other ideas of the original author, Miss Canby.

The New York Times lawsuit against OpenAI and Microsoft reminded me of this story.

I wonder what the lawsuit means in today’s ethical landscape. Occasionally, I come across articles on the internet that seem to be instances of intentional plagiarism. Between that and the mysteries of how our brains work, how can one tell what the original source of an idea is?

In Rust, we often don’t need to work directly with futures or implement low-level async APIs, but sometimes they are inevitable. In this post, I report and reflect on my experience implementing an async API. I am going to compare four different implementations, and walk you through the use of futures in a non-trivial example. I will discuss the “when” and “why” of using pin_project.

The following table summarizes the four alternative implementations, and includes references to the corresponding code in the associated GitHub repository.

In each alternative, we will implement AsyncWrite for a struct. In the first alternative, the struct does not hold any internal state. Inside the AsyncWrite implementation, we call an async function, which, as we will see, requires pinning. In the second implementation, we do the same for a struct that maintains internal state. We will see how some refactoring can help us address a class of lifetime issues. The third variant wraps the async call inside a custom Future. Thus, the AsyncWrite implementation does not have to deal directly with the async call. Here, we will observe the difference between Unpin and !Unpin types. Finally, in the fourth variant, we will completely remove the async call, and will additionally discuss the use of pin_project.

These four scenarios are not the only imaginable combinations. However, not all combinations are interesting or feasible. For example, the last row in the table represents a combination that won’t work without resorting to unsafe code, potentially introducing unsoundness.

Why custom implement AsyncWrite?

Having a rough idea about how futures work is different from actually using them to build a complex working system! A while ago, I had to work with futures at a low level for the first time. With a lot of cursory readings here and there, particularly about pinning, I ended up with a working solution 🎉 Looking at that solution now, after a few months, I see that my solution could have been simplified, if I had a deeper understanding of some concepts. Notably, in my solution, I had used pin_project unnecessarily!

Before diving into implementing AsyncWrite, let me tell you about my use case. In my project, I needed to compress large amounts of data and upload the compressed content to AWS S3. For S3, I used the aws_sdk_s3 crate, which is async and relies on tokio. For compression, I used ZstdEncoder`, from async_compression.

Async was chosen first and foremost to achieve better performance, but also because I could not find a reliable alternative for aws_sdk_s3.

ZstdEncoder has a suspiciously simple API. It took me a while to understand how it works. Looking at the examples was necessary for an aha moment!

The API for ZstdEncoder says:

This structure implements an AsyncWrite interface and will take in uncompressed data and write it compressed to an underlying stream.

ZstdEncoder provides a few methods for creating instances of it (e.g., new). These methods take an AsyncWrite object (inner) as input! This is the underlying stream that the compressed data is written into.

To feed uncompressed data into an instance of ZstdEncoder, the most convenient and efficient way is to wrap it in a BufWriter and call write_all, or any other method from tokio::io::AsyncWriteExt on it.

To retrieve the compressed data from ZstdEncoder, you must find a way to access the underlying inner stream that you provided at the creation time. ZstdEncoder takes ownership of the inner stream, and during the operation, only gives you immutable access to it. With this restricted access, and considering the available AsyncWrite types (e.g., Vec) as the underlying stream, you cannot consume the compressed data, in chunks, during the operation. In other words, to consume the compressed data, you have to wait until the end of the operation, and call into_inner on the ZstdEncoder instance, to gain ownership of the inner stream and access the compressed data. The following snippet shows this scenario.

async fn compress(in_data: &[u8]) -> Result<Vec<u8>> {
    let mut encoder = ZstdEncoder::new(Vec::new());
    encoder.write_all(in_data).await?;
    encoder.shutdown().await?;
    Ok(encoder.into_inner())
}

For a small example this works fine, but if you are using async, particularly in a streaming manner, you are most likely working with large amounts of data. You cannot simply wait until the end of the process to examine or use the results.

This was exactly the case in my project. I was processing terabytes of data, and needed to upload the compressed results to S3. More specifically, I wanted an end-to-end in-memory streaming solution. As the input stream of uncompressed data came in, it would be compressed, split into chunks and uploaded to S3 (using its multi-part upload API, with create, upload, and complete methods).

Despite my best efforts to avoid implementing AsyncWrite, it proved inevitable as no combination of existing structs and APIs could fulfill the task.

Implementing AsyncWrite

Let’s start with a quick look at AsyncWrite’s API:

pub trait AsyncWrite {
    // Required methods
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8]
    ) -> Poll<Result<usize, Error>>;

   fn poll_flush(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>
    ) -> Poll<Result<(), Error>>;

   fn poll_shutdown(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>
    ) -> Poll<Result<(), Error>>;

    // Provided methods
    fn poll_write_vectored(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &[IoSlice<'_>]
    ) -> Poll<Result<usize, Error>> { ... }

   fn is_write_vectored(&self) -> bool { ... }
}

The first thing you notice about this API is that you need to work with pinned objects! A large part of this post is dedicated to discussing pinning options. From this API, methods that need to be implemented are poll_write, poll_flush, and poll_shutdown. My focus will be on poll_write.

In this article, I am using tokio::io::AsyncWrite; futures::io::AsyncWrite is very similar.

AsyncWrite for a stateless struct

Let’s say we have an AsyncUploader struct that can asynchronously upload content to a remote server. We want to use this uploader in contexts where an AsyncWrite is needed, and will implement AsyncWrite for it. This uploader could be as simple as the following struct:

// Impl-1.a: pinned.rs

struct AsyncUploader {}

impl AsyncUploader {
   fn new() -> Self {
       Self {}
   }

   async fn upload(&self, _data: &[u8]){
       // TODO: Upload the given bytes.
   }
}

Inside upload, I am going to call some async functions (for instance, functions from aws_sdk_s3), so it has to be async. But for now, we will leave the body of upload empty.

To see AsyncUploader in action, I will use the following simple caller code.

// Impl-1.b: pinned.rs

#[tokio::main]
async fn main() {
   let mut uploader = AsyncUploader::new();

   // In reality, we will read from an input stream instead of this for loop!
   for i in 0..10 {
       let msg = format!("{i}: Hello, world!\n");
       uploader.write_all(msg.as_bytes()).await.unwrap();
   }
   uploader.shutdown().await.unwrap();
}

Futures under the hood

The most important part of the code above is the call to write_all (from AsyncWriteExt). Internally, write_all calls poll_write (from AsyncWrite). The following sequence diagram shows the objects that are involved in executing write_all, including uploader and a Future on which poll_write and poll are called, respectively.

This diagram is simplified and not very accurate. I have left out the details of tokio workers, and tasks. This is because in our simple example, we have only one task, and one worker. Adding them to the diagram is not worth the added clutter. The sequence diagram does not cover all different paths corresponding to the return value of poll_write, or the termination condition for the loop. But I hope it helps with understanding the big picture.

As shown in the diagram, in the main function, the call to write_all on the uploader results in the creation of a Future of type WriteAll. This Future is then executed by a worker in the tokio runtime. The runtime periodically calls poll on the Future, until it is ready. Each call to poll on WriteAll results in a call to poll_write (in more general cases where a buffer is involved, it may result in multiple calls to poll_write). This continues until poll_write returns Ready. This completes one iteration of the for loop in the main function. Note that the call to poll_write may sometimes return Ready immediately. In that case, we don’t really enter the loop in the sequence diagram.

Now that we’ve seen the big picture, let’s implement AsyncWrite for AsyncUploader, with a focus on poll_write. Inside poll_write, one way or another, we have to call AsyncUploader::upload, which is an async function. Normally, we’d call .await on an async function (e.g., self.upload(buf).await), but poll_write is not an async function. So inside it, we cannot use .await! Instead, we have to create a Future from the call to upload, and call poll on it. This is possible, because the call to an async function (in this case upload) returns a Future.

Futures and pinning

Recall that Future is a trait with the following API:

pub trait Future {
    type Output;

    // Required method
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>;
}

Note that poll takes a pinned mutable pointer to Self as its receiver. Pin is a smart pointer that wraps another pointer in itself, and expresses a guarantee about the immovability of the object (pointee in the following diagram) that the inner pointer points to. Pin also restricts mutable access to the pointee.

To be able to call poll on the Future returned from the call to upload, we have to first pin a pointer to that Future. Here is an attempt at it:

// Impl-1.c: pinned.rs

impl AsyncWrite for AsyncUploader {
   fn poll_write(
       self: std::pin::Pin<&mut Self>,
       cx: &mut Context<'_>,
       buf: &[u8],
   ) -> Poll<Result<usize, Error>> {
       Let mut fut  = self.upload(buf);
       let mut future = Pin::new(&mut fut);
       // TODO: Use the future.
       Poll::Pending 
   }
}

Compiling this results in the following error:

error[E0277]: `{async fn body@src/bin/main.rs:20:41: 22:6}` cannot be unpinned
   --> src/bin/pinned.rs:45:35
    |
20  |     async fn upload(&self, _buf: &[u8]) {
    |     ----------------------------------- within this `impl futures::Future<Output = ()>`
...
45  |         let mut future = Pin::new(&mut fut);
    |                          -------- ^^^^^^^^ within `impl futures::Future<Output = ()>`, the trait `Unpin` is not implemented for `{async fn body@src/bin/pinned.rs:20:41: 22:6}`
    |                          |
    |                          required by a bound introduced by this call
    |
    = note: consider using the `pin!` macro
            consider using `Box::pin` if you need to access the pinned value outside of the current scope
note: required because it appears within the type `impl Future<Output = ()>`
   --> src/bin/pinned.rs:20:5
    |
20  |     async fn upload(&self, _buf: &[u8]) {
    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
note: required by a bound in `Pin::<P>::new`
   --> /Users/razieh/.rustup/toolchains/nightly-x86_64-apple-darwin/lib/rustlib/src/rust/library/core/src/pin.rs:484:23
    |
484 | impl<P: Deref<Target: Unpin>> Pin<P> {
    |                       ^^^^^ required by this bound in `Pin::<P>::new`
...
503 |     pub const fn new(pointer: P) -> Pin<P> {
    |                  --- required by a bound in this associated function

The compiler is telling us that the Future created from the call to upload cannot be unpinned! Meaning that this Future does not implement Unpin. Unpin is a marker trait that tells the compiler that an object does not care about its location in memory. Most types are Unpin. A prominent exception are Future types created from calling async functions.

We are getting this compiler error because the input to Pin::new has to be a pointer to a pointee that is Unpin, but the Future we have is not Unpin.

We have a few solutions here. If we look at Pin’s documentation, we see another method, new_unchecked, for creating Pin that does not require Unpin. However, this method is unsafe. Let’s avoid that for now. The other alternative is Box::pin, which is suggested by the compiler.

// Impl-1.d: pinned.rs

use futures::Future;
...   
impl AsyncWrite for AsyncUploader {
   fn poll_write(
       self: std::pin::Pin<&mut Self>,
       cx: &mut Context<'_>,
       buf: &[u8],
   ) -> Poll<Result<usize, Error>> {
       let mut fut = Box::pin(self.upload(buf));


       match fut.as_mut().poll(cx) {
           Poll::Pending => Poll::Pending,
           Poll::Ready(_) => Poll::Ready(Ok(buf.len())),
       }
   }
}

Now that we have a pinned Future, we can call poll on it. Note that we have to call as_mut to get a pinned mutable reference to the Future object. This is because fut is a Pin<Box<...>>; not exactly the type we can call poll on.

This implementation creates a new Future every time poll_write is invoked. This is not at all useful! For a Future to advance, one must regularly invoke poll on it. In this setup, however, we only call poll once on each Future created inside poll_write! These Futures go out of scope and get destroyed when we return from poll_write. We never get to poll on them again or make any progress with them!

In fact, if you try this implementation with an upload function that implements a long-running task (as below), you will see that the program never terminates, and never produces any output!

// Impl-1.e: pinned.rs

struct AsyncUploader {}

impl AsyncUploader {
   ...
   async fn upload(&self, data: &[u8]) -> Result<(), Error> {
      // Sleep for a bit to simulate a long-running operation.
      tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
      println!(
         "Uploaded {}",
         String::from_utf8(data.to_vec())
           .map_err(|e| Error::new(ErrorKind::Other, format!("{e}")))?
     );
     Ok(())
   }
}

For simplicity, and to avoid getting into complications of interacting with a remote server, in this implementation of upload, I have used sleep (the non-blocking version from tokio) to simulate a long-running task.

The execution of main gets stuck, because with this implementation of upload, poll_write never returns anything other than Poll::Pending!

A long-living Future

To fix this, we have to change our design to give the Future created from the call to upload a longer lifetime, so that we can keep track of its progress, and repeatedly call poll on it. To do this, we have to bind the Future to the AsyncUploader instance, by declaring the Future as a field. The following is the updated sequence diagram.

This diagram is again not very precise. It only shows the scenario where the first call to poll_write returns Poll::Pending, which is not always the case. In some cases, the call to write_all may complete with a single call to poll_write that immediately returns Poll::Ready.

To implement this design, we have to add a new field (which we call fut) to AsyncUploader. But what exactly should be the type of this field?

We know the following:

1. It is a Future, and poll will be called on it. Since poll’s receiver is a Pin, one way or another, we have to pin this Future.
2. The Future is the result of calling an async function. We don’t know the concrete type of this Future, and have to refer to it as dyn Future<Output = T>, where T is the return type of the upload method.
3. When we create an instance of AsyncUploader, we don’t yet have an instance of this Future. To handle this case, fut has to be an Option.
4. Since dyn Future<Output = T> is not sized, we have to keep a pointer to it. It is more convenient, and more ergonomic to use Box and keep the Future on the heap.

There is a type alias for Pin<Box<dyn Future<Output = T>>>, called BoxFuture, which is exactly the type we need here. In addition to T, BoxFuture takes a lifetime parameter, which we can set to ‘static.

Putting all this together, we update AsyncUploader as follows:

// Impl-1.f: pinned.rs

struct AsyncUploader {
   fut: Option<BoxFuture<'static, Result<(), Error>>>,
}

With this new field, inside poll_write, we first check if self.fut is None. If that is the case, we (1) call upload, passing the content of the input buf to it, (2) update the value of self.fut with the resulting Future, and (3) return Poll::Ready. Since we have consumed the incoming bytes in buf, it is important to return Poll::Ready. If we already have a Future in fut, then we poll on it. If the result of poll is Poll::Ready with success, we do as we did above: call upload again, update fut, and return Poll::Ready. Otherwise, we either return Poll::Pending, or Poll::Ready with an error. This implementation is shown below.

// Impl-1.g: pinned.rs

impl AsyncWrite for AsyncUploader {
   fn poll_write(
       mut self: std::pin::Pin<&mut Self>,
       cx: &mut Context<'_>,
       buf: &[u8],
   ) -> Poll<Result<usize, Error>> {
       if self.fut.is_none() {
           self.fut = Some(Box::pin(self.upload(buf)));
           // The future we created is not necessarily completed yet, but
           // we successfully used all the input bytes, and are now ready
           // to accept new bytes. The poll_write contract requires us to
           // return Ready.
           return Poll::Ready(Ok(buf.len()));
       }

       // We have an on-going Future. We can only upload the given bytes
       // in buf, and create a new Future, if the previous one is
       // complete. Otherwise, we return Pending or Ready with an error.
       match self.fut.as_mut().unwrap().as_mut().poll(cx) {
           Poll::Pending => Poll::Pending,
           Poll::Ready(Ok(())) => {
               self.fut = Some(Box::pin(self.upload(buf)));
               Poll::Ready(Ok(buf.len()))
           }
           Poll::Ready(Err(err)) => {
               let _ = self.fut.take();
               Poll::Ready(Err(err))
           }
       }
   }
}

If we run this, we get the following error:

error: future cannot be sent between threads safely
  --> src/bin/pinned.rs:61:29
   |
61 |             self.fut = Some(Box::pin(self.upload(buf)));
   |                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^ future returned by `upload` is not `Send`
   |
   = help: the trait `Sync` is not implemented for `(dyn futures::Future<Output = Result<(), std::io::Error>> + std::marker::Send + 'static)`
note: captured value is not `Send` because `&` references cannot be sent unless their referent is `Sync`
  --> src/bin/pinned.rs:35:9
   |
35 |         &self,
   |         ^^^^^ has type `&AsyncUploader` which is not `Send`, because `AsyncUploader` is not `Sync`
   = note: required for the cast from `Pin<Box<impl futures::Future<Output = Result<(), std::io::Error>>>>` to `Pin<Box<dyn futures::Future<Output = Result<(), std::io::Error>> + std::marker::Send>>`

The error says that the Future that we have wrapped inside BoxFuture is not Send. One might be tempted to fix this by declaring a new type based on BoxFuture, but without the Send trait bound:

pub type BoxFut<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;

If we replace BoxFuture with BoxFut in the declaration of fut, we make the compiler very unhappy:

error[E0502]: cannot borrow `self` as mutable because it is also borrowed as immutable
  --> src/bin/pinned.rs:61:13
   |
61 |             self.fut = Some(Box::pin(self.upload(buf)));
   |             ^^^^            ---------------------------
   |             |               |        |
   |             |               |        immutable borrow occurs here
   |             |               cast requires that `self` is borrowed for `'static`
   |             mutable borrow occurs here

error: lifetime may not live long enough
  --> src/bin/pinned.rs:61:29
   |
53 |         mut self: std::pin::Pin<&mut Self>,
   |                                 - let's call the lifetime of this reference `'1`
...
61 |             self.fut = Some(Box::pin(self.upload(buf)));
   |                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^ cast requires that `'1` must outlive `'static`

error[E0597]: `self` does not live long enough
  --> src/bin/pinned.rs:61:38
   |
53 |         mut self: std::pin::Pin<&mut Self>,
   |         -------- binding `self` declared here
...
61 |             self.fut = Some(Box::pin(self.upload(buf)));
   |                             ---------^^^^--------------
   |                             |        |
   |                             |        borrowed value does not live long enough
   |                             cast requires that `self` is borrowed for `'static`
...
84 |     }
   |      - `self` dropped here while still borrowed

error[E0502]: cannot borrow `self` as mutable because it is also borrowed as immutable
  --> src/bin/pinned.rs:76:17
   |
76 |                 self.fut = Some(Box::pin(self.upload(buf)));
   |                 ^^^^            ---------------------------
   |                 |               |        |
   |                 |               |        immutable borrow occurs here
   |                 |               cast requires that `self` is borrowed for `'static`
   |                 mutable borrow occurs here

error[E0597]: `self` does not live long enough
  --> src/bin/pinned.rs:76:42
   |
53 |         mut self: std::pin::Pin<&mut Self>,
   |         -------- binding `self` declared here
...
76 |                 self.fut = Some(Box::pin(self.upload(buf)));
   |                                 ---------^^^^--------------
   |                                 |        |
   |                                 |        borrowed value does not live long enough
   |                                 cast requires that `self` is borrowed for `'static`
...
84 |     }
   |      - `self` dropped here while still borrowed

The main problem is that in our implementation, we are violating one of Rust’s most important memory management rules. We are borrowing self, simultaneously, as both mutable and immutable. Additionally, we have a lifetime issue. This comes from the fact that (1) we declared fut to have a static lifetime, then (2) bound its lifetime to that of self (by setting it to be the result of self.upload), which (3) is dropped at the end of poll_write.

Let’s backtrack the change from BoxFuture to BoxFut. In fact, the Send trait bound is required on BoxFuture because in a multithreaded context, calls to Future::poll may run on different threads. But the borrowing and lifetime issues are genuine problems that we have to fix.

Refactoring to the rescue

We are getting the “future cannot be sent between threads safely” error because the Future we are creating by the call to upload is capturing a reference to self, which is not Sync. Luckily, inside upload, we never use self. So we can make upload an ordinary function, and remove self from its signature. Let’s do this for now, and call the new function upload_bytes. I’ve also changed the type of data from &[u8] to Vec<u8> to solve a similar lifetime problem. Note that this kind of refactoring might not be as straightforward in more complex situations, but is often the cleanest way to address lifetime issues when working with Futures from async calls.

Here is our upload_bytes function:

// Impl-1.h: pinned.rs

/// This is the slow upload function, which we want to use inside the
/// AsyncWrite implementation.
async fn upload_bytes(data: Vec<u8>) -> Result<(), Error> {
   // Sleep for a bit to simulate a long-running operation.
   tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
   println!(
       "Uploaded {}",
       String::from_utf8(data)
           .map_err(|e| Error::new(ErrorKind::Other, format!("{e}")))?
   );
   Ok(())
}

With this change, we can implement poll_write as follows.

// Impl-1.i: pinned.rs

impl AsyncWrite for AsyncUploader {
   fn poll_write(
       mut self: std::pin::Pin<&mut Self>,
       cx: &mut Context<'_>,
       buf: &[u8],
   ) -> Poll<Result<usize, Error>> {
       if self.fut.is_none() {
           // We need a vec to avoid lifetime issues.
           let data = buf.to_vec();
           self.fut = Some(Box::pin(upload_bytes(data)));
           // The future we created is not necessarily completed yet, but
           // we successfully used all the input bytes, and are now ready
           // to accept new bytes. The poll_write contract requires us to
           // return Ready.
           return Poll::Ready(Ok(buf.len()));
       }

       // We have an on-going Future. We can only upload the given bytes
       // in buf, and create a new Future, if the previous one is
       // complete. Otherwise, we return Pending or Ready with an error.
       match self.fut.as_mut().unwrap().as_mut().poll(cx) {
           Poll::Pending => Poll::Pending,
           Poll::Ready(Ok(())) => {
               self.fut = Some(Box::pin(upload_bytes(buf.to_vec())));
               Poll::Ready(Ok(buf.len()))
           }
           Poll::Ready(Err(err)) => {
               // Clean up fut to avoid a "resumed after completion"
               // panic, if for whatever reason the caller decides to
               // ignore this error and proceed with another write or a
               // call to shutdown.
               let _ = self.fut.take();
               Poll::Ready(Err(err))
           }
       }
   }
}

Note that I have made the first argument of poll_write (self: std::pin::Pin<&mut Self>) mut! This is allowed, and not a violation of the AsyncWrite::poll_write’s signature, because we are consuming self, and can mutate it anyway. Equivalently, we could redeclare self as mutable in the body of the method. We need self to be mutable, because, we are modifying its fut, and polling on it, which requires mutable access to fut and therefore self.

For this implementation to work correctly, we need a similar code in poll_shutdown, except that in poll_shutdown, we don’t call upload_bytes, as we don’t have any bytes to upload! See the final code on github.

If we run the main function, we can see that it prints out ten “Uploaded {i}: Hello World!” messages, with one-second delays between them!

Now that we have this simple AsyncUploader working, let’s make it more interesting. Specifically, the AsyncUploader currently does not have any internal state. In reality, we need to work with structs that have some internal state.

AsyncWrite for a stateful struct

In my project, I had to upload content to S3. To speed up the process, I used S3’s multi-part upload API, which allows splitting a file (or object) into parts, and uploading them concurrently. S3 multipart upload requires each part to be at least 5M bytes, except for the last part. Addressing this requirement in the implementation can be achieved in various ways. One approach is to task AsyncUploader with ensuring that this condition is met. In such a case, AsyncUploader will invoke the upload_bytes method only if it has a minimum number of bytes to upload. Otherwise, it will simply append the bytes to an internal buffer, which is the internal state we are going to add to AsyncUploader:

// Impl-2.a: stateful.rs

struct AsyncUploader {
   /// Minimum size of the chunks to upload.
   min_size: usize,
   /// Internal buffer.
   buf: Vec<u8>,
   /// Future corresponding to the long-running upload operation.
   fut: Option<BoxFuture<'static, Result<(), Error>>>,
}


impl AsyncUploader {
   fn new(min_size: usize) -> Self {
       Self {
           min_size,
           buf: Vec::with_capacity(2 * min_size),
           fut: None,
       }
   }
   ...
}

I’ve initialized the buffer with a capacity twice the minimum size. This is just a heuristic to avoid reallocating the buffer too many times, as we add bytes to it.

To make poll_write more readable after this change, I refactored it slightly, and moved the main logic to an upload method, which (1) adds the incoming bytes to the inner buffer if it does not have enough bytes, (2) calls uploade_bytes if it has enough bytes, and (3) updates fut.

// Impl-2.b: stateful.rs

impl AsyncUploader {
   ...
   /// Precondition: This method should only be called if self.fut is None
   /// or contains a completed future.
   fn upload(&mut self, buf: &[u8]) -> usize {
       // Here we could check that the stated precondition about self.fut
       // holds. Checking the completion of the Future in self.fut is only
       // possible via a call to poll, which requires passing to it a
       // Context object. Doing that requires changing the signature of
       // the method, and passing to it an additional Context object.

       let mut size = 0;
       if self.buf.len() < self.min_size {
           self.buf.extend_from_slice(buf);
           size = buf.len();
       }

       if self.buf.len() >= self.min_size {
           let data = self.buf.drain(..).collect();
           self.fut = Some(Box::pin(upload_bytes(data)));
       } else {
           // It is safe to set self.fut to None because of the
           // precondition. Without setting fut to none, we'd continue
           // polling on it after completion, which results in a panic.
           self.fut = None;
       }

       size
   }
}

There are a few things to note about this implementation:

1. The receiver of upload is &mut self; as simple as that! We did not have to pin it, because we are not calling poll in this method.
2. There is a precondition for calling upload, as stated in the doc comment. According to this precondition, upload can only be called if self.fut is None, or its inner Future is Ready. I added this condition to keep the logic simple, without violating the contract of poll_write.

The interested reader is encouraged to modify the code, remove the precondition, and make the implementation more performant. But make sure not to violate the contract of poll_write when you call self.upload. Namely, if bytes are used and added to the inner buffer, poll_write must not return Pending.

Here is the updated poll_write, which calls self.upload, after checking that the precondition holds:

// Impl-2.c: stateful.rs

impl AsyncWrite for AsyncUploader {
   fn poll_write(
       mut self: std::pin::Pin<&mut Self>,
       cx: &mut Context<'_>,
       buf: &[u8],
   ) -> Poll<Result<usize, Error>> {
       if self.fut.is_none() {
           let size = self.upload(buf);
           // We successfully added all input bytes to the inner buffer,
           // and are now ready to accept new bytes. The poll_write
           // contract requires us to return Ready.
           return Poll::Ready(Ok(size));
       }

       // We have an on-going future. We can only upload the given bytes
       // in buf, and create a new Future, if the previous future
       // is complete. Otherwise, we return Pending or an error.
       match self.fut.as_mut().unwrap().as_mut().poll(cx) {
           Poll::Pending => {
               // For simplicity, we don't use the incoming bytes if we
               // have an on-going future. For a more performant
               // implementation, one should add incoming bytes to the
               // inner buffer as long as its capacity allows, and only
               // return Pending if the buffer is full, and the Future in
               // fut is not yet complete.
               Poll::Pending
           }
           Poll::Ready(Ok(())) => {
               let size = self.upload(buf);
               Poll::Ready(Ok(size))
           }
           Poll::Ready(Err(err)) => {
               // Clean up fut to avoid a "resumed after completion"
               // panic, if for whatever reason the caller decides to
               // ignore this error and proceed with another write or a
               // call to shutdown.
               let _ = self.fut.take();
               Poll::Ready(Err(err))
           }
       }
   }
}

In this version, we are invoking self.upload inside poll_write without a problem. Inside self.upload, the Future in self.fut is updated, but without holding a reference to self. We achieved this by draining the inner buffer to a separate vector before calling upload_bytes. This way, we removed the restriction on the lifetime of self imposed by that of the Future inside self.fut.

Now that we have an inner buffer, we have to update poll_shutdown to upload any remaining bytes in the internal buffer, before termination. Normally, with an inner buffer, you’d have to flush the buffer in poll_flush too. But in our case, flushing the inner buffer, if we don’t have enough bytes, will violate the minimum upload size condition. So we leave the implementation of poll_flush empty.

If we again run the main function after these changes, we see that this time our “Uploaded {i}: Hello world!” messages are written in batches of 8, due to the internal buffering.

Custom Future to hide Future from async

In both of the examples above, the main logic happens inside an async function (i.e., upload_bytes). The difficulty with async functions is that they result in unnamed Future types, whereas it is often easier to work with concrete known types. In this example, we are going to introduce a custom Future type, and push the async function call away from our AsyncUploader!

A custom Future type

In the following snippet, I have changed upload_bytes to return a custom Future named UploadFuture:

// Impl-3.a: pinless.rs

struct UploadFuture {
   /// Inner Future, resulted from an async call.
   fut: BoxFuture<'static, Result<(), Error>>,
}

impl Future for UploadFuture {
   type Output = Result<(), Error>;
   fn poll(...) -> Poll<Self::Output> {...}
}

/// This is the slow upload function, which we want to use inside the
/// AsyncWrite implementation.
fn upload_bytes(data: Vec<u8>) -> UploadFuture {
   let fut = async {
       // Sleep for a bit to simulate a long-running operation.
       tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
       println!(
            "Uploaded {}",
           String::from_utf8(data).map_err(|e| {
               Error::new(ErrorKind::Other, format!("{e}"))
           })?
       );
       Ok(())
   };

   UploadFuture { fut: Box::pin(fut) }
}

I have simply wrapped the entire logic in an async-block, and have wrapped the Future from the async-block in an instance of UploadFuture, which is what is returned from upload_bytes. So far, I have only introduced some nesting in the structure!

To complete the nesting, we have to update the declaration of fut in AsyncUploader:

// Impl-3.b: pinless.rs

struct AsyncUploader {
   /// Minimum size of the chunks to upload.
   min_size: usize,
   /// Internal buffer.
   buf: Vec<u8>,
   /// Future corresponding to the long-running upload operation.
   fut: Option<UploadFuture>,
}

Additionally, we have to tweak a few lines in AsyncUploader::upload. In particular, we have to change self.fut = Some(Box::pin(upload_bytes(data))); to self.fut = Some(upload_bytes(data));.

Disentangled lifetimes

One interesting thing we can do now is to make upload_bytes a method in AsyncUploader, with &mut self as its receiver:

// Impl-3.c: pinless.rs

impl AsyncUploader {
   ...
   fn upload_bytes(&mut self, data: Vec<u8>) -> UploadFuture {
       let fut = async {/* same as before */ };
       UploadFuture { fut: Box::pin(fut) }
   }
}

Recall that in Impl-2, achieving this was not possible with the async version of upload_bytes. Our attempt led to an error stating, “future cannot be sent between threads safely”. By encapsulating the result of the async block in a custom Future type, we have effectively decoupled self from the Future that we have to poll on, thus eliminating the error.

To complete the code, we have to implement Future for UploadFuture:

// Example 3.d: pinless.rs

impl Future for UploadFuture {
   type Output = Result<(), Error>;
   fn poll(
       mut self: Pin<&mut Self>,
       cx: &mut Context<'_>,
   ) -> Poll<Self::Output> {
       match self.fut.as_mut().poll(cx) {
           Poll::Pending => Poll::Pending,
           Poll::Ready(o) => Poll::Ready(o),
       }
   }
}

Safe on-the-fly pinning of Unpin objects

Finally, we have to fix poll_write, which calls poll on self.fut. For that, we need pinned access to the Future in self.fut. But in the new version of AsyncUploader, we have not declared fut as Pin. While we could have accomplished this by declaring fut as Option<Pin<UploadFuture>>, such a declaration is unnecessary. Instead, we can dynamically create a Pin<&mut UploadFuture> from self.fut when invoking poll, as demonstrated below:

// Impl-3.e: pinless.rs

impl AsyncWrite for AsyncUploader {
   fn poll_write(
       mut self: std::pin::Pin<&mut Self>,
       cx: &mut Context<'_>,
       buf: &[u8],
   ) -> Poll<Result<usize, Error>> {
       if self.fut.is_none() {
           let size = self.upload(buf);
           return Poll::Ready(Ok(size));
       }

       // We have to pin fut.
       let pinned = Pin::new(self.fut.as_mut().unwrap());
       match pinned.poll(cx) {
           Poll::Pending => {
               Poll::Pending
           }
           Poll::Ready(Ok(())) => {
               let size = self.upload(buf);
               Poll::Ready(Ok(size))
           }
           Poll::Ready(Err(err)) => {
               let _ = self.fut.take();
               Poll::Ready(Err(err))
           }
       }
   }
}

An attentive reader may wonder why we found it unnecessary to declare AsyncUploader::fut as Pin, whereas in all previous examples, the Future was pinned via BoxFuture. The reason is that UploadFuture is Unpin. So to satisfy the Pin’s contract, we don’t have to keep the pinning promise about UploadFuture uninterrupted for the entirety of its lifetime. However, an unnamed Future that we get from a call to an async function is !Unpin (not Unpin). For such a Future to work correctly, it must persist in the same location in memory throughout its entire lifetime. Using BoxFuture is the only way we can guarantee this safely.

Unsafe pinning and pin-projection

Currently, inside the body of upload_bytes, we still have an async block. Interestingly, we can remove it! The only source of async in upload_bytes is the call to tokio::sleep. The return type of tokio::sleep is a Future named Sleep. Instead of awaiting on it in an async-block, we can directly wrap it inside UpdateFuture:

// Impl-4.a: asyncless.rs

struct UploadFuture {
   content: String,
   fut: Sleep,
}

...

impl AsyncUploader {
   ...

   /// This is the slow upload function, which we want to use inside the
   /// AsyncWrite implementation.
   fn upload_bytes(&self, data: &[u8]) -> Result<UploadFuture, Error> {
       let content = String::from_utf8(data.to_vec())
           .map_err(|e| Error::new(ErrorKind::Other, format!("{e}")))?;
       Ok(UploadFuture {
           content,
           fut: tokio::time::sleep(tokio::time::Duration::from_secs(1)),
       })
   }
}

Note that I did not wrap Sleep in a BoxFuture. I could, but let’s see what happens if we deviate from that pattern.

A little bit of unsafe code

When implementing Future for UploadFuture we need pinned access to Sleep, because that is what we are going to call poll on. However, Sleep is !Unpin and we cannot use Pin::new to pin it. Instead, I am going to use the unsafe method Pin::map_unchecked_mut, to map self to a pinned mutable reference to its inner Sleep:

// Impl-4.b: asyncless.rs

impl Future for UploadFuture {
   type Output = ();
   fn poll(
       self: Pin<&mut Self>,
       cx: &mut Context<'_>,
   ) -> Poll<Self::Output> {
       // Make a copy of content before moving out of self.
       let content = self.content.clone();
       let pinned = unsafe { self.map_unchecked_mut(|s| &mut s.fut) };
       match pinned.poll(cx) {
           Poll::Pending => Poll::Pending,
           Poll::Ready(_) => {
               println!("Uploaded {}", content);
               Poll::Ready(())
           }
       }
   }
}

Owned !Unpin calls for pin_project

Since Sleep is !Unpin, after changing UploadFuture::fut from a BoxFuture to Sleep, UploadFuture too has become !Unpin. As a result, we can no longer simply pin it using Pin::new, as we did in poll_write in Impl-3. Here is the complete error message:

error[E0277]: `PhantomPinned` cannot be unpinned
   --> src/bin/asyncless.rs:82:31
    |
82  |         let pinned = Pin::new(self.fut.as_mut().unwrap());
    |                      -------- ^^^^^^^^^^^^^^^^^^^^^^^^^^ within `(PhantomData<&()>, PhantomPinned)`, the trait `Unpin` is not implemented for `PhantomPinned`
    |                      |
    |                      required by a bound introduced by this call
    |
    = note: consider using the `pin!` macro
            consider using `Box::pin` if you need to access the pinned value outside of the current scope
    = note: required because it appears within the type `(PhantomData<&()>, PhantomPinned)`
    = note: required for `Sleep` to implement `Unpin`
note: required because it appears within the type `UploadFuture`
   --> src/bin/asyncless.rs:39:8
    |
39  | struct UploadFuture {
    |        ^^^^^^^^^^^^
note: required by a bound in `Pin::<P>::new`
   --> /Users/razieh/.rustup/toolchains/nightly-x86_64-apple-darwin/lib/rustlib/src/rust/library/core/src/pin.rs:484:23
    |
484 | impl<P: Deref<Target: Unpin>> Pin<P> {
    |                       ^^^^^ required by this bound in `Pin::<P>::new`
...
503 |     pub const fn new(pointer: P) -> Pin<P> {
    |                  --- required by a bound in this associated function

To fix this, we need some more gymnastics similar to what we did above. This time, instead of writing unsafe code, I am going to use pin_project. The main benefit of using pin_project is that it saves us from writing unsafe code. The first step is to annotate AsyncUploader with the pin_project::pin_project attribute (note that, for simplicity, I’ve based this version on Impl-1, without the internal buffer in AsyncUploader):

// Impl-4.c: asyncless.rs

#[pin_project::pin_project]
struct AsyncUploader {
   #[pin]
   fut: Option<UploadFuture>,
}

The pin_project::pin_project attribute declares a new struct, which is very similar to AsyncUploader but instead of fut has a pinned field of type Pin<&'pin mut (Option<UploadFuture>)>. The declaration of this field is instructed by the #[pin] attribute that we added to fut. Moreover, pin_project::pin_project adds a project(Pin<&’a mut Self>) method to AsyncUploader, which returns an instance of the newly declared struct. Inside poll_write, we use self.project().fut to get access to a pinned version of UploadFuture.

// Impl-4.d: asyncless.rs

impl AsyncWrite for AsyncUploader {
   fn poll_write(
       self: std::pin::Pin<&mut Self>,
       cx: &mut Context<'_>,
       buf: &[u8],
   ) -> Poll<Result<usize, Error>> {
       // Method `project` consumes `self`. Before calling it on self,
       // let's create the future that we might need later.
       let fut = self.upload_bytes(buf);

       let mut this = self.project();

       if this.fut.is_none() {
           this.fut.set(fut.ok());
           // The future we created is not necessarily completed yet, but
           // we successfully used all the input bytes, and are now ready
           // to accept new bytes. The poll_write contract requires us to
           // return Ready.
           return Poll::Ready(Ok(buf.len()));
       }

       // We have an on-going future. We can only upload the given bytes
       // in buf, and create a new Future, if the previous future is
       // complete. Otherwise, we return Pending or Ready with an error.
       match this.fut.as_mut().as_pin_mut().unwrap().poll(cx) {
           Poll::Pending => Poll::Pending,
           Poll::Ready(()) => match fut {
               Ok(fut) => {
                   this.fut.set(Some(fut));
                   Poll::Ready(Ok(buf.len()))
               }
               Err(err) => Poll::Ready(Err(err)),
           },
       }
   }
}

A few things to note:

1. We could have used pin_project for UploadFuture (in Impl-4.a) in a similar manner to avoid the unsafe code in UploadFuture::poll.
2. When you have your fields pinned, as in Impl-1, Impl-2, and Impl-3, adding pin_project won’t offer much more.

As a result of Sleep and UploadFuture being !Unpin, now AsyncUploader is !Unpin too. This means that we cannot easily get mutable access to it. To account for this, I had to change the receiver of upload_bytes from &mut self to &self. If mutable access is really needed, a good solution is to refactor the code.

See the full implementation of this version on GitHub. There is also a more complex variant with a FileUploader, which is still simplified, but closer to what you may need to implement in a real-world project.

When to use pin_project?

In the four implementations we discussed in this article, we used different approaches in the design of AsyncUploader to pin the Future we needed to poll on. In each case, the choice was made based on the type of the Future wrapped in AsyncUploader. The following diagram summarizes these choices.

The focus is on designing the Outer data structure (e.g., AsyncUploader), considering that it will be used as a pinned value in its context (e.g., poll_write). The question is, what is needed to gain a pinned access to the inner, in that context, so that we can call poll on it.

If the inner is Unpin, we can simply access and pin it on-the-fly, on the call site, without having to write any unsafe code.

Otherwise, if Outer owns the inner, we can achieve our goal with some unsafe code (e.g., Impl-4). In this case, it is easy to reason about the soundness of the unsafe operation. Since outer is pinned, and since it owns inner, we can guarantee that inner does not change its location in memory during its lifetime. Because outer, being pinned and !Unpin, does not change its location in memory during its lifetime. In this case, to avoid writing unsafe code, we can use pin_project.

If Outer does not own the inner (e.g., Impl-1), then we cannot reason about the immovability of the inner so easily. Such situations often arise when working with Futures originated from invocations of async functions. Such a Future does not have a known concrete type, and thus neither a known size, at compile time. As a result, Outer cannot directly own the Future, but can only have a reference to it. To ensure that the referenced Future remains immovable and is not unsoundly mutated, we need to keep a pinned pointer to it for the entirety of its lifetime. Therefore, we need to use BoxFuture when declaring such an inner in the Outer. The pinning requirement cannot be achieved using pin_project in this case.

However, there is another scenario where it might be beneficial to annotate the outer with pin_project. Suppose that you have an Outer that will be used as pinned, but you don’t necessarily need pinned access to its fields. In these cases, especially if your struct is generic, it is advised to add pin_project to your struct to avoid potential unsound implementations of Unpin for your struct. This is because in addition to providing pin-projection functionalities, pin_project adds a conditional Unpin implementation to your struct. If I add impl Unpin for AsyncUploader {} in Impl-4 above, I get the following error:

error[E0119]: conflicting implementations of trait `Unpin` for type `AsyncUploader`
  --> src/bin/asyncless.rs:16:1
   |
16 | #[pin_project::pin_project]
   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^ conflicting implementation for `AsyncUploader`
...
22 | impl Unpin for AsyncUploader {}
   | ---------------------------- first implementation here
   |
   = note: this error originates in the derive macro `::pin_project::__private::__PinProjectInternalDerive` (in Nightly builds, run with -Z macro-backtrace for more info)

Interestingly, if you search for pin_project in GitHub, you can find many uses of pin_project without any fields being declared #[pin], or without a call to project. Some of these don’t have any generic types, and use BoxFuture (meaning that the Outer type is already unconditionally Unpin). I honestly don’t know what is the point of using pin_project in those cases. Some could be sloppy copy-pasting. But I hope some are intentional and have a good justification. Without any why-documentation it is hard to tell which one is the case!

Summary of Errors and Fixes!

Error: “future cannot be sent between threads safely”

• Likely root-cause: Future captures self, which is neither Send nor Sync.
• Fix: Refactor your struct so that the Future you have to poll on does not depend on self.
• Example: See Impl-1.g

Error: “lifetime may not live long enough”

• Likely root-cause: Future captures references, potentially causing a cyclic ordering among lifetimes!
• Fix: If possible, refactor your code so that the Future owns everything it needs, or declare a custom Future with well-thought generic lifetimes.
• Example: A hypothetical variant of Impl-3.a, with &[u8] as the type of data passed to upload_bytes.

Error: “T cannot be unpinned”

• Fix: Use BoxFuture or pin_project.
• Example: See Impl-1.g, and Impl-4.b

Error: “cannot borrow data in dereference of Pin<&mut T> as mutable”

• Likely root-cause: Pinned object is not Unpin.
• Fix: Refactor to avoid the need for a mutable receiver, or make your type Unpin.
• Example: A hypothetical version of Impl-4.a, with &mut self as the receiver.

There is more!

While crafting this article, I realized I needed a deeper understanding of the concept of pinning. This realization led me to an extensive exploration of the topic, and a write-up that I plan to publish as a follow-up of this post. Stay tuned for a deep dive into pinning!

Related Articles

1. The Async Rust book
2. Pin and Suffering: A very good step-by-step intro to futures and pinning
3. Comprehensive-Rust; the section and exercises on Async Rust

The best way to understand the behavior of any function is to plot it. But softmax is no ordinary function. It takes a vector and maps it to another vector. In more technical terms, softmax can be regarded as a multivariate and vector-valued function. There is no general method for visualizing such a function in 2-D, or even 3-D, spaces, as we typically do for functions over scalars. Yet, for a given input vector, we can plot the output of softmax over that input. There is a caveat though. If the input vector contains repeated elements, the resulting plot may not convey the complete picture. Additionally, when plotting softmax in this manner, we essentially disregard the order of the elements. Plotting the softmax of [1, 2] will be the same as that for [2, 1]. Nevertheless, this is the best way we have for visualizing softmax, and I will frequently use it in this post. The following diagram shows softmax applied to 3 integer vectors [0, …, 4], [0, …, 9], and [0, …, 19].

Softmax is one of the most important and popular activation functions used in machine learning. Activation functions serve a crucial role in artificial neural networks. Their primary purpose is to introduce non-linearities into the network, enabling it to learn complex patterns and relationships in the data. Intuitively, softmax maps a vector of real numbers to a probability distribution. That is why softmax is often used as the last layer activation function in many neural networks, including LLMs.

Suppose that you have the following vector:

This could be a datapoint you found somewhere, or the output from a dense layer in a neural network. When applying softmax to this vector, you get a vector of the same size, but each of the values in the new vector will be between 0 and 1, and their sum will be 1:

Apart from 0.79913, which corresponds to the largest value in the input vector, all other values are very small.

Mathematical formulation of softmax

To get deeper into understanding softmax, let’s take a look at its mathematical formulation. As we just saw, softmax transforms a vector of real numbers to another vector of the same size. The values of the output vector are all in the [0, 1] range, or more accurately the (0, 1] range, as we will see. Mathematically, we write:

where $S_k$ denotes the softmax function over the domain $\mathbb{R}^k$, $\mathbb{R}$ denotes the set of real numbers, and $k$ denotes the length of the input, and output vectors. For simplicity, in the following, we often drop the $k$ subscript from $S_k$ , and use $S$ to refer to the softmax function.

For an input vector $\mathbf{x}$, each element of $S(\mathbf{x})$ is calculated using the following formula:

In other words, we compute the $i$th element of the output vector by dividing the exponential of the $i$th element of the input (i.e., Euler’s number $e$ raised to the power of $x_i$) by the sum of exponentials of all elements in the input vector. As pointed out in the original softmax paper, softmax is a Normalized Exponential transformation.

From the formula, we can see that each element of $S(\mathbf{x})$ is in fact in the range (0, 1]. Additionally, we can see that the sum of all elements of $S(\mathbf{x})$ is 1. This is why the output of softmax is interpreted as a probability distribution. This behavior makes softmax suitable for multiclass classification problems, where each class has a probability associated with it.

A special case is when all $x_i$’s in the input vector $\mathbf{x}$ are the same. In this case, all $S(\mathbf{x})_i$’s will be the same as well, and since they add up to 1, they will all be equal to $\frac{1}{k}$, where $k$ is the length of vector $\mathbf{x}$. In this case, the output vector $S(\mathbf{x})$ represents a discrete uniform probability distribution.

Softmax magnifies differences

It is often said that softmax is sensitive to differences in input values, and magnifies those differences. In this section, I want to understand what that actually means.

If I apply softmax to $\mathbf{x} = [-10.0, -9.5, -9.0, …, 0, …, 9.5, 10.0]$, the resulting vector will contain very small values in the range [8.11e-10, 0.39]. What is magnifying about that?

Let’s see if visualization can help. For the vector $\mathbf{x}$ above, the following diagram shows $S(\mathbf{x})$ plotted over $\mathbf{x}$.

The graph resembles an exponential curve. We can express this curve using a function $f(.)$ such that $f(x_i) = S(\mathbf{x})_i$. Using the formula of softmax, we can see that $f(x_i) = \frac{e^{x_i}}{A}$, where $A$ is the sum of exponentials of elements of vector $\mathbf{x}$ above, and is equal to $55980.13$. Note that $f(.)$ is different from $S(.)$. The latter takes a vector as input, while the former takes a scalar as input and is a function that we can easily plot. The curve shown in the diagram above is an excerpt of the graph of $f(.)$ in the [-10, 10] range. The slope of this curve changes as we move along the x-axis. In fact, since $A$ is constant, it is easy to see that the derivative of $f(.)$ (which gives the slope of the curve at each point) is equal to $f(.)$ itself! We have $f’(.) =f(.)$. Note that we are not computing the derivative of softmax, just that of the curve shown in the diagram. At $x_i=-10$ the value and slope of the function are both $8.11e-10$ (almost zero), whereas at $x_i=10$ they are both $0.39$ (note that the x:y aspect ratio of the diagram is 1:10, so the curve looks a lot steeper than $0.39$)!

So even though values on the y-axis are all very small, they grow exponentially with changes in $x_i$. We can see this mathematically by looking at the ratios of elements of $S(\mathbf{x})$:

Each unit of change in the input vector, results in a change by a factor of $e$ in the output vector. This is what they refer to as the magnifying effect of softmax on differences in the input vector.

As mentioned above, softmax is sensitive to differences in the input vector. So if we increase or decrease all elements of the input vector a by a scalar, thereby maintaining the relative differences, the softmax of the resulting vector should be the same as the softmax of the original vector. In other words:

We can show that this is true by applying the formula of softmax, and using the fact that $e^{x_i+c}=e^{x_i} \times e^c$.

In practice, $S(\mathbf{x}+c) = S(\mathbf{x})$ plays an important role in stabilizing softmax implementations. Many softmax implementations internally compute $S(\mathbf{x}-max(\mathbf{x}))$ (see example), as opposed to $S(\mathbf{x})$, to avoid numerical overflow or underflow that may happen when computing exponentials, if $\mathbf{x}$ contains very large positive, or very small negative values.

Scaling the input vector

How about multiplying each element of the input vector by a constant $c$?

It is not so straightforward to specify $S(c\mathbf{x})$ in terms of $S(\mathbf{x})$, but we can compare the magnifying effect of softmax on $\mathbf{x}$ and $c\mathbf{x}$. Let $\mathbf{s} = S(\mathbf{x})$ and $\mathbf{w} = S(c\mathbf{x})$. In this case we have:

To express this in words, multiplying elements of the input vector by a constant results in an exponential growth in the magnifying effect of softmax on differences. Note that $c$ can be larger or smaller than 1. In the former case ($c>1$), smaller values in the input will be squashed more quickly. We see that in the following diagram, which shows $S(c\mathbf{x})$ plotted over $\mathbf{x}$, for $\mathbf{x}=[1, 2, 3, 4, 5, 6]$. As we go from left to right, we double the value of $c$ in each step, from 1 up to 8.

As we increase $c$, the tail (corresponding to smaller values in $\mathbf{x}$) is squashed and flattened more dramatically. Here, each unit of change in vector $\mathbf{x}$ results in a change by a factor of $e^c$ in the output vector.

The last diagram is particularly of interest. The output of softmax in this case is the vector $[1.6e-28, 1.4e-21, 1.3e-14, 1.1e-07, 9.99999887e-01]$, which is approximately equal to $[0, 0, 0, 0, 0, 1]$. In fact, if we increase $c$ even further, in the limit, $S(c.\mathbf{x})$ will be $[0, 0, 0, 0, 0, 1]$. This vector is the one-hot representation of $5$, the index of the largest element of $\mathbf{x}$ (and $c.\mathbf{x}$), or the argmax of $\mathbf{x}$ (and $c.\mathbf{x}$).

Mathematically speaking, argmax of a function $f(.)$ over a subset $Z$ of its domain is a member $z \in Z$ for which $f(z)$ has the maximum value of $f(.)$ over $Z$:

In machine learning literature and use cases, argmax is often applied to a vector (as opposed to a function), and returns the index of the largest element in the vector. As hinted above, this index can be represented as a one-hot vector.

Going back to our example above, we see that with a sufficiently large $c$ and a bit of approximation, softmax behaves like argmax with one-hot representation. In fact, a more accurate name for softmax that is used in some literature is softargmax.

In the diagram above, we could also go in the opposite direction and reduce $c$. By reducing $c$, or multiplying the input vector by a factor smaller than 1 ($c<1$), we can turn argmax into a top-k selector.

Does that ring a bell?

Suppose that you have a neural network, with softmax as its last layer activation function, and have trained the model on a classification task to predict the best matching class for any given input. Ideally, you have trained the model such that the output of the softmax activation function, when plotted over its input (vector $\mathbf{z}$ in the following diagram), resembles the rightmost plot in the diagram above (i.e., the plot with the lowest entropy).

At the inference time, you can simply use this for top-k sampling by multiplying the input to softmax (vector $\mathbf{z}$ above) by a parameter $c<1$, to make softmax behave more like the leftmost plot (with the highest entropy), in the previous diagram. In LLMs, a similar mechanism is used via the temperature parameter to control the diversity of the generated output.

Does softmax have a fixed point?

As mentioned earlier, softmax is a normalization function, and its output can be interpreted as a discrete probability distribution. Most normalization functions, such as z-score normalization and min-max normalization, when applied to data that is already normalized using the same method, do not produce new results. In other words, for these normalization functions, $N(N(\mathbf{x})) = N(\mathbf{x})$. Does the same property hold for softmax? In other words, if we apply softmax to a vector that already is a probability distribution will the result be the same as the input probability distribution? To answer this question, in this section we are going to compute the fixed point of softmax. In other words, we want to find vectors for which the following equation holds:

We know that $\mathbf{x}$ is a vector. For $S(\mathbf{x}) = \mathbf{x}$ to hold, all elements of $S(\mathbf{x})$ must be equal to their corresponding elements in $\mathbf{x}$. In other words, for each element $x_i$ of $\mathbf{x}$, the following equation must hold:

We can rewrite this as follows:

This equation is not very easy to solve. So let’s start with building some intuition with a simple example. For this, I decided to repeatedly apply softmax to its output, with [0, 1] as the starting input vector. In each iteration, the result of applying softmax is a 1-D vector of length 2, just like [0, 1] is. If we plot these vectors as points in the cartesian plane, we get the following diagram.

Each point in the diagram represents $S^n([0, 1])$, where $S^n(.)$ denotes the composition of softmax with itself $n$ times. We also define $S^0(\mathbf{x})$ to be $\mathbf{x}$.

As we can see in the diagram, repeatedly applying softmax, starting from $[0, 1]$ brings the output closer and closer to $[0.5, 0.5]$. In fact, $[0.5, 0.5]$ is the fixed point of softmax for 1-D arrays of length 2. There is nothing special about $[0, 1]$, and if we’d started from any other vector ($[0.5, 1]$ is a fun one), and repeatedly applied softmax, again we’d end up at $[0.5, 0.5]$.

Can we extend this result?

Indeed, we can.

We can prove this mathematically. Let $S_k(.)$ denote softmax on domain $\mathbb{R}^k$. We want to prove that:

Or

We can show this by plugging in the definition of softmax:

If you are not convinced by this proof, and want to see how we arrive at this result starting from $S(\mathbf{x})=\mathbf{x}$, there is a longer proof and discussion at the end of this section.

Now let’s take a closer look at what softmax does when it is repeatedly applied to its output. As briefly pointed above, each step of the transformation takes the result closer to a uniform distribution. However, softmax will maintain the relative ratios of the values in the output vector. We can see this in the following diagram, where I have applied softmax to vector $\mathbf{x}= [1/32, 1/16, 1/8, 1/4, 1/2, 1]$, up to three times (i.e., for $S(\mathbf{x})$, $S(S(\mathbf{x}))$, $S(S(S(\mathbf{x})))$).

Note that while the last three diagrams have the same shape, their ranges are different. From left to right, the ranges are $[0.116, 0.306]$, $[0.158, 0.191]$, and $[0.165, 0.171]$, respectively. Each step of the transformation takes the lower and upper bounds closer to $\frac{1}{6} = 0.1666$, and the standard deviation closer to zero. The standard deviations for $S(\mathbf{x})$, $S^2(\mathbf{x})$ and $S^3(\mathbf{x})$ are $0.0665$, $0.0116$, and $0.00195,$ respectively. If we zoom out and change the plots to have the same aspect ratio, we see that the curve approaches a straight horizontal line.

The choice of the initial input vector is not important, as we see in the following diagram, which shows the result of applying softmax to $\mathbf{x} = [1, 2, 4, 8, 16, 32]$, up to 3 times:

Again, the shapes have remained the same, but the ranges have shrunk.

So applying softmax on a probability distribution does not buy us anything! It only gets us closer to a uniform distribution, which means a smaller standard deviation for the output. Not that any of the neural networks I’ve seen uses softmax in this fashion, but in a deep neural network, a very small standard deviation often contributes to vanishing gradients, whereas larger standard deviations can result in more resistance against noise. Although seemingly not directly applicable anywhere in the design of a neural network, this experiment did help me better understand the behavior of softmax.

The impact of softmax on mean and standard deviation

Above, we briefly touched upon how starting from a vector, and repeatedly applying softmax on its output affects the standard deviation. Now let’s see if we can work out a formula for computing the standard deviation of $S(\mathbf{x})$.

When it comes to computing standard deviation, we always start from the mean. For any input vector $\mathbf{x}$ of length $k$, the mean of $S(\mathbf{x})$ is $\frac{1}{k}$. This is because, as we saw at the beginning, the sum of elements of $S(\mathbf{x})$ is 1.

With $\mu(S(\mathbf{x}))=\frac{1}{k}$, we can compute the standard deviation of $S(\mathbf{x})$, or rather its variance, as follows:

Expanding the quadratic term, and moving the terms around, we arrive at the following formula for the variance of $S(\mathbf{x})$:

Let $\mathbf{s}=S(\mathbf{x})$, and $\mathbf{r} = S(\mathbf{s}) = S(S(\mathbf{x}))$. Using the formula above, we have:

In the general case $\sigma^2(\mathbf{s}) \neq \sigma^2(\mathbf{r})$. This is in contrast to many normalization functions that, when applied to an already normalized vector, do not produce a new output, and thus, do not affect the variance and standard deviation of the data. For softmax, $\sigma(\mathbf{s})$ and $\sigma(\mathbf{r})$ will be equal, only when $\mathbf{s}$ is a uniform distribution.

To better understand the impact of softmax on the standard deviation, the following diagram shows how the standard deviation of $S^n(\mathbf{x})$ changes as $n$ grows from 1 to 16 (recall that $S^n(.)$ is the result of applying $S(.)$ to its output, $n$ times). The diagram shows changes in standard deviation for four starting vectors of length 3, as we increase $n$. We see that, as expected, in all cases, the standard deviation converges to 0.

We can repeat the experiment for $S(c.\mathbf{x})$. Reusing the four starting vectors as in the previous experiment, the following diagram shows how the standard deviation of $S(c.\mathbf{x})$ changes as we increase $c$ from 1 to 16.

For three out of the four vectors, the standard deviation converges to $0.47$. This is because, as we saw earlier, for large values of $c$, $S(c.\mathbf{x})$ becomes argmax, and its output approaches a one-hot vector. Standard deviation of a one-hot vector of length 3 (e.g., $[0, 0, 1]$) is indeed $0.47$. The blue line is the only exception to the convergence rule. This is because in this case, the elements of the input vector are all the same, the output of softmax is [$\frac{1}{3}$, $\frac{1}{3}$, $\frac{1}{3}$], and the scaling coefficient, $c$, has no impact on this output vector (i.e., does not bring it closer to a one-hot vector). Thus in this case, the standard deviation remains the same (i.e., zero) across different $c$’s.

We could do the same for $S(\mathbf{x}+c)$, but we know that $S(\mathbf{x}+c) = S(\mathbf{x})$. So the standard deviation would not change with changes in $c$.

Gradient and partial derivatives of softmax

In the realm of machine learning, the most important property of a function is its gradient. The gradient of a multivariate function that maps $\mathbb{R}^k$ to $\mathbb{R}$ is a vector composed of the partial derivatives of the function with respect to each of its input variables. Since softmax is a vector-valued function, we can regard each element of its output vector as a multivariate function, and form a gradient vector for it. By putting all these gradients together, we get a Jacobian matrix.

Here, we are going to stick with the simpler and more familiar concept of partial derivatives, which are the building blocks of both the gradient vector and the Jacobian matrix. More specifically, we are interested in computing:

for all $1 \leq i, j \leq k$, where $k$ is the length of the input vector $\mathbf{x}$.

We have to distinguish between two cases: (1) $i = j$, and (2) $i \neq j$. Let’s start with $i=j$:

We can simplify this to:

Since we have $0<S(\mathbf{x})_i \le 1$, the derivative of the $i$th element of $S(\mathbf{x})$ with respect to the $i$th element of $\mathbf{x}$ is always in the range [0, 1). It becomes zero only when $\mathbf{x}$ has length 1 (i.e., it is a scalar value). Now let’s compute the derivative when $i \neq j$. We have:

We can write this as:

As expected, this partial derivative is always negative. Putting these two cases together, if we increase one of the elements in the input vector, its corresponding element in the output will increase, while all other elements of the output will decrease.

Ok! This part was dense with formulas. Let’s look at a diagram now. It is not possible to plot these partial derivatives in the general case. So we will look at a simple example, for the sake of building intuition. Let’s consider the vector $\mathbf{x} = [0, 1, …, 31]$. We are going to take $x_0$ (I am using zero-based indexing here to be consistent with the python code that generated the diagram) and modify it, by increasing its value from -100 to 100, while keeping all other values in $\mathbf{x}$ fixed. In the following, the graphs on the top show how $S(\mathbf{x})_0$ and ${S(\mathbf{x})_3}_1$ change as we increase $S(\mathbf{x})_0$. I have chosen ${S(\mathbf{x})_3}_1$ because it gives the best visual appearance.

The lower graphs show how the derivatives of $S(\mathbf{x})_0$ and ${S(\mathbf{x})_3}_1$ change as we increase $x_0$.

We can take this one step further, and compute the second partial derivatives. For $i=j$, we have:

This formula yields a positive value, if $S(\mathbf{x})_i$ is less than $0.5$, and becomes negative only when $S(\mathbf{x})_i$ is larger than $0.5$. We can also see this in the diagram. In this example, $S(\mathbf{x})_0$ is larger than $0.5$ for $x_0>31$. At this point the derivative of $S(\mathbf{x})_0$ starts to fall, as we see in the lower diagram. Note that as $S(\mathbf{x})_0$ grows larger than $0.5$, ${S(\mathbf{x})_3}_1$ falls below $0.5$. Intuitively, this is because a probability distribution (e.g., vector $S(\mathbf{x})$) cannot have more than one value larger than $0.5$.

The second partial derivative of softmax in the $i \neq j$ case is given by the following formula:

According to this formula, the second derivative of ${S(\mathbf{x})_3}_1$ with respect to $x_0$ is positive only when $S(\mathbf{x})_0>1/2$. As we saw above, this first happens after $x_0=31$ (i.e., at $x_0=32$), and that is where ${dS(\mathbf{x})_3}_1/dx_0$ starts to grow, as we see in the lower diagram. Note that despite this, the value of ${dS(\mathbf{x})_3}_1/dx_0$ remains negative, and ${S(\mathbf{x})_3}_1$ continues to decline after this point, as shown by the red line in the upper diagram.

The “S”-shaped curve in the upper diagram might have reminded you of sigmoid. Here are the formulas for sigmoid ($\sigma(x)$), and its first and second derivatives:

Despite the seeming resemblance between sigmoid and softmax, they are very different functions. Sigmoid operates on scalars, and maps them to scalars. Even though in the machine learning world the input to sigmoid is often a vector, sigmoid is applied to each element of the vector independently. Changes to one element of the input vector has no impact on the outputs of other elements. This is in contrast to softmax where, as we see in the red curves in the diagram above, changes to one element of the vector affects the output for all other elements.

We could continue studying other properties of softmax, for instance its inverse (does it even have an inverse?). But I am going to stop this part here!

Softmax in higher dimensions

So far, we have been only considering 1-dimensional arrays. But in general, we could apply softmax to any n-dimensional (also written as n-D) array. In fact implementations of softmax in popular python packages take an n-D array as input. In tensorflow terminology, an n-D array is called a rank-n tensor.

The simplest n-D array is a 0-D array, or a scalar. Applying softmax to any scalar gives 1. If we plot $S(\mathbf{x})$ for $\mathbf{x} \in \mathbb{R}$, we’ll get the horizontal line $y=1$. Note that in practice, a 0-D array behaves the same as a 1-D array of length 1.

Now let’s move on to the 2-D space, and apply softmax to points on the unit circle! Because, why not?!

If we have $k$ points, then the input vector will be a 2-D array of the shape $(2, k)$ or $(k, 2)$, depending on how you want to arrange the points. For this example, I use $(2, k)$, to keep the code for generating the points and plotting them simpler. The domain of softmax in this case is $\mathbb{R}^{2 \times k}$. The following snippet generates such an input vector:

x = np.arange(-1, 1.01, 0.02)                # 101 points in [-1, 1].
y = np.sqrt(np.absolute(1 - np.power(x, 2))) # The upper semicircle.

x = np.concatenate((x, np.flip(x)), axis=None)
y = np.concatenate((y, np.flip(-y)), axis=None)
v = np.array([x, y])

As shown in the code, we have: $y = \pm \sqrt{1- x^2}$. The vector we want to pass into softmax is $\mathbf{v} = [\mathbf{x}, \mathbf{y}]$. In the code snippet above, vector $\mathbf{v}$ has shape $(2, 202)$.

With such an input vector, we have a few options when applying softmax. We can either treat the entire content of the array as a single vector, and apply softmax on the flattened 1-D vector, or apply softmax on each dimension independently. In scipy.special, which is the package I am using, this is done using the axis argument to the softmax function. Note that regardless of our choice of axis, softmax from scipy.special returns an array with the same shape as the input.

Let’s start with applying softmax to the entire vector:

from scipy.special import softmax

s = softmax(v, axis=None)
print(s.sum()) # Prints 1.0
print(s.shape) # Prints (2, 202)

In this code snippet, we use $\mathbf{s}$ to denote the result of applying softmax to $\mathbf{v}$. As expected, after the transformation, the sum of all the elements in $\mathbf{s}$ is 1. To see the effect of applying softmax on $\mathbf{v}$, the following diagram shows two curves. The one on the left plots $\mathbf{y}$ (i.e., $\mathbf{v}[1]$) over $\mathbf{x}$ (i.e., $\mathbf{v}[0]$). As expected, the result is a circle. On the right, we do the same for elements of $\mathbf{s}$. We plot $\mathbf{s}[1]$ on the y-axis, over $\mathbf{s}[0]$ on the x-axis.

After the transformation we have an asymmetrical and skewed egg-shaped curve. In our original input, the points on the upper circle have positive values on the y-axis. These are sines of angles between 0 and $\pi$. After the transformation, these values have become smaller, but they are proportionately larger than the y-values of the lower half of the transformed shape. The points on the lower semicircle, on the left, have negative values on the y-axis. After the transformation, these values have become small positive values, all in the range $[0.0007, 0.00195]$. The x-axis also shrank from [-1, 1] to $[0.0007, 0.0053]$.

The following table reports the min, max, mean, and standard deviations of $\mathbf{x}$, $\mathbf{y}$, $\mathbf{v}$ and their softmax transformations, $\mathbf{s}$, $\mathbf{s}[0]$, $\mathbf{s}[1]$.

Vector $\mathbf{v}$ has $404$ elements. So the mean of its softmax transformation, $\mathbf{s}$ , is $\frac{1}{404} (=0.00247)$. But $\mathbf{s}[0]$ and $\mathbf{s}[1]$ have different means, which should not be surprising, because the input vectors $\mathbf{x}$ and $\mathbf{y}$ don’t have the exact same values, as we can see from the difference in their standard deviations. Note that the ranges for $\mathbf{s}$, $\mathbf{s}[0]$ and $\mathbf{s}[1]$ would not be so small if we had fewer points in $\mathbf{v}$. But then the diagrams would not be as smooth.

From the table, we see that $\mathbf{s}[0]$ and $\mathbf{s}[1]$ have the same min and max, and these are in fact softmax transformations of -1 and 1 for our input vector $\mathbf{v}$. In other words:

Visually, the highest and lowest points of the original circle have become the highest and lowest points in the new shape, respectively. Likewise, the leftmost and rightmost points of the original circle have become the leftmost and rightmost points in the new shape.

However, their alignment with the center has changed. In the circle, the x-coordinate of the highest and lowest points is the same as the x-coordinate of the center, which is at $(0, 0)=(\mu(\mathbf{x}), \mu(\mathbf{y}))$. In the transformed shape, the x-coordinate of the highest and lowest points is $0.00195$ (under the transformation $0$ is mapped to $0.00195$), which is different from the x-coordinate of the new center at $(0.00230, 0.00265)=(\mu(\mathbf{s}[0]), \mu(\mathbf{s}[1]))$. We see the same misalignment for the y-coordinate of the leftmost and rightmost points in the new shape.

In addition to this misalignment with respect to the center, the points on the curve no longer have the same distance to the center. The following diagram shows the distance of each point on the transformed shape to its center (indicated by the red dot in the diagram above), starting from the leftmost point (corresponding to $(-1, 0)$ in the circle) and visiting all points on the curve in a clockwise traversal.

We see the first peak at the point corresponding to $(0, 1)$ in the original circle. The second peak corresponds to $(1, 0)$. These two points, although being reflections of each other with respect to the $x=y$ line, have different distances from the center. This is because the x- and y-coordinates of the center in the transformed shape are different.

We can talk more about this curve, and the curious reader is invited to analyze it further, but for now, let’s move on to applying softmax along each row of our 2-D vector $\mathbf{v}$, independently.

This is effectively the same as applying softmax to $\mathbf{x}$ and $\mathbf{y}$ vectors separately. As shown in the following snippet, we do this by setting axis=1.

s = softmax(v, axis=1)
print(s.sum()) # Prints 1.9999999999999998
print(s.shape) # Prints (2, 202)

The sum of all elements of $\mathbf{s}$ in this case is $2.0$, instead of $1.0$ that we saw in the previous case. When we plot $\mathbf{s}[1]$ over $\mathbf{s}[0]$, we see a similar shape as before, but the ranges of the transformed values are different. On the x-axis we have $[0.00154, 0.0114]$, and on the y-axis we have $[0.00135, 0.0099]$.

As before, let’s take a quick look at min, max, mean, and standard deviation after the transformation:

Comparing this table with the previous one we see that the min and max values are roughly twice that of the previous case. Also, the means of all vectors in this case are equal to $\frac{1}{202}$. Can you explain why?

The last option we have for applying softmax to our vector $\mathbf{v}$, is to apply it on each point, independently.

s = softmax(v, axis=0)
print(s.sum()) # Prints 202.0
print(s.shape) # Prints (2, 202)

The sum of all elements of $\mathbf{s}$ in this case is $202.0$ (we have $202$ points in vector $\mathbf{v}$). Since the output of softmax is always a valid probability distribution, for each index $i \in [0, 201]$, we have:

In other words, coordinates of the points on the transformed shape can be specified as:

The softmax transformation has turned our circle into a line!

I hope you find these experiments and discussions as much fun and useful as I did. I am still trying to get my head around the egg-shaped curves. Do these shapes have a name, or any interesting properties? Let me know if you have any ideas.

Here is the jupyter notebook I used for the experiments and generating the diagrams.

Last updated: 2024-01-08

Lately, I’ve been thinking a lot about the qualities of great leaders. I’ve read a lot about the topic; and almost everything I read seems to be written from the perspective of somebody in a leadership position, assuming authority, and carrying a huge ego! In this post, I want to offer a different perspective, reflecting on my own experience reporting to supervisors, managers and tech leads that I’ve worked with in the past 15 years.

I have identified two important traits that one must have to be a great leader, when working with highly intelligent, self-motivated, and autonomy-craving people. In my own work, I use my experience and these two traits as a guidance when I lead people and projects. Of course, depending on the situation, the type of tasks, and the people you work with, other qualities might become more important.

Ability to recover from failures

The first distinguishing characteristic of a great leader is her confidence in her own ability to recover from failures, and to handle unanticipated situations. This kind of confidence comes with years of experience working on various projects (not necessarily leading them), facing challenges, and correcting mistakes.

Such a leader knows that small mistakes or misses of the team will not easily damage her reputation or threaten her success. She also knows that making mistakes is the only way to grow. As a result, she can give her reports the space to be themselves, make their mistakes, find their own ways, and grow! She knows if individuals in the team grow, the team grows, and their success is her success.

Of course, this only works if you have self-motivated and capable people in your team. A great leader is good at discovering talent, skills, and unique qualities in others, assigning them the right tasks, and giving them credit for their good work. A great leader with these qualities can focus on her main job, which is to empower her team and help them grow and bring their best selves to work.

An experienced leader capable of recovering from failures has things under control without having to micromanage. This gives her time to develop herself as a leader. Her involvement, therefore, is often in the form of reducing chaos, reducing workloads, and building a safe space for everyone to cooperate.

I have been lucky enough to work with one or two leaders like this during my career.

Ability to lead more technically adept people

A great leader is comfortable with and capable of leading people with expertise and skills beyond her own! She can happily and respectfully delegate responsibilities to those more technically capable than herself; and is nevertheless capable of integrating the outcomes of their efforts into a larger cohesive whole. This is the only way that a team can achieve anything beyond the capabilities of a single person.

This quality is crucial for effective leadership, but is difficult to acquire and develop. On one hand, a robust personality is needed to successfully navigate interactions with individuals across hierarchical levels, varying technical expertise, and diverse backgrounds, all the while demonstrating compassion, patience, and humility. On the other hand, an effective leader must have a depth of technical knowledge, breadth of experience, and the capacity to judiciously respond to an array of potentially unforeseen scenarios. If you have done only one thing in your career, it is unlikely that you’ll be able to make the call when you face a new situation.

Many people in leadership and management roles struggle with leading more experienced or more technically adept individuals. They often perceive leadership as delegating tasks one could manage independently but lacks the time for! While there is merit in delegating such tasks, its effectiveness depends on considering the interests and career aspirations of the team members receiving these assignments. Failing to do so can result in a demoralized team. To have a productive and successful team, you have to go beyond delegating the tasks you can do yourself. You have to allow and encourage your team to do things you cannot do, or cannot even imagine! You have to give them the space to be creative in their own ways.

If as a leader, you ever find yourself in a situation where you think you can do something better, or know exactly how something should be done, then do it yourself; nicely and quietly, without even thinking about delegating it or discussing it with anyone! Unless you are delegating the work to a team member who is passionate about doing that particular piece of work; perhaps because of a personal interest, a learning opportunity, or career development opportunity.

In my career, I’ve seen inexperienced managers, whose notion of leadership was merely a mechanical task delegation. At times, I’ve felt trapped, my skills wasted, and my effort ineffective. I’ve felt that the knowledge and capabilities I could bring to the team were capped by those of my manager.

If this is your leadership style, take some time and rethink your approach because you really need to change it. If you make your knowledge and capabilities a cap for the knowledge and capabilities that your team brings to work (perhaps because you want to have full control over what they do), then the collective achievements of your team will be confined to what you alone could accomplish – perhaps a bit faster, if you are lucky. If, instead, you are open to letting your team bring all their talent, ideas and capabilities to the project, and can support them in contributing in ways you could not imagine, then your team will accomplish much more than you could ever imagine, and those achievements will be yours too.

One of the managers that I loved working with, and still consider the best manager I ever known, always said that she was not a technical person, and left the technical work to me. Indeed she was not deep into technical details as much as I was, but she was very bright, had a breadth of technical understanding, and was always on top of everything, with full control and clarity over the overall direction and progress of the project. She was particularly great at giving me exactly the kind of help I needed. When I had a new idea, and needed to know who could help, my manager put me in touch with the right people and teams. We worked in harmony, and were successful because we brought complementary skills and knowledge to the team. I had the technical knowledge, and she had the organizational knowledge and connections that I lacked. She considered her main job to serve her team, and it is that mindset that made her so great and successful.

What can you do?

I hope this post helped you see leadership in a new light. Whether you are in a leadership role or not, if you are not happy with your situation or team dynamics, the best chance you have at changing your circumstances is to talk to your team or manager, in a candid manner. Regardless of your roles, it is important to see each other as equals, who happen to have taken different responsibilities, based on their talents, skills, experiences, and hopefully interests. Keep in mind that others, whether in a leadership role or not, like yourself, are allowed to have flaws, space to grow, and things to learn.

Update 26 Nov 2023: I finally got the time to read through How to Lead a Team (Google SWE book). If you want to learn more about how to effectively lead a team, and positive patterns and antipatterns of leadership, this is a good read. If you are in a leadership position, this is a must read.

Last updated: 2023-11-27

A lot has been said about Generative AI since that interview and its impact on society, the future of work, and in particular the future of creative work. Many, including Sam Altman, have pointed out that these tools are here to assist the human professionals, rather than replacing them. They can amplify you, and unblock your creativity, and can have a transformative impact on the creative industry, and other jobs.

Obviously, this is not the first time that AI has achieved superhuman results. Earlier we saw AlphaGo mastering the game of Go, AlphaFold solving the protein folding problem, and simpler AI-powered systems before them assisting professionals in medicine, face recognition, and other domains. What is new is that transformer-based general-purpose generative models, which are closer to general AI, have now reached unexpected success in creative tasks.

This unforeseen success raises some questions. The one that I have been thinking about is the question of complexity. Are creative tasks more complex (e.g., writing) than routine mechanical tasks (e.g., truck driving, or working in assembly line)? Has AI now achieved a point where it can perform fundamentally more complex tasks?

Many mechanical tasks can be very complex. In particular, when precision and accuracy is required to correctly complete each step of the work. Consider driving for instance. There are many situational constraints and changing conditions that should be taken into account when driving a car.

When it comes to creative work, it is difficult to talk about its complexity, as creativity spans a wide range of tasks and activities. Nevertheless, creativity feels complex because we know very little about how it works. One aspect of creativity is novelty: to piece together new things, and in new ways. However, novelty is considered creative only if it results in something useful. The authors in 1 state that: “We distinguish creative ideas from bizarre ideas, which are also novel, because creative ideas take into account the parameters of a situation, the constraints. Novel productions that are in some way “useful” or “worthwhile” in a given context are, thus, creative.”

When it comes to creative content generated by AI, people often mean textual or visual content generated by the AI system in response to a prompt provided by the user. Luckily, for most of this content, there is no clear line between right or wrong, and valid or invalid. For textual content for instance, the main constraints that apply, and can rule out the usefulness of the generated content, are grammatical rules and a degree of relevance to the prompt. Consider instead generating a novel, perhaps more succinct, solution to a math problem. This one is a lot more complex, as satisfying the constraints of the problem is more difficult. You know exactly what is right and what is wrong. The space of possible solutions is much slimmer.

This lack of a clear line between right and wrong, is, in my opinion, the first factor that contributes to the success of AI in creative tasks. It significantly reduces the complexity of the task.

The second factor is that the usefulness and worth of creative and artistic content is a matter of taste and subjective opinion. Creative contents are greatly open to interpretation and their aesthetic appeal changes from one person to another. The human brain, as the consumer of the creative content, fills in the gaps and participates in assigning meaning to the creative content and making it worthwhile. This, which I call the smart recipient effect, is the second contributor to the early success of AI in creative tasks.

In the following, I’ll elaborate on these two factors! But the short answer to my two questions above is that (1) the creative tasks that we are talking about in the context of generative AI are not that complex, and (2) it does not seem like these large models can solve dramatically more complex problems for us.

No right or wrong

When it comes to artistic content, there is often no right or wrong. You don’t care about completeness, accuracy, or precision. The artist can violate certain constraints and rules, and still produce a worthwhile piece of art, maybe even introduce a new form of art!

The value of a piece of art is in the idea behind it and the feeling it provokes in you. In this case, AI is a great assistant. It has not replaced human producers, but is a great aid to amplify their abilities. It is a very good tool for expanding your imagination, and presenting you with new ideas.

Consider a paragraph generated by ChatGPT. What makes it worthwhile or useful? Certainly relevance to the prompt, good grammar, a coherent flow of thoughts, and factfullness. Apart from good grammar, the other criteria are not very well defined or easy to measure. For factfulness in particular, unless you are an expert on the topic; you cannot say how factual the generated text is. Often if the text is well-organized and is relevant to the topic at a high level, you are happy with it.

How about images generated by AI? What makes them worthwhile? You may say relevance to the prompt, good ratios and perspectives, lighting effects, clear shape boundaries, etc. In my experience, as soon as you have an unusual prompt, the generated image fails in all these aspects. Nevertheless, for all intents and purposes, a creative content producer might find the AI-generated image exactly what he or she has been looking for.

Here is a little experiment, from 22 April 2023, where I asked ChatGPT to generate a prompt for DALL·E 2:

Me: DALLE is an AI system that generates images from textual prompts. Produce a prompt that can be used as input to DALLE.

ChatGPT response:

Here's a prompt that can be used as input to DALLE:
"Imagine a peaceful countryside landscape with rolling hills, tall trees, and a
river flowing gently through the valley. The sun is setting in the distance,
casting a warm orange glow across the scene. A small cabin sits on the edge of
the river, smoke gently rising from the chimney. In front of the cabin, a wooden
dock extends out into the water, and a rowboat bobs gently in the current. A
family of ducks swims lazily by, enjoying the tranquility of the evening."

I had no idea you could provide such long inputs to DALL·E 2! Honestly, I could never come up with such a beautiful description of a scene. But there is nothing particularly novel about the scene described in the prompt. Anyway, I gave it to DALL·E 2, and here are the four images it produced:

The images look a lot better than any images I ever got from DALL·E 2 with my prompts! However, there are a lot of details in the prompt that are not present in the generated images. I don’t see smoke from a chimney, nor a bobbing boat, nor a family of ducks in any of the images. It is difficult to say which one of these photos ChatGPT had in mind, but I think you can agree with me that all these images somewhat fit the description.

I thought I’d ask ChatGPT for another prompt:

Me: That was very artistic. Do you have another one?

ChatGPT’s response:

Sure, here's another prompt for DALLE:
"Picture a futuristic cityscape with towering skyscrapers reaching towards the
sky, interlaced with networks of sleek flying vehicles weaving in between the
buildings. A bustling street below is filled with people of all ages and
backgrounds, moving quickly along the sidewalks. Neon lights and holographic
advertisements fill the air with color and light, and in the distance, a
massive artificial intelligence network can be seen, humming quietly in the
background. Amidst all this, a lone figure stands at the edge of a rooftop,
looking out over the vast expanse of the city and contemplating its infinite
possibilities."

While the previous prompt described a rather customary scene, this one has got more novelty in it. There are few things that are particularly interesting about this description. The humming massive artificial intelligence in the background, and the lone figure standing at the edge of a rooftop. To my disappointment, I don’t see any of them in the images generated by DALL·E 2. Nor are there any flying vehicles. I suppose this is because of the limited training samples. There has never been that much variety or volume for such futuristic cityscape images. Parts of the generated images are also blurred, which again I think is because of the smaller sample space.

The most interesting part however is that ChatGPT provides details of sound (the quietly humming AI) in its description. I don’t see why ChatGPT thinks this is relevant, but here is what ChatGPT has to say about that:

Me: How do you expect DALLE to draw the artificial intelligence network and its humming sound?

ChatGPT’s response:

As an AI language model, I cannot predict with certainty how DALLE will generate
images from textual prompts, as the specific details and style of the image will
depend on the model's training data and internal weights.

However, it's possible that DALLE could use its knowledge of common visual
representations for technology and sound to generate an image of an artificial
intelligence network, such as rows of computer servers, glowing data cables, or
a digital interface displaying lines of code.

As for the sound of the network humming, DALLE might generate an abstract
representation of sound, such as waves or lines emanating from the network, or
a visual metaphor such as the depiction of the network as a living organism that
emits a constant hum. Ultimately, the output of DALLE will be influenced by the 
way it has been trained on vast amounts of data and the creative interpretation
of the language prompt.

Regardless of its relevance and usefulness, the idea of including a description of a background noise in a prompt for DALL·E 2 is new to me, but it does not necessarily mean that it is an original idea that ChatGPT came up with. In fact, “The extent to which the work produced is novel can vary from being original only for the person who completed the work (this is the notion of reinventing ideas known already in the larger social context) to being original for a limited social group to being original for all of humanity” [1].

Finally, we have to try Bard too. This is what Bard said in response to the first prompt I gave to ChatGPT:

Bard’s response:

Sure, here are some examples of prompts that can be used as input to DALL-E:

* A photorealistic painting of a cat sitting on a couch in a living room with a
  coffee mug in its hand.
* A detailed infographic of the human respiratory system.
* A vintage, grainy, sepia, photograph of an adorned Native American man,
  Hassleblad, f4, 150mm, natural light.
* An iPhone made of terracotta from Harappa, Indus Valley Civilization, Pakistan;
  2600-1900 BC, studio light.
* A colorful surreal seamless pattern of a cat and a dog playing together.
* A film still from the movie "The Matrix" of Neo dodging bullets in slow motion.
* A detailed map of the city of New York with all the major landmarks labeled.
* A realistic 3D rendering of the planet Earth from space.
* A beautiful abstract painting in the style of Jackson Pollock.
* A funny cartoon of a dog chasing its tail.

These are just a few examples, and there are many other possibilities. The possibilities are endless!

These look more like what I’d expect people to provide as their prompts to DALL·E 2. I suppose this is because Bard has access to online content. There are some technical details in some of the prompts that I know nothing about, and cannot verify if DALL·E 2 could get that right! But I fed the second prompt to DALL·E 2, and the results are hilarious. It’s been known that DALL·E 2 cannot do text!

The smart recipient effect

“A stupid man’s report of what a clever man says can never be accurate, because he unconsciously translates what he hears into something he can understand.” – Bertrand Russell

This is true in the other direction too! In fact, similar to the quote above, we can say: A human brain’s report of what ChatGPT says can never be accurate, because it unconsciously fills in the gaps and assigns meaning to the words as it sees fit.

I believe, the ultimate measure of understanding something, is the ability to do something with it. In a conversation, or exchange of content, the smarter the recipient, the more likely for it to grasp the meaning in the content, and the easier for the producer to produce content that the recipient can act on.

A smart recipient fills in the gaps and automatically fixes inaccuracies based on its knowledge and model of the world. It actively participates in assigning more meaning to the content at hand, and asks for more details when needed. I have demonstrated this in an earlier post on playing with Lego with ChatGPT. I had to actively interpret ChatGPT’s instructions, making tweaks and adjustments along the way using my knowledge and observations.

If we were to ask ChatGPT to produce some code, for instance in Rust, while the resulting code would seem correct to the human eye at the first glance, and in simple cases the user would be able to run through the instructions and calculate the output of the code, it is very likely that the Rust compiler would not be able to do much with the code (unless it is a very small snippet without many dependencies). The Rust compiler is not as smart as the human user, and in particular cannot fill in the gaps.

For ChatGPT to produce input for the Rust compiler, it has to be familiar with all the rules of the language – the syntax and the semantics – or be able to interact¹ with the compiler the same way the Rust programmers do. Most Rust programmers do not master all details of the language, but can interact with the compiler, and learn from it, by adjusting and improving their mental model of the language. This need for precision and details, can be generalized to interaction with any other compiler, or tool that humans have built.

And that I think is the biggest difference between the creative industry and other industries: people working in the creative industry produce content for human interpretation, while people in other professions produce content that must obey strict rules, or be consumed by dumb tools.

One last thing to note is that, in a conversation, ChatGPT, like us, plays both the recipient and the producer roles. It is the recipient, when processing our requests, and has proven to be a very smart one. It has succeeded in understanding very complex prompts, and correctly following the instructions given to it. When responding, it plays the producer’s role, and has proven to produce content that the human brain can interpret and enjoy.

What does this mean for our creativity?

ChatGPT is now being branded as an assistant tool that can speed up or improve the quality of your work.That is true and great. Recently I asked ChatGPT to rewrite some of my every-day writings, and its rewrites were really better than my own. ChatGPT is certainly better at using phrases and expressions than I am. For my more creative writings, however, I have not been happy with ChatGPT’s rewrites. In those cases it felt like ChatGPT was trying to get the novelty out of the text and make it closer to what is deemed normal. This is the main danger I see in the use of these Generative AIs. There does not seem to be a great operational definition of creativity, let alone a measure of it. But I believe like all other natural phenomena there is a bell curve of creativity somewhere. Too much use of generative AIs will push the standard deviation of that curve to zero!

Last updated: 2023-05-04

tl;dr: GPT-3.5 and GPT-4 models can generate Lego instructions that appear plausible. However, they may be incomplete or include occasional incorrect or imprecise steps. Therefore, it’s unlikely that you could build a model using these instructions alone. Nevertheless, the fact that these models can generate textual Lego instructions at all is impressive, as such instructions are notoriously difficult to find.

A few weeks ago, I noticed my Lego sets, and felt a strong urge to play with them! Since our friendly AIs now have a reputation of being good at creative tasks, I thought I’d invite them to the game!

DALL.E on the Task

First, I asked DALL.E to generate instructions for making a Lego car. Here is the prompt I gave to it, and the result that I got. Not at all what I was hoping for.

Me: step-by-step Lego instructions for building a small car from Lego bricks

I could not think of how to improve my prompt, so instead I thought I’d ask it to generate IKEA-style instructions for building a table with three legs. This is what I got:

Again, not what I was looking for! Normally, I’d tweak my prompt to see if I could get any better results, but this time I wanted to get to my Lego playing, and was not in the mood to argue with DALL.E. So I moved on to ChatGPT!

ChatGPT Knows the Game

This is the free version of ChatGPT, with the GPT-3.5 model.

Me: Please generate descriptive Lego instructions for building a small car from Lego bricks.

The response from ChatGPT was very detailed and long:

This looks really good. I like how ChatGPT starts by listing the ingredients, and has numbered the steps. Let’s see if I can follow these instructions to build a Lego car!

Step 1: Gathering Pieces

Step 1: Gather Your Pieces
…
4 wheels
2 2x4 bricks
2 2x2 bricks
1 2x3 brick
1 2x1 brick
1 2x2 sloped brick
1 2x1 sloped brick
1 1x2 plate
1 1x1 plate
1 1x2 tile
1 1x1 tile
1 steering wheel

Figure 1-a below shows the pieces I gathered.

1. In all Lego sets I’ve seen, the tires and the bolts are separate pieces and you have to put them together to get your wheels. ChatGPT did not point that out, which is different from normal Lego instructions. It also did not specify the wheel size.
2. I could not find a 1x1 plate in any of my Lego sets, so I am instead using a 1x1 plate with a horizontal clip!

Step 2: The Base of the Car

Step 2: Build the Base of the Car

To start building the base of the car, take your 2 2x4 bricks and attach them
to each other using a 2x2 brick. This will create a flat base for your car.

1-a	1-b

Figure 1-b shows the result of this step. I guess you could do it differently, but I am fine with this one.

Step 3: Add the Wheels

Step 3: Add the Wheels

Next, attach your 4 wheels to the bottom of the base. Place 2 wheels on each
side, towards the front and back of the car.

This is where things get complicated! With the pieces from step 1, and the instructions in this step there is no way to directly attach the wheels to the base of the car. Or if there is, my imagination cannot go that far! I need an additional piece to attach the wheels to. I asked ChatGPT for help.

Me: How do I attach the wheels to the base of the car?

There was no mention of an axle in the first step. That is a new piece we need. Ignoring that part, I think ChatGPT is telling me to build what I’ve shown in Figure 2-a! The additional advice at the end is marvelous. I am not sure what kind of car, axle, wheel, or angle it is talking about now.

2-a	2-b	2-c

Anyway, to make progress, I am using a wheel holder piece to attach the wheels. This is an additional piece that was not mentioned in the first step. I need two of them (Figure 2-b). And now I can attach the wheels, as shown in Figure 2-c. It is good to note that the title of this step says, “Add the Wheels” as opposed to “Attach the Wheels”!!

Step 4: The Front of the Car

Step 4: Build the Front of the Car

Take your 2x3 brick and attach it to the front of the car, just behind the
front wheels. This will create the front of the car.

This step was easy. See Figure 3-a.

Step 5: The Windshield

Step 5: Build the Windshield

To create the windshield, attach your 2x2 sloped brick to the top of the 2x3
brick. Then, attach your 2x1 sloped brick to the top of the 2x2 sloped brick.
This will create a slanted windshield for your car.

My slanted windshield is shown in Figure 3-b. I am not sure what the purpose of the 2x1 sloped brick is. The car is starting to look like a rooster to me!

Step 6: Add the Steering Wheel

Step 6: Add the Steering Wheel

Take your steering wheel and attach it to the center of the 2x2 sloped brick,
just below the windshield.

Another tricky step, with the keyword Add. There are a few things that are odd with this step. First of all, what is the center of a 2x2 sloped brick? Secondly, the 2x2 sloped brick is already attached to other pieces. There is no space for attaching anything below it, or the windshield. Finally, there is not enough space to fit the steering wheel anywhere in that area.

I know this is a lot to expect a language model to have an accurate spatial imagination of a construction like this, but it is interesting to see where it fails.

To make progress, I am just “putting” the steering wheel close to the windshield, instead of attaching it (Figure 3-c).

3-a	3-b	3-c

Step 7: Back of the Car

Step 7: Build the Back of the Car

To create the back of the car, take your 2x2 brick and attach it to the back of
the base. Then, attach your 2x1 brick to the back of the 2x2 brick.

This part went well. Now the car, in Figure 4-a, resembles a rooster even more!

Step 8: Add the License Plate

Step 8: Add the License Plate

Take your 1x2 plate and attach it to the back of the car, just below the 2x1
brick. Then, attach your 1x1 tile to the center of the plate. This will create
a license plate for your car.

I cannot figure out how to attach the license plate. Let’s for now, very suspiciously, put the tile just below the center of the 1x2 plate.

The Finished Car

Step 9: Finish the Car

Finally, take your 1x1 plate and attach it to the top of the 2x1 brick at the
back of the car. This will create a small spoiler. Your Lego car is now
complete!

Really?! Anyway, Figure 4-c shows my completed Lego car, following these instructions!

4-a	4-b	4-c

ChatGPT never used the 1x2 tile that was in the list of required pieces. I did not follow up on that with it.

I prefer my car to be one single piece, so I made some tweaks, but my tweaks did not fix the rooster look!

While this is impressive, and ChatGPT got many steps right, it failed at more complicated steps. In particular, where some spatial reasoning was required. This experiment is a good demonstration of ChatGPT’s performance on tasks were reasoning (in this case spatial reasoning) and precision are required.

In one of my earliest experiments with ChatGPT, I asked it to draw the truth table of a simple logical formula. While it drew very nice tables, with all the rows, to my disappointment, it failed to correctly calculate the truth values for all the rows. So I tried to walk it through the process. At some point during the debugging, ChatGPT produced the following explanation for its calculation:

In other words, the logical or operation is true if at least one of the
operands is true, and it is false only if both operands are false.

In the case of (a2 v a4), both operands are false. However, since the logical
or operation is true if at least one of the operands is true, the overall
expression (a2 v a4) is true even though both operands are false.

I want you to focus on the last sentence: “However, since the logical or operation is true if at least one of the operands is true, the overall expression (a2 v a4) is true even though both operands are false.”

I don’t know how to argue with this! If a person said it to me, I’d leave their presence and never get anywhere near them! To be honest though, I think, the average person does struggle with simple logic. So perhaps it should not be surprising that ChatGPT produces text like this, as it is largely trained on text produced by ordinary people. In my opinion (and I am not an expert), fixing the reasoning problem requires serious engineering work on the design and building of these generative models.

GPT 4 on the Task

About a month later, I repeated the experiment with the GPT-4 model. To my disappointment, GPT-4 was not any better.

It still struggles with the attachment of the pieces. Consider the following instruction:

2. Start with the base: Attach the 4 wheel axles to the underside of the 4x12
base plate, placing two axles at each end. Then, attach the wheels to the axles.

You cannot just attach an axle to a plate! You need some connectors. So ChatGPT (still the GPT-4 model) corrected that step:

Bard Does not Feel Like It!

Before finishing this post, let’s see how Bard performs on the task.

Me: Please generate descriptive Lego instructions for building a small car from Lego bricks.

Bard, as usual, gives a very high level answer. I like how it advises me on using my imagination. I was hoping to get a taste of its own imagination. So I asked it to be a bit more specific, and give me step-by-step instructions.

Me: Can you be a bit more specific? I was looking for step-by-step instructions.

Not as detailed and structured as ChatGPT’s response, but at least now it is talking about specific Lego bricks and how to attach them.

I particularly like the instruction for building the wheels:

To build the wheels, you will need 4x4 bricks and 2x2 bricks. Start by
building a rectangle with 4x4 bricks. Then, add 2x2 bricks to the top and
bottom of the rectangle to create a wheel.

This is a nice approximation of a wheel, if all you have are bricks. And after this, Bard is back to its style of giving a high level description and no specifics.

Final Remarks

It should not be surprising that these language models don’t have a good grasp of the interface between the Lego pieces. Lego instructions are visual, with hardly any text on them. The only virtual Lego design tool that I know of is the Bricklink Studio, which stores Lego designs as .io files that do not have a textual format and can only be parsed by the Bricklink Studio software. So there is essentially no reliable text on Lego designs to train these language models on.

The good news for Lego designers, and IKEA designers too, is that their jobs are safe!

Last updated: 2023-05-03