fnands

Quick change detection with geospatial embeddings

Ferdinand Schenck — Sun, 23 Nov 2025 00:00:00 GMT

If you work in the Geospatial Machine Learning (or AI¹) domain, you’ll know that embeddings are all the rage at the moment. I used the meme below in an internal presentation a few months ago, but I think it is relevant here:

Meme from Zoolander with Mugatu saying ‘So Hot Right Now’

But what exactly are embeddings?

What are embeddings?

There is a little bit of heterogeneity in embeddings, but effectively they are compressed representations of data that have been “embedded” into some feature space, usually by a learned model. This is not unique to geospatial data, and can be done for text, images, videos, graphs, or whatever data structure you can train a model for. You might have heard of RAG (retrieval augmented generation) in the language domain, where documents are embedded by a language model, and the embeddings are stored in some way (often a vector database) from which they can be retrieved via a semantic search. A search engines image search functionality often works in a similar fashion, but using image models.

For the geospatial case, embeddings are created by training a foundation model on large amounts of (usually) satellite data, and then embeddings are created for a large corpus of satellite data. Some examples are the embeddings created using the Clay Foundation model, the EarthIndex embeddings created by EarthGenome, the Tessera Embeddings, the Major TOM embeddings by ESA, and the ones I used today, the AlphaEarth Foundations dataset by Google (Deepmind).

There are some differences between all these embeddings: the sources of data are different (most use at least Sentinel-2, but often combine it with other data), and some do inference on image patches (e.g. 224x224 pixel patches of th earth), while some do it per pixel, i.e. release an embedding for every 10x10 meters on earth. The AlphaEarth dataset does the latter.

From here on, when I talk about embeddings I will generally mean “embeddings created using a geospatial foundation model on global satellite data”.

Why embeddings?

Once you have embeddings, there is a actually a lot of fun stuff you can do with them. This tutorial on the Google Earth Engine page goes through a couple examples, ranging from similarity search, to supervised and unsupervised classification.

Essentially, a lot of work goes into training these large foundation models, and doing inference with these models is not cheap either. See this Github issue I stumbled upon where the Clay team discuss what it would cost to do an inference run for the entire world:

chip size (x10m/px)	cost unit (km^2/$/h/worker)	cost to run the world
50x50 px	1000	$510K
100x100 px	4000	$127K

But once you have created the embeddings dataset, you have amortized the cost, and everything you do downstream can now benefit from the work that has been done in creating these embeddings. This allows you to cheaply do large scale calculations for a fraction of the cost of if you had to run the encoder every time.

30DayMapChallenge: Day 23 - Process

Being a geospatial company, we at LiveEO are doing the 30DayMapChallenge. I’ve wanted an excuse to play around with embeddings a little bit more, so I thought I’d pick day 23. The theme of today is to document the process of making the map, hence this blog post.

What I wanted to do is to see how I could use embeddings to extract some information that can be mapped. What I’m doing today is creating a change map of Berlin for the earliest and latest dates available in the AlphaEarth embeddings dataset, i.e. 2017 and 2024.

I did this by calculating (one minus) the cosine similarity between all pixels that overlap with Berlin for 2017 and 2024, and using that too see what has changed.

Conceptually, embeddings that are similar will have a high cosine similarity, and embeddings that are very different will have a low cosine similarity, so by doing this we can find the places where the embeddings between 2017 and 2024 have changed the most. As the AlphaEarth embeddings have been unit normed, this means that the formula above actually simplifies to just taking the dot product, which is pretty cheap to calculate.

Plotted, this looks like:

Change Map of Berlin between 2017 and 2024 calculated using the AlphaEarth embeddings

As you can see, the areas with high change stand out as bright spots, while the areas with low change stay dark.

A few highlights

Adlershof

I studied in Adlershof (it hosts the science campus of HU-Berlin), and during the time I was there a lot of construction happened. Looking at the map, you can clearly see a cluster of large changes in that area:

Adlershof Changes

Sentinel-2 Before

Sentinel-2 After

According to the Economic Affairs, Energy and Public Enterprises for Berlin, over half a billion Euros have been invested by the Berlin government to improve Adlershof in the last few years.

Tegel

Another point that stood out was the Tegel Airport, which was shut down in 2021:

Tegel Changes

Sentinel-2 Before

Sentinel-2 After

The temporary buildings around the terminal are part of a refugee camp set up in 2022 to house people mostly fleeing the war in Ukraine.

The area of the former airport is being turned into a business hub as well as a living area with around 5000 new homes.

I was wondering what was going on on the runway, and a quick search brought up a thread on the r/berlin subreddit asking the same question, which had this comment:

Comment
byu/Dragonogard549 from discussion
inberlin

So pretty much the norm for doing construction in the heavily bombed Berlin: you gotta remove the bombs before you can build new things.

Grunewald

Another one that stood out to me in was a blob in the Grunewald forest in the western side of Berlin:

Grunewald Changes

Sentinel-2 Before

Sentinel-2 After

It’s from a large fire that occurred in 2022 when a bomb storage site caught fire and sparked a fire in the forest. This one impressed me as it is not that apparent on the imagery, at least not in the RGB bands.

Method

If you want to make a change map like this, it’s rather easy if you have access to Google Earth Engine.

I did this in the GEE code editor, but you can also call the API.

The first step is to define the AlphaEarth Embeddings as your dataset:

// Load collection.
var dataset = ee.ImageCollection('GOOGLE/SATELLITE_EMBEDDING/V1/ANNUAL');

Next, define your AOI. You can get more complicated, by pulling in a shapefile or some other geometry, but I just took a bounding box of Berlin:

// Area of interest - Berlin, Germany
var aoi = ee.Geometry.Polygon([
  [
    [13.08, 52.33],  // Southwest corner
    [13.76, 52.33],  // Southeast corner
    [13.76, 52.68],  // Northeast corner
    [13.08, 52.68],  // Northwest corner
    [13.08, 52.33]   // Close the polygon
  ]
]);

Next, we get the embeddings for the years we are interested in. The AlphaEarth embeddings are annual, i.e. one per year, and span the range of 2017 to 2024 so far, and I hope they will continue adding these every year:

// Get embeddings for 2017
var startDate2017 = ee.Date.fromYMD(2017, 1, 1);
var endDate2017 = startDate2017.advance(1, 'year');
var embeddings2017 = dataset
    .filter(ee.Filter.date(startDate2017, endDate2017))
    .filter(ee.Filter.bounds(aoi))
    .mosaic();

// Get embeddings for 2024
var startDate2024 = ee.Date.fromYMD(2024, 1, 1);
var endDate2024 = startDate2024.advance(1, 'year');
var embeddings2024 = dataset
    .filter(ee.Filter.date(startDate2024, endDate2024))
    .filter(ee.Filter.bounds(aoi))
    .mosaic();

Next we will calculate the cosine similarity between our dates for every pixel. We define the change as :

// Calculate cosine similarity per pixel across all 64 bands
// Note: The embeddings have unit length by design, so cosine similarity 
// simplifies to just the dot product

// Dot product: sum of element-wise multiplication across all 64 bands
var cosineSimilarity = embeddings2017.multiply(embeddings2024)
    .reduce(ee.Reducer.sum());

// Change = 1 - cosine similarity (higher values = more change)
var change = ee.Image(1).subtract(cosineSimilarity);

// Clip to AOI
var changeClipped = change.clip(aoi);

// Debug: check statistics
print('Change statistics:', changeClipped.reduceRegion({
  reducer: ee.Reducer.minMax(),
  geometry: aoi,
  scale: 1000,
  maxPixels: 1e9
}));

Next we want to visualize this. This is optional, as you can also at this point just export the raster, but it’s helpful for testing purposes:

// Visualize change map
Map.addLayer(
  changeClipped,
  {min: 0, max: 0.5, palette: ['black', 'blue', 'cyan', 'yellow', 'red']},
  'Change 2017-2024 (brighter = more change)'
);

// Optional: visualize embeddings for context
var visParams = {min: -0.3, max: 0.3, bands: ['A01', 'A16', 'A09']};
Map.addLayer(embeddings2017.clip(aoi), visParams, '2017 embeddings', false);
Map.addLayer(embeddings2024.clip(aoi), visParams, '2024 embeddings', false);

Map.centerObject(aoi, 8);
Map.setOptions('SATELLITE');

And voilà, you have a change map! The code above runs in seconds!

You can export this to Google drive with:

// Export full change map
Export.image.toDrive({
  image: changeClipped,
  description: 'change_map_2017_2024',
  folder: 'GEE_Berlin_Analysis',
  region: aoi,
  scale: 10,
  maxPixels: 1e9,
  fileFormat: 'GeoTIFF'
});

And now you can visualize these changes in your favorite GIS software (I used QGIS).

If you want the full code I used to also export the matching Sentinel-2 images, expand the code below:

Show code

Translating Cython to Mojo, a first attempt

Ferdinand Schenck — Mon, 28 Jul 2025 00:00:00 GMT

Ever since I heard about Mojo I (and presumably most other people) thought it would be a good language to speed up functions to be called from Python. Everyone knows that vanilla Python can be slow, but one of the reasons that Python programs can be reasonably fast in practice is because Python often leans on libraries written in more performant languages, predominantly C/C++, but increasingly also Rust.

Until recently, there has been no real way to call Mojo code from Python, but about a month ago (in Max release 25.4) the ability to call Mojo from Python was added as a beta feature. It’s not fully cooked yet, and it will likely still change a lot, but I wanted to give it a look just to get an idea of where things are heading.

One specific idea that I had when I heard about Mojo was that Mojo might be a good replacement for Cython and apparently I was not the only one to have had this thought:

Convergent thinking on HN

The comments are from the HackerNews discussion on Vincent Warmerdam’s blog post titled “Python can run Mojo now” which made it to the front page of HN a while ago.

So where can I find a lot of Cython code?

Scikit-learn

Scikit-learn implements a bunch of machine learning algorithms and related utilities, and makes heavy use of Cython. How hard would it be to translate some of the Cython code in scikit-learn to Mojo?

I wanted a piece of code that was relatively simple, both just as I didn’t want to jump into the deep end, but also because there are some restrictions on Mojo functions being called from Python, namely (from the known limitations section of the Mojo/Python interop):

Functions taking more than 3 arguments. Currently PyTypeBuilder.add_function() and related function bindings only support Mojo functions that take up to 3 PythonObject arguments: fn(PythonObject, PythonObject, PythonObject).

A simple case: dbscan_inner

An example I found that satisfies this criteria is the inner loop of DBSCAN that assigns points to clusters. It’s relatively short and takes exactly three arguments.

This is a classic case of a place where you would usually want to call speed up a tight inner loop in Python, in this case written in Cython:

# Fast inner loop for DBSCAN.

# Authors: The scikit-learn developers
# SPDX-License-Identifier: BSD-3-Clause

from libcpp.vector cimport vector

from ..utils._typedefs cimport uint8_t, intp_t


def dbscan_inner(const uint8_t[::1] is_core,
                 object[:] neighborhoods,
                 intp_t[::1] labels):
    cdef intp_t i, label_num = 0, v
    cdef intp_t[:] neighb
    cdef vector[intp_t] stack

    for i in range(labels.shape[0]):
        if labels[i] != -1 or not is_core[i]:
            continue

        # Depth-first search starting from i, ending at the non-core points.
        # This is very similar to the classic algorithm for computing connected
        # components, the difference being that we label non-core points as
        # part of a cluster (component), but don't expand their neighborhoods.
        while True:
            if labels[i] == -1:
                labels[i] = label_num
                if is_core[i]:
                    neighb = neighborhoods[i]
                    for i in range(neighb.shape[0]):
                        v = neighb[i]
                        if labels[v] == -1:
                            stack.push_back(v)

            if stack.size() == 0:
                break
            i = stack.back()
            stack.pop_back()

        label_num += 1

It’s not a complicated algorithm, and it labels core points and propagates that label to the neighbors of the core points.

Translating to Mojo

For the most part I just copied over the Cython code verbatim. There is a bit of boilerplate we need to add to the .mojo file to make the function callable:

from python import PythonObject
from python.bindings import PythonModuleBuilder

from os import abort

@export
fn PyInit__dbscan_inner_mojo() -> PythonObject:
    try:
        var m = PythonModuleBuilder("dbscan_inner_mojo")
        m.def_function[dbscan_inner]("dbscan_inner", docstring="Fast inner loop for DBSCAN.")
        return m.finalize()
    except e:
        return abort[PythonObject](String("error creating Python Mojo module:", e))

but other than that, the translation was actually surprisingly straightforward, see if you can spot the differences in the Mojo and Cython versions:

fn dbscan_inner(is_core: PythonObject,
                 neighborhoods: PythonObject,
                 labels: PythonObject) raises:
    var i: Int = 0
    var label_num: Int= 0
    var v: Int = 0

    var stack: List[Int] = []

    for i in range(labels.shape[0]):
        if labels[i] != -1 or not is_core[i]:
            continue

        # Depth-first search starting from i, ending at the non-core points.
        # This is very similar to the classic algorithm for computing connected
        # components, the difference being that we label non-core points as
        # part of a cluster (component), but don't expand their neighborhoods.
        while True:
            if labels[i] == -1:
                labels[i] = label_num
                if is_core[i]:
                    neighb = neighborhoods[i]
                    for i in range(neighb.shape[0]):
                        v = Int(neighb[i])
                        if labels[v] == -1:
                            stack.append(v)

            if len(stack) == 0:
                break
            i = stack.pop()

        label_num += 1

I defined stack as a Mojo List[Int] to replace the C++ vector[intp_t] implementation in Cython. Other than the changes related to stack, the only other changes were the initializations of the variables, and casting the entries in neighbors to integers.

It was honestly quite a bit simpler than I thought it would be, and the fact that both Cython and Mojo’s syntax is based on Python means a lot of the code “just works”.

As part of this experiment, my goal was to change the Python code as little as possible, and all I needed to do in _dbscan.py was add:

import max.mojo.importer
import sys
sys.path.insert(0, "")

from _dbscan_inner_mojo import dbscan_inner

The sys.path.insert(0, "") is a bit clunky, but the Mojo devs have said this is a temporary workaround.

I then ran pytest and all all the dbscan tests passed:

============================================================== test session starts ==============================================================
platform linux -- Python 3.12.9, pytest-8.4.1, pluggy-1.6.0
rootdir: /fast/Workspace/scikit-learn
configfile: pyproject.toml
plugins: anyio-4.9.0
collected 30 items                                                                                                                              

tests/test_dbscan.py ..............................                                                                                       [100%]

======================================================== 30 passed, 10 warnings in 0.54s ========================================================

The performance however is a bit lacking, presumably because Mojo is iterating over PythonObjects for which it can’t properly optimize:

Cython average time: 2.78e-05 seconds
Mojo average time: 0.0227 seconds

That’s around 800 times slower than Cython. We can however, make some minor tweaks to improve this.

Improving performance

Let’s look at what is being passed to dbscan_inner:

core_samples is a numpy array of bytes (np.uint8) signifying whether on not a sample is considered a core sample.
neighborhoods is a list of numpy arrays of integers that specify which points neighbor each point. Effectively the edges of a graph.
labels is a numpy array of integers, initialized to -1, signifying that the points are currently unlabeled.

We can transform labels and is_core into Mojo Spans (thanks to Owen Hilyard on the Modular Discord for the hints):

var labels_ptr = labels_py.ctypes.data.unsafe_get_as_pointer[DType.index]()
var labels = Span(labels_ptr, Int(labels_py.shape[0]))

var is_core_ptr = is_core_py.ctypes.data.unsafe_get_as_pointer[DType.bool]()
var is_core = Span(is_core_ptr, Int(is_core_py.shape[0]))

Not the prettiest, but this creates the Spans without copying over the data.

The final code looks like:

fn dbscan_inner(is_core_py: PythonObject,
                 neighborhoods_py: PythonObject,
                 labels_py: PythonObject) raises:
    var label_num: Int= 0
    var v: Int = 0

    var labels_ptr = labels_py.ctypes.data.unsafe_get_as_pointer[DType.index]()
    var labels = Span(labels_ptr, Int(labels_py.shape[0]))

    var is_core_ptr = is_core_py.ctypes.data.unsafe_get_as_pointer[DType.bool]()
    var is_core = Span[mut=False](is_core_ptr, Int(is_core_py.shape[0]))


    var stack: List[Int] = []

    for i in range(len(labels)):
        if labels[i] != -1 or not is_core[i]:
            continue

        # Depth-first search starting from i, ending at the non-core points.
        # This is very similar to the classic algorithm for computing connected
        # components, the difference being that we label non-core points as
        # part of a cluster (component), but don't expand their neighborhoods.
        while True:
            if labels[i] == -1:
                labels[i] = label_num
                if is_core[i]:
                    var neighb = neighborhoods_py[i]

                    for j in range(len(neighb)):
                        v = Int(neighb[j])
                        if labels[v] == -1:
                            stack.append(v)

            if len(stack) == 0:
                break
            i = stack.pop()

        label_num += 1

Testing the performance now, we get:

Cython average time: 2.9e-05 seconds
Mojo average time: 8.59e-05 seconds

So around 3x slower than Cython, but a lot faster than before.

Ideally we would also translate neighborhoods into a Mojo type, but it gets a bit tricky here as neighborhoods is a list of numpy arrays, which can all have different sizes, so simply assigning them to a type is a bit hard. There might be some solution out there, although likely changing the input to dbscan_inner to something that can more easily be mapped to Mojo is likely the most sensible answer, but that’s beyond the scope of this little test.

Even so, the overall performance of DBSCAN as a whole is unchanged, as this inner function isn’t really the slow part of the algorithm (benchmarking code adapted from HDBSCAN):

dbscan Cython vs Mojo

The performance is identical (lines overlap almost exactly), and it’s the other parts of DBSCAN, like the neighborhood calculation, that take up the majority of the time:

=== Component Breakdown ===
Data validation: 0.0003s (0.0%)
NearestNeighbors fit: 0.0304s (2.1%)
Radius neighbors computation: 1.3573s (95.1%)
Core sample identification: 0.0019s (0.1%)
Cluster assignment: 0.0002s (0.0%)
Total: 1.4278s

In the future, I’d like to look into translating the slower parts of DBSCAN into Mojo, as the neighborhood radius calculation that takes up the most time can probably be parallelized, maybe even on the GPU.

Conclusions

I chose this example not because it makes a lot of sense to translate it to Mojo, but just because it was easy to do so in a short amount of time. Right now, the Python interop is still a little too bleeding edge to do anything serious with, but at the pace that the language is evolving I doubt that this will be true for long.

It was however promising just how simple it was in this case, and most of the effort was in translating the PythonObjects into appropriate Mojo types to allow the compiler to reason about them. If I could request something from the Modular team it would be a “cheat-sheet” for best practices for translating common Python/numpy types into Mojo.

A more wholistic approach would be to also reconsider what is being passed to Mojo to make your life a bit easier when it comes to doing these translations.

Future plans

Once the Python interop stabilizes a little I want to see if I can rewrite some more substantial part of scikit-learn in Mojo, and preferably some algorithm that’s amenable to vectorization, possibly even on the GPU, so that I can really play into the strengths of Mojo. If you have any suggestions for an algorithm that is in need of some speeding up, let me know.

I think moving a lot of scikit-learn’s more computationally intensive code to Mojo could be an interesting project. There is a project called Mojmelo which is effectively the Mojo ecosystem’s answer to scikit-learn, however, almost no-one uses Mojo just yet.

On the other hand, scikit-learn was downloaded 100 Million times last month, so if you can speed up some of scikit-learn’s algorithms you can have a positive impact for a lot of users.

Citation

BibTeX citation:

@online{schenck2025,
  author = {Schenck, Ferdinand},
  title = {Translating {Cython} to {Mojo,} a First Attempt},
  date = {2025-07-28},
  url = {https://fnands.com/blog/2025/sklearn-mojo-dbscan-inner/},
  langid = {en}
}

For attribution, please cite this work as:

Schenck, Ferdinand. 2025. “Translating Cython to Mojo, a First Attempt.” July 28. https://fnands.com/blog/2025/sklearn-mojo-dbscan-inner/.

Can multi-sensor foundation models be more than the sum of their parts?

Ferdinand Schenck — Sun, 22 Jun 2025 00:00:00 GMT

Geospatial foundation models (GFMs) have been on my mind recently, partially because I attended the ESA-NASA International Workshop on AI Foundation Model for EO and partially because I’ve been working on fine tuning some GFMs for downstream use at work for the last while. This post was in part prompted by two recent LinkedIn posts, one by Christopher Ren and the other by Madeline Lisaius, both of which express some amount of skepticism about the way in which current GFMs are trained, although from somewhat different angles. Christopher Ren also wrote an expanded blog post on the subject, which takes aim mostly at IBMs new TerraMind GFM, but it is worth reading the responses from one of the TerraMind authors at the bottom of the post as that adds some nuance to the arguments.

It’s somewhat clear that GFMs are a hot topic in the Earth Observation (EO) space at the moment, and it is fair to question whether the hype is warranted. At the ESA-NASA workshop one of the points made was that there seems to be much more activity in the creation of GFMs than actual downstream use of them so far, and there were some interesting discussions as to why this might be the case at the moment.

A recent post from Bruno Sanchez-Andrade Nuño (director of the Clay project) also made me think that there is a rough bifurcation in the GFM space appearing: one branch goes deep and the other goes wide. I think it is best if we understand which branch a model fits into and not judge one by the standards of the others.

I’m not going to directly respond to the other conversations going on: I’m just adding my two cents to the mix, and I want to be clear that the work I am doing definitely falls into the “go deep” branch, and my opinions are very much coloured by that fact.

Should GFMs go more general or specialize?

On the surface this might seem like a slightly odd question seeing as one of the principle reasons people are interested in GFMs (and FMs in general) is better generalization: EO is after all often a global endeavour, and it is desirable to have a foundation model that will help your downstream tasks generalize across geographies, illumination conditions, imaging angles etc. But there are many aspects to generalization, some of which don’t apply to all sensors. An example is the time of day an image was taken at. This can strongly affect what your image looks like, as shadows and illumination levels can vary greatly by time of day. This however does not really affect missions like Sentinel-2, where the orbit has been selected in such a way that the mean local solar time when the image is taken is always approximately 10:30 am, leading (by design) to very consistent illumination levels. Similar arguments go for viewing angles.

One of the ways that people have been trying to get more general is to train foundation models on multiple sources. Examples of this are the Clay foundation model (at least V1) which was trained on a wide range of sensors from MODIS, with a 500 m GSD, to aerial imagery of New Zealand at under 50 cm GSD:

Platform	Spatial Coverage	Spectral bands	GSD (meters)	Samples
Landsat 8 and 9	Global	6 optical bands	30	11,617,984
Sentinel 2 L2A	Global	10 optical bands	10	18,683,945
Sentinel 1 RTC	Global	2 radar bands	10	16,133,394
NAIP	USA	4 optical bands	< 1	20,984,171
LINZ	New Zealand	3 optical bands	< 0.5	3,299,006
MODIS	Global	7 bands	500	1,350,864

Another example of this is DOFA, which takes a similar approach to variety in input sensors, this time including hyperspectral data at 224 spectral bands:

Platform	Spatial Coverage	Spectral bands	GSD (meters)	Samples
Sentinel-1	Global	2 SAR bands (VV, VH)	5 x 20	4,642,353
Sentinel-2	Global	9 optical bands (0.49-2.15µm)	10	977,774
Gaofen	China (cities)	4 bands (RGB + NIR)	4	117,450
NAIP	USA	3 bands (RGB)	1	2,332,351
EnMAP	Various	224 hyperspectral bands	30	11,483

The DOFA paper is worth a read, and putting on my scientist hat: this is really interesting work, and it’s really interesting to see the different solutions that these authors have come up with to make a single model deal with such varied inputs. But putting back on my engineer hat I have to ask: what do we gain from this?

The engineering perspective

One of the points made in the DOFA paper:

The increasing number of specialized foundation models makes it difficult to select the most appropriate one for a specific downstream task.

On the surface this sounds fair, but is it really that hard to go to PapersWithCode, find the most similar dataset to your downstream task and select a model based on that? I can’t really think of a scenario where you would not just spend a day or two searching through the literature for the most fitting model for your particular use case. The one case I can maybe think this might be the case is if you are a geospatial person with no ML skills and have a model that was set up for you as a black box behind some interface and squeezing every last bit of performance out is not critical to you.

When implementing a model for a specific product need, one often focuses on the specific sensor, or at least a set of very similar sensors, e.g. sub-meter GSD sensors with at least the four usual spectral bands. When building a product that will utilize exclusively Sentinel-1 data, does the model really gain anything from being trained on Sentinel-2 and aerial imagery as well?

With all that being said, if you do have multiple sensors available at inference time (e.g. Sentinel-2 and Sentinel-1 data), it does seem to make sense to train/infer on multiple modalities at once. See e.g. table 2 in the TerraMind paper.

A concrete example

A while ago we were testing out a few foundation models as backbones for a product we are developing, which boils down to bi-temporal change detection using Planet’s SkySat constellation. We chose the main backbone we are using based on benchmarks, but I did have the nagging question of how much do we really gain from this, and if other backbones might offer better performance? This was basically the theme of my talk at the aforementioned ESA-NASA workshop.

I ran a few test using a variety of FM backbones, some trained on remote sensing data, and some just on natural images, just to see how much the pre-training dataset really matters. To make the comparison fair, the backbones used all had around 100 M parameters, but I did throw in a smaller baseline (ChangeFormer), as well as a 300 M version of the best performing network just to see if size matters (spoiler: it does).

Quick FM Benchmark

One of the most interesting comparisons here is DINOv2: I used two variations, one using the original weights trained on natural images from Meta, and another with weights from Keumgang Cha, which were trained on the MillionAID and SkyScript datasets. MillionAID is exclusively aerial imagery, while SkyScript contains mostly aerial imagery, plus some SkySat, Sentinel-2 and Landsat images. It’s abundantly clear that the same architecture trained on remote sensing images greatly improve downstream performance compared to a the variant that was trained on natural images. This is expected, but it’s impressive to see how large this gap is.

The best model we tested was trained mostly on aerial imagery, showing the domain gap isn’t so much whether or not your sensor is in space or on a plane, but has more to do with similar resolutions.

The models were all trained for the same number of epochs, on the same relatively small dataset (around 1600 patches of 512 x 512 pixels) with the same optimizer etc. The encoders were not frozen, but trained with a lower learning rate than the decoders, as is common practice in most transfer learning scenarios. I will caveat this all with saying that I didn’t do a massive amount of hyperparameter tuning for this test, but I think the differences are significant enough that it probably wouldn’t make too much of a difference.

How can we test this?

What I would need to see to be convinced is that when training a foundation model on multiple sensors that it would perform better on downstream tasks on each of the sensors than it would be if it was trained exclusively on the specific sensor to be used. I.e. one would need to show that the model would be more than the sum of it’s parts.

The question is pretty much, given the same architecture, compute budget and dataset size, can a model learn something from one sensor that improves its performance on another? Or could it be that we need to throw everything into a big bucket and burn a lot of compute in the fashion of the current big LLMs that are so popular right now in order to really see generalization?

I’m definitely not ruling out the possibility that there might be some case (e.g. the sensor you are targeting doesn’t have a lot of data available), but I have the feeling that the further away in GSD and spectral characteristics you go the less helpful pre-training becomes.

How specialized can we go?

It’s fairly obvious that the best GFM you can choose will likely be the one trained on the exact sensor you are targeting for your downstream task. This is fairly easy for sensors like the Sentinel missions or the Landsat missions, where anyone with a bit of compute to burn can easily download tons of data from those missions and train a model. Even for aerial imagery there is a lot of open data available, with the caveat that the data is not as global, and aerial sensors do have some sensor to sensor variability.

Where this gets tricky is in the commercial domain, where data isn’t freely available and providers put strict licenses on their data¹. To train a foundation model on commercial data requires you to dump somewhere between hundreds of thousands up to millions of Euros on data alone, which is infeasible for most researchers, and a significant investment for most companies. The only case that I know of so far of someone creating a sensor specific foundation model is a Pleiades Neo foundation model created by Disaitek, which was made possible by being granted access to Pleiades Neo imagery through a “Call for Innovation” from Aribus and CNES. Disaitek of course does not make this model public, as this presumably gives them a bit of an edge over their competitors, and as the model was trained on data covering France only, it is questionable of how much use it would be in other parts of the world.

So what can be done in the commercial space? Most companies don’t have access to enough data to easily train a foundation model, and those who do are unlikely to share it as it gives them an edge over their competition. The only players with both the access to the data and the incentive to make these models available to others are the imagery providers themselves, i.e. Planet, Airbus, Maxar, Blacksky, Capella etc.

Do I think these providers will just open these models for all to use? I doubt it, but they might offer it as a perk to their customers. I.e. something along the lines of “buy at least X Euro worth of imagery per year and get access to our FM”.

The competition in the 30 cm class imagery space seems to be heating up, with several players building up large constellations of satellites in this resolution range, like Maxar’s Legion, Planet’s Pelican and BlackSky’s Gen-3. One way these providers can differentiate their offerings would be by offering a foundation model trained on their specific sensor. Whether I think it’s likely that they do this is another question.

Afterword

Please take this post for what it is: the opinionated rant of someone who works in a somewhat privileged niche of the EO domain where I have a lot of expensive VHR data to play with. The problems I am trying to solve and the constraints I have are likely quite different from those that others might encounter.

With that being said, if you find yourself if a similar boat to me and are wondering which foundation model to pick for your particular task: pick the one trained on the closest thing you can find to the sensor you are targeting.

I am kind of hoping that someone does prove me wrong, and I will happily write an apology post if some does so.

Footnotes

The one exception here is Umbra, who have a very generous open data program, and probably have enough data there that anyone can just train a decently sized model on their data.↩︎

Citation

BibTeX citation:

@online{schenck2025,
  author = {Schenck, Ferdinand},
  title = {Can Multi-Sensor Foundation Models Be More Than the Sum of
    Their Parts?},
  date = {2025-06-22},
  url = {https://fnands.com/blog/2025/multi-sensor-gfms/},
  langid = {en}
}

For attribution, please cite this work as:

Schenck, Ferdinand. 2025. “Can Multi-Sensor Foundation Models Be More Than the Sum of Their Parts?” June 22. https://fnands.com/blog/2025/multi-sensor-gfms/.

A quick first look at GPU programming in Mojo

Ferdinand Schenck — Sat, 12 Apr 2025 00:00:00 GMT

The day has finally arrived. Well actually, the day arrived in February, but who’s counting. The Mojo language has finally publicly released the ability to do GPU programming - if you have a reasonably modern NVIDIA GPU. Luckily for me, I have an RTX 3090, and although it isn’t officially supported, it is basically an A10, which is. Looking at some of the comments on the nightly releases, it does seem that AMD support is on the way as well.

The Modular team publicly released the ability to do GPU programming in Mojo in release 25.1, with further support and documentation in release 25.2. Fun fact: release 25.2 also saw my first (tiny) contribution to the Mojo standard library.

This is a really important step for Mojo, a language that bills itself as a language designed to solve a variety of AI development challenges, which in this day and age basically means programming an increasingly heterogeneous stack of hardware. Today this mostly means GPUs, but there is an explosion of new accelerators like the ones from Cerebras, Groq and SambaNova, not to mention the not-so-new TPU from Google.

As DeepSeek showed the world recently: if you’re willing to put the work in, there is a lot more to be squeezed out of current-gen hardware than most people thought. Now, I don’t think every ML engineer or researcher should be looking for every possible way to get more out of their compute, but there are definitely some wins to be had. As an example, I’m really fascinated by the work of Tri Dao and his collaborators, who work on deeply hardware aware improvements in machine learning, e.g. FlashAttention, which is mathematically equivalent to the attention mechanism that powers all transformer models, but with hardware aware optimizations that take into account the cost of memory access in GPUs. This does make me wonder what other optimizations are out there to be discovered. This however is not easy, as the authors note in the “Limitations and Future Directions” section of the FlashAttention paper:

Our current approach to building IO-aware implementations of attention requires writing a new CUDA kernel for each new attention implementation. This requires writing the attention algorithm in a considerably lower-level language than PyTorch, and requires significant engineering effort. Implementations may also not be transferrable across GPU architectures. These limitations suggest the need for a method that supports writing attention algorithms in a high-level language (e.g., PyTorch), and compiling to IO-aware implementations in CUDA

What makes GPU programming in Mojo interesting is that you don’t need the CUDA toolkit to do so, and compiles down to PTX which you can think of as NVIDIA’s version of assembly. If Mojo (and Max in general) can make it easier to write GPU kernels in a more user-friendly language, it could be a game changer.

If you want to get started, there is a guide for getting started with GPU programming in Mojo from Modular (the company behind Mojo), which I strongly recommend.

A little example

I learn by doing, so I wanted to try to implement something with relatively simple using the GPU.

The example idea I chose is to transform an RGB image to grayscale, which is an embarrassingly parallel problem without a lot of complexity. I was halfway through writing this post before I realized that there was already an example of how to do grayscale conversion in the Mojo repo, but oh well. I basically just start with what’s in the documentation, but I added another example that I did do myself.

To start, let’s read in an image using mimage, an image processing library I am working on.

Let’s take a look at our input image:

import mimage

original_image = mimage.imread("treed_brain_512.png")
print(original_image.spec())

512x512x4xuint8

The image is represented here as a rank three tensor with the dimensions being width, height and channels, and the data type is an unsigned 8-bit integer. In this case we have four channels: red, green, blue and alpha (transparency), the latter being 255 for all pixels.

print(original_image[0], original_image[1], original_image[2], original_image[3])

73 63 97 255

So what we want to do here is to sum together the RGB values for each pixel, using the weights 0.299, 0.587 and 0.114 for red, green and blue respectively. If you want to know why we are using these weights, read this article.

The first thing we need to do is convert the tensor to float32 to prevent any overflow. I’m throwing away the alpha channel for simplicity:

Code: some imports

Speeding up CRC-32 calculations in Mojo

Ferdinand Schenck — Sat, 05 Oct 2024 00:00:00 GMT

In a previous post on parsing PNG images in I very briefly mentioned cyclic redundancy checks, and posted a rather cryptic looking function which I claimed was a bit inefficient.

In this post I want to follow up on that a bit and see how we can speed up these calculations.

For reference, this post was done with Mojo 24.5, so a few language details have changed since my last post (e.g. math.bit got moved to the top-level bit and a few of it’s functions have been renamed).

I actually wrote most of this post in June with Mojo 24.4, but ran into a bit of an issue which has now been resolved. It even resulted in a new unit test for Mojo, so thanks to Laszlo Kindrat for fixing the issue, and for soraros and martinvuyk for helping figure out what the actual issue was.

A bit of context

But first, let’s go through a bit of background so we know what we’re dealing with.

Cyclic redundancy checks

CRCs are error detecting codes that are often used to detect corruption of data in digital files, an example of which is PNG files. In the case of PNGs for example the CRC-32 is calculated for the data of each chunk and appended to the end of the chunk, so that the person reading the file can verify whether the data they read was the same as the data that was written.

A CRC check technically does “long division in the ring of polynomials of binary coefficients ()” 😳.

It’s not as complicated as it sounds. I found the Wikipedia article on Polynomial long division to be helpful, and if you want an in depth explanation then this post by Kareem Omar does an excellent job of explaining both the concept and implementation considerations. I won’t go deep into the explanations, so I recommend you read at least the first part of Kareem’s post for more background. I pretty much use his post as a guide.

Did you read that post? Then welcome back, and well continue from there.

But tl;dr: XOR is equivalent to polynomial long division (over a finite field) for binary numbers, and XOR is a very efficient operation to calculate. Essentially what a CRC check does in practice is to run through a sequence of bytes, and iteratively perform a lot of XORs and bit-shifts. By iteratively going through each bit, one can come up with a value that will (nearly certainly) change if data is corrupted somehow.

The original

The CRC-32 check from my previous post looked something like this:

from bit import bit_reverse

fn CRC32(data: List[UInt8], dummy_table: List[UInt32]) -> UInt32:
    var crc32: UInt32 = 0xffffffff
    for byte in data:
        crc32 = (bit_reverse(byte[]).cast[DType.uint32]() << 24) ^ crc32
        for _ in range(8):
            
            if crc32 & 0x80000000 != 0:
                crc32 = (crc32 << 1) ^ 0x04c11db7
            else:
                crc32 = crc32 << 1
                
    return bit_reverse(crc32^0xffffffff)

I’ll step through this in a moment, but the first thing you might notice here is that I am reversing a lot of bits here. The table argument is a bit of future proofing, which we won’t need for now, but will become apparent soon. Just ignore it for now.

This is because when I was implementing this function (based off a C example), I implemented a little-endian version of the algorithm, while PNGs are encoded as big-endian. It’s not a huge deal, but it does mean that I am constantly reversing bytes, and then reversing the output again.

The correct bit order

We can make this better by implementing the big-endian version:

fn CRC32_inv(data: List[UInt8], dummy_table: List[UInt32]) -> UInt32:
    """Big endian CRC-32 check using 0xedb88320 as polynomial."""

    # Initialize crc32 as all 1s
    var crc32: UInt32 = 0xffffffff

    # Step though all bytes in bytestream
    for byte in data:

        # XOR new byte with crc32 
        # We need to cast it to a UInt32 first as Mojo doesn't do implicit casts
        crc32 = (byte[].cast[DType.uint32]() ) ^ crc32

        # Step though crc32 8 times
        for _ in range(8):
            
            # If leading bit is 1, bitshift by 1 and XOR with the generator polynomial
            # otherwise just shift by one 
            if crc32 & 1 != 0:
                crc32 = (crc32 >> 1) ^ 0xedb88320
            else:
                crc32 = crc32 >> 1

    # Invert upon return
    return crc32^0xffffffff

This is very similar, and just entails that we use a the reverse of the polynomial we did before (if you bit reverse 0x04c11db7 you get 0xedb88320). This also saves us one 24-bit bit-shift, as we are now working on the bottom 8 bits of the UInt32 instead of the top 8.

Just to verify that these implementations are equivalent, let’s do a quick test:

var example_list = List[UInt8](5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201,
                                           5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201,
                                           5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201,
                                           5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201, 
                                           5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201,
                                           5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201,
                                           5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201,
                                           5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201, 42)

print(hex(CRC32(example_list, List[UInt32](1))))
print(hex(CRC32_inv(example_list, List[UInt32](1))))

0x382aa34e
0x382aa34e

And there we go, a more elegant version of the CRC-32 check I implemented last time.

As the theme of today’s post is trying to speed things up, let’s do a little bit of benchmarking to see if this change has saved us any time. As we are doing one fewer bit reverse and bit shift per byte, as well as a final reverse, we should see a bit of a performance uplift.

A bit of benchmarking:

So let’s define a benchmarking function. This function will take two version of the CRC32 function and benchmark their runtimes.

We have to a little bit of work first. The functions we test need to return None, so we need to wrap our functions in functions with no return value. Note the benchmark.keep(a) calls: the compiler will realize that a is never used and will compile this away unless you instruct it to keep them:

import benchmark
from testing import assert_equal

alias data_table_func = fn(data: List[UInt8], table: List[UInt32]) -> UInt32


fn wrap_func[testing_function: data_table_func, data: List[UInt8], table: List[UInt32] ]():
    var a =  testing_function(data = data, table = table)
    benchmark.keep(a)


fn bench[function_1: data_table_func, function_2: data_table_func,
         test_list: List[UInt8], test_table: List[UInt32]](quiet: Bool = False) raises -> Tuple[Float64, Float64, Float64]:


    var CRC32_1 = function_1(test_list, test_table) 

    alias function_1_wrapped = wrap_func[function_1]
    var report_1 = benchmark.run[function_1_wrapped[test_list, test_table]](max_runtime_secs=1.0
    ).mean(benchmark.Unit.ms)
    

    var CRC32_2 = function_2(test_list, test_table) 
    
    assert_equal(CRC32_1, CRC32_2)

    alias function_2_wrapped = wrap_func[function_2]
    var report_2 = benchmark.run[function_2_wrapped[test_list, test_table]](max_runtime_secs=1.0
    ).mean(benchmark.Unit.ms)

    var speedup = report_1/report_2

    if not quiet:
        print("Function 1 runtime (ms): \t", report_1)
        print("Function 2 runtime (ms): \t", report_2)
        print("Speedup factor: \t\t", (speedup))


    return report_1, report_2, speedup

Next we need a test case. I’m not sure if there is a nicer way to fill a List with random values yet, but for now we can list create an UnsafePointer, alloc some space, and fill it with random numbers using rand. Then we can init a List using that data:

from random import rand

alias fill_size = 2**20
alias g = UnsafePointer[UInt8].alloc(fill_size)
rand[DType.uint8](g, fill_size)
alias rand_list = List[UInt8](unsafe_pointer = g, size = fill_size, capacity = fill_size)

alias dummy_table = List[UInt32](1)

And finally we are ready to benchmark:

original_time, improved_time, speedup = bench[CRC32, CRC32_inv, rand_list, dummy_table]()

Function 1 runtime (ms):     9.437418430000001
Function 2 runtime (ms):     6.8247150400000001
Speedup factor:          1.3828296675666039

Nice! So just by avoiding a few bit reversals and bit shifts we get about a 30% uplift in performance depending on the run.

Note: I am doing this in a Jupyter notebook, so there is a bit of variance from run to run.

A quick Python speed-check

While we’re checking performance, let’s see how this implementation would perform in Python:

%%python

def py_crc32(data: bytearray) -> int: 
    crc32 = 0xffffffff
    
    for byte in data:
        crc32 = byte ^ crc32

        for i in range(8):
            if crc32 & 1 != 0:
                crc32 = (crc32 >> 1) ^ 0xedb88320
            else:
                crc32 = crc32 >> 1

    return crc32^0xffffffff

And let’s just do a sanity check to assure ourselves that we produce the same CRC-32 value given the same bytestream:

%%python

py_example_list = [5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201,
    5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201,
    5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201,
    5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201, 
    5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201,
    5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201,
    5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201,
    5, 78, 138, 1, 54, 172, 104, 99, 54, 167, 94, 56, 22, 184, 204, 90, 201, 42]
    
py_example_bytearray = bytearray(py_example_list)

print(hex(py_crc32(py_example_bytearray)))

0x382aa34e

%%python
from timeit import timeit
import numpy as np


py_rand_list = np.random.randint(0, 255, size=2**20, dtype=np.uint8)

naive_python_seconds = timeit(lambda: py_crc32(py_rand_list), number=10)/10

print(f"Naive Python runtime (ms): {naive_python_seconds*100}")

Naive Python runtime (ms): 123.02583360986318

That’s pretty slow in fact this means that:

print("Speedup over Python: Mojo is", (naive_python_seconds*100/improved_time), "faster")

Speedup over Python: Mojo is 18.026515816235925 faster

So that’s a nice little speedup we can get by writing essentially the same logic in Mojo.

Now this is a bit unfair. If you actually wrote the this function in Python for anything other than educational purposes, then you are using Python wrong, but I’ll get back to how you would actually do it later.

Adding a table

Now, the majority of CRC-32 implementations you see will usually use a table.

This is called the Sarwate algorithm and exploits the fact that as we are operating on data one byte at a time, there are only 256 unique values that the calculation in the inner loop of the algorithm can take. This is a small enough table (1024 bytes for 256 unsigned 32-bit integers) that it can easily be kept in the L1 cache during computation.

Let’s try it out to see how much faster this gets us.

First, we need to pre-compute the table:


fn fill_table() -> List[UInt32]:
    var table = List[UInt32](capacity=256)
    for i in range(256):

        var key = UInt8(i)
        var crc32 = key.cast[DType.uint32]()
        for i in range(8):
            if crc32 & 1 != 0:
                crc32 = (crc32 >> 1) ^ 0xedb88320
            else:
                crc32 = crc32 >> 1

        table[i] = crc32
    return table

alias little_endian_table = fill_table()

We’ve now effectively amortized the cost of the inner loop and replaced it with a table lookup which should in principle be faster.

Our CRC-32 algorithm now take this form:

fn CRC32_table(data: List[UInt8], table: List[UInt32]) -> UInt32:
    var crc32: UInt32 = 0xffffffff
    
    for byte in data:
        var index = (crc32 ^ byte[].cast[DType.uint32]()) & 0xff
        crc32 = table[int(index)] ^ (crc32 >> 8)

    return crc32^0xffffffff

# Sanity check
print(hex(CRC32(example_list, List[UInt32](1))))
print(hex(CRC32_table(example_list, little_endian_table)))

0x382aa34e
0x382aa34e

Now we can test it out!

_, _, speedup_table = bench[CRC32, CRC32_table, rand_list, little_endian_table]()

Function 1 runtime (ms):     8.5761109999999992
Function 2 runtime (ms):     1.9543971584601614
Speedup factor:          4.3881106574863127

So a speedup of around 4-5 times as fast as when we started. This is already a pretty good result, but can we do better?

Two byte table

In principle, we could load the data as a 16-bit integer at a time and use the same trick above to build a table. However, such a table would have 65536 entries, resulting in a 256 KB table, which is a lot bigger than the 48 KB of L1 cache on my machine.

It turns out, we can still do two bytes at a time.

Following the description given by Kareem Omar, we realize that if we have two successive bytes in our message, let’s call them X and Y:

X = xxxxxxxx
Y = yyyyyyyy

if we pad these messages with zeroes

Xz = xxxxxxxx00000000
zY = 00000000yyyyyyyy

then

Xz ^ zY = xxxxxxxxyyyyyyyy = M

What this means is that if we have a CRC algorithm that works on a 16-bit message called M then:

CRC(M) = CRC(Xz ^ zY) = CRC(Xz) ^ CRC(zY)

For the operation on Y, leading zeros don’t affect the calculation, so we can just use the 8-bit CRC algorithm we have developed above:

CRC(zY) = CRC(Y)

The same does not hold for trailing zeroes. However, as Xz will always be of the form xxxxxxxx00000000 and there can only be 256 unique values for Xz, we can build a new table with 256 entries to look up these values. We can build two separate 256 entry tables, or we can just build one 512 entry table,

So let’s construct this new table:

fn fill_table_2_byte() -> List[UInt32]:

    var table = List[UInt32](capacity=512)
    table.size = 512
    table.resize(256*2, 0)

    for i in range(256):

        var key = UInt8(i)
        var crc32 = key.cast[DType.uint32]()
        for i in range(8):
            if crc32 & 1 != 0:
                crc32 = (crc32 >> 1) ^ 0xedb88320
            else:
                crc32 = crc32 >> 1

        table[i] = crc32

    for i in range(256, 512):
        var crc32 = table[i-256]
        table[i] = (crc32 >> 8) ^ table[int(crc32.cast[DType.uint8]())]
    return table

var little_endian_table_2_byte = fill_table_2_byte()

We have to now update our algorithm to do two table lookups instead of one.

Additionally, we have to modify our algorithm to read the data two bytes at a time.

fn CRC32_table_2_byte(data: List[UInt8], table: List[UInt32]) -> UInt32:
    var crc32: UInt32 = 0xffffffff

    var length = len(data)//2
    var extra = len(data) % 2

    # Go through range in steps of 2
    for i in range(start = 0, end = length * 2 , step = 2):
        
        # Concatenate the two bytes
        var val: UInt32 = ((data[i + 1].cast[DType.uint32]() << 8) | data[i].cast[DType.uint32]())
        # Creat index, and look up values in table
        var index = crc32 ^ val
        crc32 =  table[int((index >> 8).cast[DType.uint8]())] ^ table[256 + int(index.cast[DType.uint8]())] ^ (crc32 >> 16)

    # Tail call to calculate final byte in case message not divisable by 2
    for i in range(2*length, 2*length + extra ):
        var index = (crc32 ^ data[i].cast[DType.uint32]()) & 0xff
        crc32 = table[int(index)] ^ (crc32 >> 8)


    return crc32^0xffffffff

# Sanity check
print(hex(CRC32(example_list, List[UInt32](1))))
print(hex(CRC32_table_2_byte(example_list, little_endian_table_2_byte)))

0x382aa34e
0x382aa34e

So let’s see if this speeds things up:

alias a_little_endian_table_2_byte = fill_table_2_byte()


_, _, speedup_table_2 = bench[CRC32, CRC32_table_2_byte, rand_list, a_little_endian_table_2_byte]()

Function 1 runtime (ms):     9.4386268461538467
Function 2 runtime (ms):     1.2856673344847436
Speedup factor:          7.3414223049670637

Nice, now we are at an around 7x speedup to where we began.

But can we go further? There’s nothing in the above that forces us to only use two bytes at a time.

4-bytes tabular

Nothing stops us from doing 4 bytes, i.e. 32-bits at a time with similar logic to above.

I’ll quickly create a function that will fill an arbitrarily sized table:

fn fill_table_n_byte[n: Int]() -> List[UInt32]:
    var table = List[UInt32](capacity=256 * n)
    table.size = 256*n

    for i in range(256 * n):
        if i < 256:
            var key = UInt8(i)
            var crc32 = key.cast[DType.uint32]()
            for _ in range(8):
                if crc32 & 1 != 0:
                    crc32 = (crc32 >> 1) ^ 0xEDB88320
                else:
                    crc32 = crc32 >> 1

            table[i] = crc32
        else:
            var crc32 = table[i - 256]
            var index = int(crc32.cast[DType.uint8]())
            table[i] = (crc32 >> 8) ^ table[index]

    return table

var little_endian_table_4_byte = fill_table_n_byte[4]()

Now to create the 4-byte function:

fn CRC32_table_4_byte(data: List[UInt8], table: List[UInt32]) -> UInt32:
    var crc32: UInt32 = 0xffffffff

    var size = 4

    var length = len(data)//size
    var extra = len(data) % size

    for i in range(start = 0, end = length*size, step = size):
        
        var val: UInt32 =  (data[i + 3].cast[DType.uint32]() << 24) | (data[i + 2].cast[DType.uint32]() << 16) | (data[i + 1].cast[DType.uint32]() << 8) | data[i].cast[DType.uint32]()
        var index = crc32 ^ val.cast[DType.uint32]()
        crc32 = table[0*256 + int((index >> 24).cast[DType.uint8]())] ^
                table[1*256 + int((index >> 16).cast[DType.uint8]())] ^
                table[2*256 + int((index >> 8).cast[DType.uint8]())] ^
                table[3*256 + int((index >> 0).cast[DType.uint8]())] 
    
    for i in range(size*length, size*length + extra ):
        var index = (crc32 ^ data[i].cast[DType.uint32]()) & 0xff
        crc32 = table[int(index)] ^ (crc32 >> 8)


    return crc32^0xffffffff

# Sanity check
print(hex(CRC32(example_list, List[UInt32](1))))
print(hex(CRC32_table_4_byte(example_list, little_endian_table_4_byte)))

0x382aa34e
0x382aa34e

And presto, it still works!

Let’s see how much faster we are now?

alias a_little_endian_table_4_byte = fill_table_n_byte[4]()


_, _, _ = bench[CRC32, CRC32_table_4_byte, rand_list, a_little_endian_table_4_byte]()

Function 1 runtime (ms):     9.9107345799999997
Function 2 runtime (ms):     0.69239487
Speedup factor:          14.313703075240866

A 14 times speedup over where we started!

But why stop there?

n-bytes implementation

We can in principle explicitly write out version that take 8, 16 or however many bytes at time. This get’s a little long-winded, so I’ll write some generic functions to make functions of arbitrary size.

fn CRC32_table_n_byte_compact[size: Int](data: List[UInt8], table: List[UInt32]) -> UInt32:
    var crc32: UInt32 = 0xFFFFFFFF

    alias step_size = 4 # Always smashing 4 bytes into an UInt32
    alias units = size // step_size

    var length = len(data) // size
    var extra = len(data) % size

    var n = 0

    for i in range(start=0, end=length * size, step=size):
        @parameter
        for j in range(units):
            var vals = (
                (data[i + j * step_size + 3].cast[DType.uint32]() << 24)
                | (data[i + j * step_size + 2].cast[DType.uint32]() << 16)
                | (data[i + j * step_size + 1].cast[DType.uint32]() << 8)
                | (data[i + j * step_size + 0].cast[DType.uint32]() << 0)
            )

            if j == 0:
                vals = vals ^ crc32
                crc32 = 0 
        

            n = size - j * step_size
            crc32 = (
                table[(n - 4) * 256 + int((vals >> 24).cast[DType.uint8]())]
                ^ table[(n - 3) * 256 + int((vals >> 16).cast[DType.uint8]())]
                ^ table[(n - 2) * 256 + int((vals >> 8).cast[DType.uint8]())]
                ^ table[(n - 1) * 256 + int((vals >> 0).cast[DType.uint8]())]     
            ) ^ crc32
            

    for i in range(size * length, size * length + extra):
        var index = (crc32 ^ data[i].cast[DType.uint32]()) & 0xFF
        crc32 = table[int(index)] ^ (crc32 >> 8)

    return crc32 ^ 0xFFFFFFFF

Let’s do a quick sanity check to see if this works:

print(hex(CRC32(example_list, List[UInt32](1))))
print(hex(CRC32_table_n_byte_compact[4](example_list, little_endian_table_4_byte)))

0x382aa34e
0x382aa34e

How big can we go?

Let’s try and increase the table size as far as we can. We’ll go up in powers of two, and see how far we can go.

alias a_little_endian_table_4_byte = fill_table_n_byte[4]()

_, _, speedup_table_4 = bench[CRC32, CRC32_table_n_byte_compact[4], rand_list, a_little_endian_table_4_byte]()

Function 1 runtime (ms):     9.5889589700000002
Function 2 runtime (ms):     0.69029748700000004
Speedup factor:          13.891052988868841


_, _, speedup_table_8 = bench[CRC32, CRC32_table_n_byte_compact[8], rand_list, fill_table_n_byte[8]()]()

Function 1 runtime (ms):     10.002244664999999
Function 2 runtime (ms):     0.48423754050000001
Speedup factor:          20.655657251753283

_, _, speedup_table_16 = bench[CRC32, CRC32_table_n_byte_compact[16], rand_list, fill_table_n_byte[16]()]()

Function 1 runtime (ms):     9.065976534999999
Function 2 runtime (ms):     0.3490761703607127
Speedup factor:          25.971341800936472

_, _, speedup_table_32 = bench[CRC32, CRC32_table_n_byte_compact[32], rand_list, fill_table_n_byte[32]()]()

Function 1 runtime (ms):     9.6776499699999992
Function 2 runtime (ms):     0.22259112022425251
Speedup factor:          43.477250845631744

_, _, speedup_table_64 = bench[CRC32, CRC32_table_n_byte_compact[64], rand_list, fill_table_n_byte[64]()]()

Function 1 runtime (ms):     9.1518207500000006
Function 2 runtime (ms):     0.29884062938339201
Speedup factor:          30.6244193397774

_, _, speedup_table_128 = bench[CRC32, CRC32_table_n_byte_compact[128], rand_list, fill_table_n_byte[128]()]()

Function 1 runtime (ms):     9.4991627550000004
Function 2 runtime (ms):     0.53072807599999994
Speedup factor:          17.898361109126629

And there it is. At least on my machine, 32 bytes is the limit, maxing out at a roughly 40 times speedup over the naive implementation. After that we start to see a performance decrease.

Let’s plot this to see the trend:

Code

Building your own personal ghostwriter

Ferdinand Schenck — Thu, 03 Oct 2024 00:00:00 GMT

This past weekend I participated in a AI hackathon organized by Factory Network and {Tech: Berlin} with my friends Axel Nordfeldt and Jonathan Nye. The sponsors were Mistral, Weaviate and LumaAI, meaning we had a bunch of fun credits to play around with.

I’m an ML engineer, but I haven’t had that much exposure to the world brave new world of AI engineering that consists of calling LLM API’s and prompt engineering, so I wanted to use the hackathon as a way to get a better understanding of which tools are out there are how people are using them.

We spitballed a few ideas the week before the event, and eventually settled on building a kind of ghostwriter during the event. The idea wasn’t to get chatbot to write you a story based on a minimal prompt, but rather to have something that interviews you in depth, and uses the interview as the basis of a blog post, short story or article.

The idea was inspired by an experience Axel had where he went on a trip and later recounted it to a friend who is a talented writer. The friend then wrote up a summary of the trip and gave it to Axel, who was impressed by how well the friend had captured the essence of the trip.

Writing about my own travel experiences is something I also wish I could do better. Not to for the sake of anyone else, but just for my own memories. I went to India in 2012 and I have a handful of photos of the trip, but I didn’t write much down at the time. I’d love to have a more detailed account of the trip, but I can’t remember all the details anymore. I recently found a blog post by someone I had met on the trip which even includes a quote from me, which I hardly remember.

I struggle with journaling or keeping diaries in general, so I would like to have a tool that can help me write about my experiences. I don’t know if it is just because writing is hard (part of the reason I keep this blog is to practice writing), or if it is just because I feel awkward writing about myself.

What isn’t that hard however is telling someone about my experiences, especially someone who knows how to ask pertinent questions. So the idea was to build a tool that could ask you questions about an experience, and then write up a summary based on your answers. Effectively the way a ghostwriter would work, but instead of writing a book for you, it would write a blog post or short story.

The build

For the two-day hackathon we had access to the sponsors’ APIs. We especially wanted to play around with Mistral’s API, which had just recently added a new multi-modal model, Pixtral that can understand images as well as text.

So we basically built the app using Streamlit (called JournalAIst) as two parts: an interviewer and a writer. Basically we had a chat loop in which you could upload images (interpreted by Pixtral) and during which the interviewer model (mistral-large-2407) would ask you questions about your experience and the images that you added. After answering the interviewers questions for a while, you can choose to end the conversation, after which you are sent to the writer.

Interviewer

The writer will then be given the transcript of your conversation, as well as Pixtral’s descriptions of any images you might have uploaded. The model then uses this context to write a few paragraphs worth of text about your experience. You get to choose if you want a blog post, short story or article, each of which having a slightly different tone and viewpoint.

Writer

We had a little bit of time left at the end so we ended up also adding a section where we fed a summary of your post to LumaAI’s Dream Machine model, which generates a five second video based on this input, which was a cute little addition. Try and imagine the story that resulted in this video:

Getting the whole thing up and running was surprisingly easy given that none of us had any experience with Streamlit (shout-out to Vindiw Wijesooriya who’s mistral-streamlit-chat repo pointed us in the right direction).

The story writer worked pretty well off the bat, but getting the interviewer to work well was the biggest challenge. None of us had a lot of experience prompt engineering, and the interviewer kept repeating questions (like asking who you were with three times in a row), so messing with the interviewer prompt was one of the most time consuming steps. If you have ever been interviewed by a good interviewer you know that someone skilled at the craft of interviewing can steer a conversation in interesting directions, and that is the feeling we were going for here. Easier said than done. It definitely worked best when your initial response contains quite a bit of context, hence why our initial question to the user asks for at least a few sentences. Eventually we can up with something that gave a reasonable experience, but I feel there is a lot to be done there. Maybe coming up with a set list of questions, and only having the model do follow up questions?

We made the story writer output the story it writes as markdown so they can easily be viewed and used as blog posts. As for the images we did something janky: we saved them as image_n.jpg where n is the order in which the images were uploaded, and told the model to refer to them in that way in markdown. This felt a little weird at first, but it seems to work well enough. The model will just dump a ![Picture description](picture_n.jpg) in the markdown when it wanted to refer to the images, and it worked most of the time.

A few sample stories.

In case you were wondering, no I didn’t use JournalAIst to write this blog post. The language is clearly not flowery enough. To give you an idea of the quality of the writing, here are a few samples:

Story 1: A Wild Rendezvous in Kruger National Park

I recounted my experience of being in the Kruger National Park with a couple of friends last year. For this case I basically just threw in a handful of details and sent the model off to write. See the video at the end that we created while the LumaAI credits were still good. The video seems to try and combine a Lion, Elephant and Zebra, with the Elephant getting the Elephant’s share of the mix. Also, note the abomination in the background.

Blog Post: Kruger Story 1

Despite the minimal details, it did summarise what I told it quite truthfully, but with a healthy dash of breathless wonder, which doesn’t sound like me. This might be as the prompt was write in the style of Bill Bryson, but even there I wouldn’t say it sounds a lot like him either. The prompt also requests the model to write in a “humorous and engaging” manner, so I guess that might be the reason it sounds as it does.

Story 2: A Wild Ride through Kruger: A Tale of Elephants, Lions, and Fearless Honey Badgers

I tried giving it another go (caveat, this one was generated after the free credits expired, so we switched the model to mistral-small-2409), but this time added more context and chatted for a bit longer.

Blog Post: Kruger Story 2

Again, the story is reasonable, although I wished it was maybe a bit longer? Or maybe it is mercifully short, depending on your perspective.

Story 3: A Weekend of Wonders: Our Hackathon Adventure

I also asked it to write up my experience of the hackathon with a few pictures I took during the event. I tried to continue the interview for a while to give as much info as possible.

Journal Entry: Hackathon Story

The story is OK. Everything it stated was factual, and as the prompt this time was to write a journal entry and not a blog post, it’s not as fantastical as before. It’s a little short though, and I could likely have written the same if not more myself.

Conclusions

Our project was one of the few that actually fully worked by the end of the hackathon, and our first showing really impressed the judges (they told me afterwards). On the strength of that we made it to the final six teams (out of about 25 that decided to enter). The pitch went fine, and we again did a live demo in the four minutes allotted to us. We didn’t end up winning any prizes, but we were just happy to be in the finals. Maybe we needed a salesperson on our team and not just a bunch of engineers 😉.

I spoke to some of the judges afterwards, and one of the criticisms was that while initially impressive, they had worries about hallucinations: in our pitch we had highlighted that this is a way to hopefully tell an authentic and meaningful story, but with hallucinations, you might spend as much time fixing the story as you do talking to the model. This is extremely fair criticism and was also mirrored our worries. It did seem the longer you chatted to it the more factual the story ended up being (unsurprisingly), but at what point do you spend more time chatting then it would just have taken you to write it?

On the last point: I would still argue that the job of a ghostwriter is to take your unstructured thoughts and edit them down into something coherent.

My feeling is that the chat interface, or at least typing is not as free as I would have liked. One of the judges suggested a voice interface, which is something we considered as well (i.e. using something like Whisper to turn speech into text), but would have been hard to get working in the little time we had.

My feeling is that there is something interesting here, and that making the interaction with the interviewer more natural is probably where we should focus our efforts.

The quality of the interviewer questions could also be better. Maybe we can prompt it by giving it some transcripts of interviews by Louis Theroux or some other good interviewer.

Code

After the hackathon the API keys provided expired, but Mistral recently opened up a free tier, so we bumped the main model down to one of the smaller ones supported on the free tier.

In any case, the code is here if you wish to play around with it. The code is what I would call “Hackathon Quality” so beware. You’ll need a Mistral API key to run it, but it will work with the free tier, so won’t cost you anything.

It’s also deployed on Streamlit cloud, which might work depending on whether or not we’ve hit the limits of Mistral’s free tier.

Citation

BibTeX citation:

@online{schenck2024,
  author = {Schenck, Ferdinand},
  title = {Building Your Own Personal Ghostwriter},
  date = {2024-10-03},
  url = {https://fnands.com/blog/2024/factory-hackathon/},
  langid = {en}
}

For attribution, please cite this work as:

Schenck, Ferdinand. 2024. “Building Your Own Personal Ghostwriter.” October 3. https://fnands.com/blog/2024/factory-hackathon/.

Parsing PNG images in Mojo

Ferdinand Schenck — Sat, 04 May 2024 00:00:00 GMT

So for the past while I’ve been trying to follow along with the development of Mojo, but so far I’ve mostly just followed along with the changelog and written some pretty trivial pieces of code. In my last post I said I wanted to try something a bit more substantial, so here goes.

I was looking at the Basalt project, which tries to build a Machine Learning framework in pure Mojo, and realized that the only images used so far were MNIST, which come in a weird binary format anyway. Why no other though? As Mojo does not yet support accelerators (like GPUs) Imagenet is probably impractical, but it should be fairly quick to train a CNN on something like CIFAR-10 on a CPU these days. The CIFAR-10 dataset is available from the original source as either a pickle archive or some custom binary format. I thought about writing decoders for these, but it might be more useful to write a PNG parser in Mojo, and then use the version of the dataset hosted on Kaggle and other places, or just transform the original to PNGs using this package. That way the code can be used to open PNG images in general.

A PNG parser in (pure-ish) Mojo

Don’t mistake this post for a tutorial: read it as someone discovering the gory details of the PNG standard while learning a new language. If you want to read more about the PNG format, the wikipedia page is pretty helpful as an overview, and the W3C page provides a lot of detail.

For reference, this was written with Mojo 24.3.0, and as Mojo is still changing pretty fast a lot of what is done below might be outdated. I actually I did basically the entire post in 24.2.1 but 24.3 got released just before I published, but it required only minor changes to make it work in the new version.

The goal here is not to build a tool to display PNGs, but just to read them into an array (or tensor) that can be used for ML purposes, so I will skip over a lot of the more display oriented details.

Reading in the data

To start let’s take a test image. This is the test image from the PIL library, which is an image of the OG programmer Grace Hopper: This is a relatively simple PNG, so it should be a good place to start.

hopper

Now that Mojo has implemented it’s version of pathlib in the stdlib, we can actually check if the file exists:

from pathlib import Path
test_image = Path('hopper.png')
print(test_image.exists())

True

We’ll also import the image via Python so we can compare if the outputs we get match the Python case.

from python import Python
var Image = Python.import_module('PIL.Image')
var np = Python.import_module('numpy')
py_array = np.array(Image.open("hopper.png"))

We’re going to read the raw bytes. I would have expected the data to be unsigned 8-bit integers, but Mojo reads them as signed 8-bit integers. There is however a proposal to change this, so this might change soon.

with open(test_image, "r") as f:
    file_contents = f.read_bytes()

print(len(file_contents))

Checking the file header

PNG files have a signature defined in the first 8-bytes, part of which is the letters PNG in ASCII. Well define a little helper function to convert from bytes to String:

fn bytes_to_string(list: List[Int8]) -> String:
    var word = String("")
    for letter in list:
        word += chr(int(letter[].cast[DType.uint8]()))

    return word

To make sure we are actually dealing with a PNG, we can check the bits 1 to 3:

png_signature = file_contents[0:8]
print(bytes_to_string(png_signature[1:4]))

PNG

Yup, it’s telling us it is a PNG file.

Reading chunks

So now we read the first “chunk”, which should be the header. Each chunk consists of four parts, the chunk length (4 bytes), the chunk type (4 bytes), the chunk data (however long the first 4 bytes said it was), and a checksum (called the CRC) computed from the data (4 bytes).

Length	Chunk type	Chunk data	CRC
4 bytes	4 bytes	Length bytes	4 bytes

When reading in data with read_bytes, the data comes as a list of signed 8-bit integers, but we would like to interpret the data as 32-bit unsigned integers. Below is a helper function to do so (thanks to Michael Kowalski) for the help.

from math.bit import bswap, bitreverse
from testing import assert_true


fn bytes_to_uint32_be(owned list: List[Int8]) raises -> List[UInt32]:
  assert_true(len(list) % 4 == 0, "List[Int8] length must be a multiple of 4 to convert to List[Int32]")
  var result_length = len(list) // 4
  
  # get the data pointer with ownership.
  # This avoids copying and makes sure only one List owns a pointer to the underlying address.
  var ptr_to_int8 = list.steal_data() 
  var ptr_to_uint32 = ptr_to_int8.bitcast[UInt32]()

  var result = List[UInt32]()
  result.data = ptr_to_uint32
  result.capacity = result_length
  result.size = result_length

  # swap the bytes in each UInt32 to convert from big-endian to little-endian
  for i in range(result_length):
    result[i] = bswap(result[i])

  return result

Reading the image header

The firs chunk after the file header should always be the image header, so let’s have a look at it:

Let’s see how long the first chunk is:

read_head = 8
chunk_length = bytes_to_uint32_be(file_contents[read_head:read_head+4])[0]
print(chunk_length)

So the first chunk is 13 bytes long. Let’s see what type it is:


chunk_type = file_contents[read_head+4:read_head+8]
print(bytes_to_string(chunk_type))

IHDR

IHDR, which confirms that this chunk is the image header. We can now parse the next 13 bytes of header data to get information about the image:

start_header = int(read_head+8)
end_header = int(read_head+8+chunk_length)
header_data = file_contents[start_header:end_header]

The first two chunks tell us the width and height of the image respectively:

print("Image width: ", bytes_to_uint32_be(header_data[0:4])[0])
print("Image height: ", bytes_to_uint32_be(header_data[4:8])[0])

Image width:  128
Image height:  128

So our image is 128x128 pixels in size.

The next bytes tell is the bit depth of each pixel, color type, compression method, filter method, and whether the image is interlaced or not.

print("Bit depth: ", int(header_data[8]))
print("Color type: ", int(header_data[9]))
print("Compression method: ", int(header_data[10]))
print("Filter method: ", int(header_data[11]))
print("Interlaced: ", int(header_data[12]))

Bit depth:  8
Color type:  2
Compression method:  0
Filter method:  0
Interlaced:  0

So the color type is Truecolor, so RGB, with a bit depth of 8.

Interesting side note: in the PIL PngImagePlugin there is a changelog:

# history:
# 1996-05-06 fl   Created (couldn't resist it)
# 1996-12-14 fl   Upgraded, added read and verify support (0.2)
# 1996-12-15 fl   Separate PNG stream parser
# 1996-12-29 fl   Added write support, added getchunks
# 1996-12-30 fl   Eliminated circular references in decoder (0.3)
# 1998-07-12 fl   Read/write 16-bit images as mode I (0.4)
# 2001-02-08 fl   Added transparency support (from Zircon) (0.5)
# 2001-04-16 fl   Don't close data source in "open" method (0.6)
# 2004-02-24 fl   Don't even pretend to support interlaced files (0.7)
# 2004-08-31 fl   Do basic sanity check on chunk identifiers (0.8)
# 2004-09-20 fl   Added PngInfo chunk container
# 2004-12-18 fl   Added DPI read support (based on code by Niki Spahiev)
# 2008-08-13 fl   Added tRNS support for RGB images
# 2009-03-06 fl   Support for preserving ICC profiles (by Florian Hoech)
# 2009-03-08 fl   Added zTXT support (from Lowell Alleman)
# 2009-03-29 fl   Read interlaced PNG files (from Conrado Porto Lopes Gouvua)

I like the comment from 2004: Don't even pretend to support interlaced files and then interlaced PNG being supported about 13 years after PNG reading was added to PIL. I have a feeling I won’t be dealing with interlaced files in this post…

The final part of this chunk is the CRC32 value, which is the 32-bit cyclic redundancy check. I don’t go into too much details, but it’s basically an error-detecting code that’s added to detect if the chunk data is corrupt. By checking the provided CRC32 value against one we calculate ourselves we can ensure that the data we are reading is not corrupt.

start_crc = int(read_head+8+chunk_length)
end_crc = int(start_crc+4)
header_crc = bytes_to_uint32_be(file_contents[start_crc:end_crc])[0]
print("CRC: ", hex(header_crc))

CRC:  0x4c5cf69c

We need a little bit of code to calculate the CRC32 value.
This is not the most efficient implementation, but it is simple.
I’ll probably do a follow up post where I explain what this does in more detail.

fn CRC32(owned data: List[SIMD[DType.int8, 1]]) -> SIMD[DType.uint32, 1]:
    var crc32: UInt32 = 0xffffffff
    for byte in data:
        crc32 = (bitreverse(byte[]).cast[DType.uint32]() << 24) ^ crc32
        for i in range(8):
            
            if crc32 & 0x80000000 != 0:
                crc32 = (crc32 << 1) ^ 0x04c11db7
            else:
                crc32 = crc32 << 1

    return bitreverse(crc32^0xffffffff)

print(hex(CRC32(file_contents[read_head+4:end_header])))

0x4c5cf69c

Great, so the CRC hexes match, so we know that the data in our IHDR chunk is good.

Reading more chunks

Now, reading parts of each chunk will get repetitive, so let’s define a struct called Chunk to hold the information contained in a chunk, and a function that will parse chunks for us and return the constituent parts:

struct Chunk(Movable, Copyable):
    var length: UInt32
    var type: String
    var data: List[Int8]
    var crc: UInt32
    var end: Int

    fn __init__(inout self, length: UInt32, chunk_type: String, data : List[Int8], crc: UInt32, end: Int):
        self.length = length
        self.type = chunk_type
        self.data = data
        self.crc = crc
        self.end = end

    fn __moveinit__(inout self, owned existing: Chunk):
        self.length = existing.length
        self.type = existing.type
        self.data = existing.data
        self.crc = existing.crc
        self.end = existing.end

    fn __copyinit__(inout self, existing: Chunk):
        self.length = existing.length
        self.type = existing.type
        self.data = existing.data
        self.crc = existing.crc
        self.end = existing.end


def parse_next_chunk(owned data: List[Int8], read_head: Int) -> Chunk:
    chunk_length = bytes_to_uint32_be(data[read_head:read_head+4])[0]
    chunk_type = bytes_to_string(data[read_head+4:read_head+8])
    start_data = int(read_head+8)
    end_data = int(start_data+chunk_length)
    chunk_data = data[start_data:end_data]
    start_crc = int(end_data)
    end_crc = int(start_crc+4)
    chunk_crc = bytes_to_uint32_be(data[start_crc:end_crc])[0]

    # Check CRC
    assert_true(CRC32(data[read_head+4:end_data]) == chunk_crc, "CRC32 does not match")
    return Chunk(length=chunk_length, chunk_type=chunk_type, data=chunk_data, crc=chunk_crc, end=end_crc)

During chunk creation the CRC32 value for the chunk data is computed, and an issue will be raised if it is different to what is expected.

Let’s test this to see if it parses the IHDR chunk:

var header_chunk = parse_next_chunk(file_contents, 8)
print(header_chunk.type)
read_head = header_chunk.end

IHDR

The next few chunks are called “Ancillary chunks”, and are not strictly necessary. They contain image attributes (like gamma) that may be used in rendering the image:

var gamma_chunk = parse_next_chunk(file_contents, read_head)
print(gamma_chunk.type)
read_head = gamma_chunk.end

gAMA

var chromacity_chunk = parse_next_chunk(file_contents, read_head)
print(chromacity_chunk.type)
read_head = chromacity_chunk.end

cHRM

var background_chunk = parse_next_chunk(file_contents, read_head)
print(background_chunk.type)
read_head = background_chunk.end

bKGD

var pixel_size_chunk = parse_next_chunk(file_contents, read_head)
print(pixel_size_chunk.type)
read_head = pixel_size_chunk.end

pHYs

The image data chunk

The IDAT chunk (there can actually be several of them per image) contains the actual image data.

var image_data_chunk = parse_next_chunk(file_contents, read_head)
print(image_data_chunk.type)
read_head = image_data_chunk.end

IDAT

Decompression

PNGs are compressed (losslessly) with the DEFLATE compression algorithm.

PNGs are first filtered, then compressed, but as we are decoding, we need to first uncompress the data and the undo the filter.

This next section is why I said in “pure-ish” Mojo: I considered implementing it, but that would be quite a lot of work, so I am hoping that either someone else does this, or that I might dig into this in the future.

So for the moment, I am using the zlib version of the algorithm through Mojo’s foreign function interface (FFI).

The following I lightly adapted from the Mojo discord from a thread between Ilya Lubenets and Jack Clayton:

from sys import ffi
alias Bytef = Scalar[DType.int8]
alias uLong = UInt64
alias zlib_type = fn(
    _out: Pointer[Bytef], 
    _out_len: Pointer[UInt64], 
    _in: Pointer[Bytef], 
    _in_len: uLong
) -> Int
fn log_zlib_result(Z_RES: Int, compressing: Bool = True) raises -> NoneType:
    var prefix: String = ''
    if not compressing:
        prefix = "un"

    if Z_RES == 0:
        print('OK ' + prefix.upper() + 'COMPRESSING: Everything ' + prefix + 'compressed fine')
    elif Z_RES == -4:
        raise Error('ERROR ' + prefix.upper() + 'COMPRESSING: Not enought memory')
    elif Z_RES == -5:
        raise Error('ERROR ' + prefix.upper() + 'COMPRESSING: Buffer have not enough memory')
    else:
        raise Error('ERROR ' + prefix.upper() + 'COMPRESSING: Unhandled exception')

fn uncompress(data: List[Int8], quiet: Bool = True) raises -> List[UInt8]:
    var data_memory_amount: Int = len(data)*10  # This can be done better. 
    var handle = ffi.DLHandle('')
    var zlib_uncompress = handle.get_function[zlib_type]('uncompress')

    var uncompressed = Pointer[Bytef].alloc(data_memory_amount)
    var compressed = Pointer[Bytef].alloc(len(data))
    var uncompressed_len = Pointer[uLong].alloc(1)
    memset_zero(uncompressed, data_memory_amount)
    memset_zero(uncompressed_len, 1)
    uncompressed_len[0] = data_memory_amount
    for i in range(len(data)):
        compressed.store(i, data[i])

    var Z_RES = zlib_uncompress(
        uncompressed,
        uncompressed_len,
        compressed,
        len(data),
    )

    if not quiet:
        log_zlib_result(Z_RES, compressing=False)
        print('Uncompressed length: ' + str(uncompressed_len[0]))
    # Can probably do something more efficient here with pointers, but eh. 
    var res = List[UInt8]()
    for i in range(uncompressed_len[0]):
        res.append(uncompressed[i].cast[DType.uint8]())
    return res

Drumroll… let’s see if this worked:

uncompressed_data = uncompress(image_data_chunk.data, quiet=False)

OK UNCOMPRESSING: Everything uncompressed fine
Uncompressed length: 49280

Now we have a list of uncompressed bytes. However, these are not pixel values yet. The uncompressed data has a length of 49280 bytes. We know we have an RGB image with 8-bit colour depth, so expect bytes worth of pixel data. Notice that , and that our image has a shape of (128, 128).
These extra 128 bytes are to let us know what filter was used to transform the byte values each line of pixels (knows as scanlines) into something that can be efficiently compressed.

Unfilter

The possible filter types specified by the PNG specification are:

   Type    Name
   
   0       None
   1       Sub
   2       Up
   3       Average
   4       Paeth

There is some subtle points to pay attention to in the specification, such as the fact that these filters are applied per byte, and not per pixel value. For 8-bit colour depth this is unimportant, but at 16-bits, this means the first byte of a pixel (the MSB, or most significant byte) will be computed separately from the second byte (the LSB, or least significant byte). I won’t go too deep into all the details here, but you can read the details of the specification here.

I’ll briefly explain the basic idea behind each filter:

0: None
- No filter is applied and each byte value is just the raw pixel value.
1: Sub
- Each byte has the preceding byte value subtracted from it
2: Up
- Each byte has the value of the byte above it subtracted from it.
3: Average:
- Each byte has the (floor) of the average of the bytes above and to the left of it subtracted from it.
4: Paeth:
- The three neighbouring pixels (left, above and upper left) are used to calculate a value that is subtracted from the pixel. It’s a little more involved than the other three.

So when decoding the filtered data, we need to reverse the above operations to regain the pixel values.

Now that we know that, let’s look at our first byte value:

print(uncompressed_data[0])

So we are dealing with filter type 1 here. Let’s decode the first row:

var filter_type = uncompressed_data[0]
var scanline = uncompressed_data[1:128*3+1]

# Decoded image data
var result = List[UInt8](capacity=128*3)
# take the first pixels as 0
var left: UInt8 = 0
var pixel_size: Int = 3
var offset: Int = 1

for i in range(len(scanline)):
    if i >= pixel_size:
        left = result[i-pixel_size] 

    # The specification specifies that the result is modulo 256
    # Silimar to the C implementation, we can just add the left pixel to the current pixel,
    # and the result will be modulo 256 due to overflow
    result.append((uncompressed_data[i + offset] + left))

And let’s confirm that the row we decoded is the same as PIL would do:

for i in range(128):
    for j in range(3):
        assert_true(result[i*3+j] == py_array[0][i][j].__int__(), "Pixel values do not match")

Now that we have the general idea of things, let’s write this more generally, and do the other filters as well.

For an idea of how filters are chosen, read this stackoverflow post and the resources it points to: How do PNG encoders pick which filter to use?

I’ve done these as functions that take 16-bit signed integers. This is important mostly for the case of the Paeth filter, where the standard states:

The calculations within the PaethPredictor function must be performed exactly, without overflow. Arithmetic modulo 256 is to be used only for the final step of subtracting the function result from the target byte value.

So basically we need to keep a higher level of precision and then cast back to bytes at the end.

I based the implementation on the iPXE implementation of a png decoder written in C.

from math import abs

fn undo_trivial(current: Int16, left: Int16 = 0, above: Int16 = 0, above_left: Int16 = 0) -> Int16:
    return current

fn undo_sub(current: Int16, left: Int16 = 0, above: Int16 = 0, above_left: Int16 = 0) -> Int16:
    return current + left

fn undo_up(current: Int16, left: Int16 = 0, above: Int16 = 0, above_left: Int16 = 0) -> Int16:
    return current + above

fn undo_average(current: Int16, left: Int16 = 0, above: Int16 = 0, above_left: Int16 = 0) -> Int16:
    return current + ((above + left) >> 1) # Bitshift is equivalent to division by 2

fn undo_paeth(current: Int16, left: Int16 = 0, above: Int16 = 0, above_left: Int16 = 0) -> Int16:

    var peath: Int16 = left + above - above_left
    var peath_a: Int16 = abs(peath - left)
    var peath_b: Int16 = abs(peath - above)
    var peath_c: Int16 = abs(peath - above_left)
    if ( peath_a <= peath_b ) and ( peath_a <= peath_c ):
        return (current + left)
    elif ( peath_b <= peath_c ): 
        return (current + above)
    else:
        return (current + above_left)

fn undo_filter(filter_type: UInt8, current: UInt8, left: UInt8 = 0, above: UInt8 = 0, above_left: UInt8 = 0) raises -> UInt8:

    var current_int = current.cast[DType.int16]()
    var left_int = left.cast[DType.int16]()
    var above_int = above.cast[DType.int16]()
    var above_left_int = above_left.cast[DType.int16]()
    var result_int: Int16 = 0

    if filter_type == 0:
        result_int= undo_trivial(current_int, left_int, above_int, above_left_int)
    elif filter_type == 1:
        result_int = undo_sub(current_int, left_int, above_int, above_left_int)
    elif filter_type == 2:
        result_int = undo_up(current_int, left_int, above_int, above_left_int)
    elif filter_type == 3:
        result_int = undo_average(current_int, left_int, above_int, above_left_int)
    elif filter_type == 4:
        result_int = undo_paeth(current_int, left_int, above_int, above_left_int)
    else:
        raise Error("Unknown filter type")
    return result_int.cast[DType.uint8]()

For the undo_filter function, I was trying to add the separate filters to some kind of Tuple or List so I could just index them (hence the uniform signatures), but wasn’t able to figure out how to do this in Mojo yet.

So let’s apply these to the whole image and confirm that we have the same results as we would get from Python:


# Decoded image data
# take the first pixels as 0
var pixel_size: Int = 3

# Initialize the previous scanline to 0
var previous_result = List[UInt8](0*128)

for line in range(128):
    var offset =  1 + 1*line + line * 128 * 3
    var left: UInt8 = 0
    var above_left: UInt8 = 0

    #var left: UInt8 = 0
    var result = List[UInt8](capacity=128*3)
    var scanline = uncompressed_data[offset:offset+128*3]

    var filter_type = uncompressed_data[offset - 1]

    for i in range(len(scanline)):
        if i >= pixel_size:
            left = result[i-pixel_size] 
            above_left = previous_result[i-pixel_size] 

        result.append(undo_filter(filter_type, uncompressed_data[i + offset], left, previous_result[i], above_left))


    previous_result = result
    for i in range(128):
        for j in range(3):
            assert_true(result[i*3+j] == py_array[line][i][j].__int__(), "Pixel values do not match")

And that’s it. If the above runs it means we’ve sucessfully parsed a PNG file, and at least get the same data out as you would by using Pillow.

Creating a tensor

Now ideally we want the above into a Tensor.

Lets write a function that will parse the image data and return a Tensor for us.

from tensor import Tensor, TensorSpec, TensorShape
from utils.index import Index
from random import rand

var height = 128
var width = 128
var channels = 3

# Declare the image tensor.
var spec = TensorSpec(DType.uint8, height, width, channels)
var tensor_image = Tensor[DType.uint8](spec)

# Decoded image data
# take the first pixels as 0
var pixel_size: Int = 3

# Initialize the previous scanline to 0
var previous_result = List[UInt8](0*128)

for line in range(128):
    var offset =  1 + 1*line + line * 128 * 3
    var left: UInt8 = 0
    var above_left: UInt8 = 0

    #var left: UInt8 = 0
    var result = List[UInt8](capacity=128*3)
    var scanline = uncompressed_data[offset:offset+128*3]

    var filter_type = uncompressed_data[offset - 1]

    for i in range(len(scanline)):
        if i >= pixel_size:
            left = result[i-pixel_size] 
            above_left = previous_result[i-pixel_size] 

        result.append(undo_filter(filter_type, uncompressed_data[i + offset], left, previous_result[i], above_left))


    previous_result = result
    for i in range(128):
        for j in range(3):
            tensor_image[Index(line, i, j)] = result[i*3+j]

I’m not entirely sure why I need to use Index while setting items, but when getting I can just provide indices:

print(tensor_image[0,1,2])
print(py_array[0][1][2])

62
62

And there we have it. I will put it all together soon but let’s finish parsing the file quickly.

Final chunks

There are a few more chunks at this point: text chunks which hold some comments, and an end chunk, which denotes the end of the file:

var text_chunk_1 = parse_next_chunk(file_contents, read_head)
print(text_chunk_1.type)
read_head = text_chunk_1.end
print(bytes_to_string(text_chunk_1.data))

tEXt
comment

var text_chunk_2 = parse_next_chunk(file_contents, read_head)
print(text_chunk_2.type)
read_head = text_chunk_2.end
print(bytes_to_string(text_chunk_2.data))

tEXt
date:create

var text_chunk_3 = parse_next_chunk(file_contents, read_head)
print(text_chunk_3.type)
read_head = text_chunk_3.end
print(bytes_to_string(text_chunk_3.data))

tEXt
date:modify

The text chunks above actually have more info, but seem to be UTF-8 encoded, and Mojo just seems to handle ASCII?

var end_chunk = parse_next_chunk(file_contents, read_head)
print(end_chunk.type)
read_head = end_chunk.end

IEND

Putting it all together.

Let’s package the logic above up a bit more nicely. I’m thinking something that resembles PIL.

Well start with a struct called PNGImage


fn bytes_to_hex_string(list: List[Int8]) -> String:
    var word = String("")
    for letter in list:
        word += hex(int(letter[].cast[DType.uint8]()))

    return word


fn determine_file_type(data: List[Int8]) -> String:

    # Is there a better way? Probably
    if bytes_to_hex_string(data[0:8]) == String("0x890x500x4e0x470xd0xa0x1a0xa"):
        return "PNG"
    else:
        return "Unknown"



struct PNGImage:
    var image_path: Path
    var raw_data: List[Int8]
    var width: Int
    var height: Int
    var channels: Int
    var bit_depth: Int
    var color_type: Int
    var compression_method: UInt8
    var filter_method: UInt8
    var interlaced: UInt8
    var data: List[UInt8]
    var data_type: DType

    fn __init__(inout self, file_name: Path) raises:

        self.image_path = file_name
        assert_true(self.image_path.exists(), "File does not exist")

        with open(self.image_path , "r") as image_file:
            self.raw_data = image_file.read_bytes()

        assert_true(determine_file_type(self.raw_data) == "PNG", "File is not a PNG. Only PNGs are supported")

        var read_head = 8

        var header_chunk = parse_next_chunk(self.raw_data, read_head)
        read_head = header_chunk.end

        self.width = int(bytes_to_uint32_be(header_chunk.data[0:4])[0])
        self.height = int(bytes_to_uint32_be(header_chunk.data[4:8])[0])
        self.bit_depth = int(header_chunk.data[8].cast[DType.uint32]())
        self.color_type = int(header_chunk.data[9])
        self.compression_method = header_chunk.data[10].cast[DType.uint8]()
        self.filter_method = header_chunk.data[11].cast[DType.uint8]()
        self.interlaced = header_chunk.data[12].cast[DType.uint8]()

        # There must be a better way to do this
        var color_type_dict = Dict[Int, Int]()
        color_type_dict[0] = 1
        color_type_dict[2] = 3
        color_type_dict[3] = 1
        color_type_dict[4] = 2
        color_type_dict[6] = 4

        self.channels = color_type_dict[self.color_type]

        if self.bit_depth == 8:
            self.data_type = DType.uint8
        elif self.bit_depth == 16:
            self.data_type = DType.uint16
        else:
            raise Error("Unknown bit depth")

        # Check color_type and bit_depth
        assert_true(self.color_type == 2, "Only RGB images are supported")
        assert_true(self.bit_depth == 8, "Only 8-bit images are supported")
        
        # Check if the image is interlaced
        assert_true(self.interlaced == 0, "Interlaced images are not supported")
        # Chack compression method
        assert_true(self.compression_method == 0, "Compression method not supported")


        # Scan over chunks until end found
        var ended = False
        var data_found = False
        var uncompressd_data = List[Int8]() 
        while read_head < len(self.raw_data) and not ended:
            var chunk = parse_next_chunk(self.raw_data, read_head)
            read_head = chunk.end

            if chunk.type == "IDAT":
                uncompressd_data.extend(chunk.data)
                data_found = True
            elif chunk.type == "IEND":
                ended = True

        assert_true(ended, "IEND chunk not found")
        assert_true(data_found, "IDAT chunk not found")
        self.data = uncompress(uncompressd_data)
            


    # In case the filename is passed as a string
    fn __init__(inout self, file_name: String) raises:
        self.__init__(Path(file_name))


    fn to_tensor(self) raises -> Tensor[DType.uint8]:
        var spec = TensorSpec(DType.uint8, self.height, self.width, self.channels)
        var tensor_image = Tensor[DType.uint8](spec)

        var pixel_size: Int = self.channels * (self.bit_depth // 8)

        # Initialize the previous scanline to 0
        var previous_result = List[UInt8](0*self.width)

        for line in range(self.height):
            var offset =  1 + 1*line + line * self.width * pixel_size
            var left: UInt8 = 0
            var above_left: UInt8 = 0

            var result = List[UInt8](capacity=self.width*pixel_size)
            var scanline = self.data[offset:offset+self.width*pixel_size]

            var filter_type = self.data[offset - 1]

            for i in range(len(scanline)):
                if i >= pixel_size:
                    left = result[i-pixel_size] 
                    above_left = previous_result[i-pixel_size] 

                result.append(undo_filter(filter_type, self.data[i + offset], left, previous_result[i], above_left))

            previous_result = result
            for i in range(self.width):
                for j in range(self.channels):
                    tensor_image[Index(line, i, j)] = result[i*self.channels+j]


        return tensor_image

Well, it’s not the prettiest, but let’s see if it works:

var hopper = PNGImage(test_image)
hopper_tensor = hopper.to_tensor()

for i in range(hopper.height):
    for j in range(hopper.width):
        for c in range(hopper.channels):
            assert_true(hopper_tensor[i, j, c] == py_array[i][j][c].__int__(), "Pixel values do not match")

If the above runs, then it means we read the image correctly!

Let’s try on a PNG image from the CIFAR-10 dataset:

cifar_image = Path('114_automobile.png')
var cifar = PNGImage(cifar_image)
cifar_tensor = cifar.to_tensor()

py_cifar = np.array(Image.open('114_automobile.png'))

for i in range(cifar.height):
    for j in range(cifar.width):
        for c in range(cifar.channels):
            assert_true(cifar_tensor[i, j, c] == py_cifar[i][j][c].__int__(), "Pixel values do not match")

This also works! Now we should be able to read the CIFAR-10 dataset!

I have a few questions about my implementation above, i.e:

How to handle 16-bit images? Do we need a separate function, i.e. to_tensor_16?
Is a struct the correct way to do this?

This is one of those points where I don’t feel I know what the idiomatic way to do this in Mojo is yet. The 🪄Mojical🪄 way, you might say 😉. Mojo is so young that I’m not sure an idiomatic way has emerged yet.

Conclusions

Reading PNGs was quite a fun topic for a blog post. It made me really get my hands dirty with some of the more low level concepts in Mojo, something I felt I didn’t fully grasp before.

I’ll admit, this ended up being a bit more work than I expected. To paraphrase JFK:

jfk_meme

It’s impressive how far Mojo has come in just a few months: when I was trying to write a bit of Mojo in September of last year it felt hard to do anything practical, while now the language seems to be quite usable.

There are still a few things I need to get used to. One thing is I always feel like I “need” to write fn functions, and not def functions. This is good practice when writing libraries and such, but it makes me wonder: when is writing def style functions appropriate, as fn will always be safer and more performant?

Addendum

I refactored the code from this blog post a bit and wrote it up into a library I am calling Mimage. The goal would be to be able to read and write common image formats in pure Mojo without having to resort to calling Python or C. Currently Mimage still requires some C libraries for the uncompress step, but I am hoping that those will be available in pure Mojo soon.

The next steps will likely be adding support to Mimage for 16-bit PNGs and JPEGs. The long term goal would be to be able to read and write all the same image formats as Python’s Pillow, but that will likely take a long time to reach. As I am on the ML side of things, I’ll try and focus on the formats and functionality needed for ML purposes, like being able to read all the images in the Imagenet dataset.

Citation

BibTeX citation:

@online{schenck2024,
  author = {Schenck, Ferdinand},
  title = {Parsing {PNG} Images in {Mojo}},
  date = {2024-05-04},
  url = {https://fnands.com/blog/2024/mojo-png-parsing/},
  langid = {en}
}

For attribution, please cite this work as:

Schenck, Ferdinand. 2024. “Parsing PNG Images in Mojo.” May 4. https://fnands.com/blog/2024/mojo-png-parsing/.

Mojo’s standard library goes open source

Ferdinand Schenck — Sun, 14 Apr 2024 00:00:00 GMT

A lot has changed in Mojo land since I last had a look. The last time I wrote one of these Mojo was at version 0.6, and now it’s at 24.2!
This is just due to a change in version numbering, so they have basically still kept the pace of about one new version per month, so the releases (ignoring bugfix releases) goes 0.6 -> 0.7 -> 24.1 -> 24.2.

Why 24.1/24.2? Modular have moved to a YY.MAJOR.MINOR numbering scheme, so 24.1 and 24.2 are the first and second releases (under this naming scheme) in 2024, respectively. The reason this was done is to keep the Mojo versioning in sync with Max, Modular’s other product. A minor worry I have about this is that unlike semantic versioning where you get a nice 1.0 release, it’s a bit hard to know now what a stable release will look like.

Another big piece of news is that the Mojo standard library is now open source, under the Apache 2.0 license (with LLVM exceptions). So far only the SDK has been available, but with the stdlib this is one step closer to the entire language being open source. Some people on HackerNews have grumbled about the language not being open source, but I do totally get why the developers want to reach some level of maturity before opening their language up to the peanut gallery.

In any case, I want to have a bit of a look at what has changed.

Release 0.7

The one object I’ve severely missed in Mojo has been a dictionary type, and in 0.7 they finally added it. It’s one of the most useful data structures in Python, and they have kept the implementation similar to the Python one (so a hash map under the hood).

A slight difference is that for Mojo dicts the types must be statically specified, unlike Python where you can just kinda do whatever you want.

So to define a dictionary:

from collections import Dict
var my_dict = Dict[String, String]()
my_dict["foo"] = "oof"
my_dict["bar"] = "rab"
print(len(my_dict))      
print(my_dict["foo"])      
print(my_dict.pop("bar"))  
print(len(my_dict))

2
oof
rab
1

So pretty straightforward stuff. Nice!

Another profiling feature that I’ve never seen anywhere before is getting not just the number of cores on your system, but the types as well.

from sys import info

print(info.num_logical_cores())
print(info.num_physical_cores())
print(info.num_performance_cores())

20
14
6

So you can see my machine has 20 total “cores”, which I believe breaks down as 6 performance cores and 8 efficiency cores, for 6 + 8 = 14 physical cores, and the 6 performance cores each having two threads, so 2 * 6 + 8 = 20 threads. I’ve never really had to consider what kind of cores my program runs on, but this will probably become an increasingly important thing to consider going forward.

Another interesting thing that caught my eye is the new Reference type, which is taking a page out of the Rust playbook. Checking the roadmap, they do state that a borrow checker is on the way. This does make a lot of sense: if you want to be able to give people access to low-level memory management, do it the Rust way. Right now you can mess around with raw pointers and the compiler will just let you, but it seems this will change in the future.

I mentioned in an earlier post that Mojo felt like writing C++ with Python syntax, but maybe it will be more like writing Rust with Python syntax soon.

Talking of comparing Mojo to Rust, ThePrimeagen made a video discussing the topic. Actually it’s discussing a blog post by Modular that is discussing an earlier video from him discussing a Mojo community blog post. It’s a whole thing. In any case, it’s an interesting video going into a lot of the rationale behind how a lot of the design decisions about Mojo are being made.

Release 24.1

Another release, another useful data structure. This time it’s Mojo’s implementation of a set:

from collections import Set

var my_set = Set[String]()
my_set.add("foo")
my_set.add("bar")
print(len(my_set))

my_set.add("foo")
print(len(my_set))

if "foo" in my_set:
    print("foo is in the set")

2
2
foo is in the set

So Set is a nice little addition, and works as you would expect. The a in b syntax from Python is also new, so every day we get a bit closer to Python.

A feature I haven’t appreciated before is unbound values and the alias declaration: so you can create an alias of an object, which is basically a partial initialization that can then be fully initialized later.

@value
struct StructWithDefault[a: Int, b: Int, c: Int = 8, d: Int = 9]: pass
alias end_unbound = StructWithDefault[3, 4, 5, _]

var my_instance_5 = end_unbound[5]()
var my_instance_10 = end_unbound[10]()

This reminds me a bit of partial in Python (however, this does happen at compile time), and might make sense if you are initializing multiple versions of large objects. On that note, the compile-time metaprogramming side of Mojo is something I still need to explore and something I haven’t fully grasped yet.

A new feature is that you can unbind any number of parameters with *_

alias most_unbound = StructWithDefault[3, *_]

var my_instance_567 = most_unbound[5, 6, 7]()
var my_instance_8910 = most_unbound[8, 9, 10]()

An issue I ran into earlier was the lack of support for iteration in DynamicVector, which has now been rectified!

Additionally (in v24.2), the DynamicVector been replaced with a List, which has a distinctly Pythonic feel to it:

var names =  List[String]("Alice", "Bob", "Charlie")
for x in names:
    x[] = str("Hello, ") + x[]
for x in names:
    print(x[])

Hello, Alice
Hello, Bob
Hello, Charlie

Even negative indexing is now supported:

print(names[-1])

Hello, Charlie

If you are wondering what the [] syntax in print(x[]) is for is that the iterator over the list returns references, and [] is the dereference operator in Mojo.

Additionally, file handling just became easier with the introduction of pathlib.Path and os.listdir, which should be familiar to anyone who uses Python.

Release 24.2

An interesting change is the removal of the let declaration, i.e. the way you declare immutable variables in Mojo. I thought this was a bit strange at first, but there has been a lot of thought put into the decision, see this post by Chris Lattner.

Partly it’s due to simplicity, as Python has no concept of immutability, but the one interesting argument that stuck out to me was: > The immutability only applies to the local value i.e. only the Pointer (or Reference) is immutable if you use let, not the value being pointed to, which was causing some confusion. It sounds like the concept is being re-evaluated and will likely make a comeback in some shape later.

print in Mojo is now very close to the Python version, so you can specify the separator and end:

print("Hello", "Mojo", sep=", ", end="!!!\n") 
print("Hello", "Mojo", sep="~~~", end="\n~~~")

Hello, Mojo!!!
Hello~~~Mojo
~~~

Other than that, there are a few changes that bring the language closer to Python, such as adding variadic keyword arguments (as **kwargs), but the biggest change in 24.2 is really the release of the stdlib.

Conclusions

Mojo is finally starting to look usable. Earlier this year there was still too much missing to do much more than a little sandboxed demo in Mojo, and it looks like this is starting to change. Someone is even starting building an ML framework (called Basalt) in Mojo, and I’m excited to see where that goes. On that note, I’m very curious to see when Mojo will finally support GPUs, as that will really be the point where it will start living up to it’s promise of being the language of machine learning.

So far I’ve basically just been writing these posts as an exercise to make myself really parse the Mojo changelog, but I’m tempted to actually start running some real world tests.

Citation

BibTeX citation:

@online{schenck2024,
  author = {Schenck, Ferdinand},
  title = {Mojo’s Standard Library Goes Open Source},
  date = {2024-04-14},
  url = {https://fnands.com/blog/2024/mojo-stdlib-open-source/},
  langid = {en}
}

For attribution, please cite this work as:

Schenck, Ferdinand. 2024. “Mojo’s Standard Library Goes Open Source.” April 14. https://fnands.com/blog/2024/mojo-stdlib-open-source/.

An intuitive introduction to pansharpening

Ferdinand Schenck — Tue, 02 Jan 2024 00:00:00 GMT

Many earth observation satellites are equipped with a higher resolution panchromatic sensor, and a lower resolution multispectral sensor. One example of this is the Pleiades satellite constellation by Airbus, for which the panchromatic sensor has a 70 cm ground sampling distance (GSD), and the multispectral sensor has a 2.8 m GSD, i.e. four times lower than the panchromatic band.

The panchromatic band is sensitive to a wide spectrum of light, usually overlapping with several of the other spectral bands, while the multispectral bands focus on narrower parts of the spectrum.

The image below shows the spectral resolution of the panchromatic as well the the blue, green, red and near infra-red (NIR) bands of the Pleiades sensor (taken from here):

Pleiades spectral bands.

One of the reasons for having different resolutions for different spectral bands is that a sensor needs to receive a certain amount of light before it can form a reliable image. If you increase the spatial resolution or decrease the spectral resolution, you reduce the amount of light reaching each pixel, meaning satellite designers often have to choose between high spectral resolution or high spatial resolution. For a more in-depth explanation see this answer on GIS StackExchange.

What this results in is that you often have one image with high spatial information, but low spectral information, and one with low spatial information but high spectral information.

But what if you want both?

This is where pansharpening comes in.

By combining the spatial information from the panchromatic image and the spectral information from the multispectral image, we can create an image with the best of both worlds.

An intuitive example

Remote sensing images are sometimes a bit unintuitive to deal with, so I will use a “normal” image as an example, in this case a logo I created with Ideogram:

import PIL.Image as Image
import numpy as np
original_image = Image.open("treed_brain_512.jpeg")
original_width, original_height = original_image.size
original_array = np.array(original_image)
original_image

To simulate the situation in satellite images we will create a “panchromatic” image by making a grayscale image:

toy_panchromatic_image = original_image.convert("L")
toy_panchromatic_array = np.array(toy_panchromatic_image)
toy_panchromatic_image

and to simulate a lower resolution multispectral image we will downsample the image to get a lower resolution version of the colour image.
In this case our original image had a resolution of 512 x 512 pixels, and we downsample each side by a factor of four to get a 128 x 128 pixel image, meaning for every pixel in the multispectral image we will have sixteen (4 x 4) pixels in the panchromatic image.

toy_multispectral_image = original_image.resize((original_width//4, original_width//4))
toy_multispectral_array = np.array(toy_multispectral_image)
toy_multispectral_image

To see the difference we can naively upsample the multispectral image and show the images side-by-side:

toy_upscaled_multispectral_image = toy_multispectral_image.resize((original_width, original_height))
toy_upscaled_multispectral_array = np.array(toy_upscaled_multispectral_image)

side_by_side = np.concatenate((toy_panchromatic_array[:,:,None].repeat(3, axis=2), toy_upscaled_multispectral_array), axis=1)
Image.fromarray(side_by_side)

As you can see we now have a high resolution “panchromatic” image and a low resolution “multispectral” image.

So how can we add back sharpness to the image on the right?

Brovey pansharpening

One of the simplest (and most widely used) methods for pansharpening is called the Brovey transformation. For reference, it is the default pansharpening operation in GDAL.

To do a Brovey transform you need to upsample the multispectral image so your images have the same resolution.

We did this above, but to confirm let us check the resolutions of our images:

print(f"Panchromatic shape: {toy_panchromatic_array.shape}")
print(f"Multispectral shape: {toy_upscaled_multispectral_array.shape}")

Panchromatic shape: (512, 512)
Multispectral shape: (512, 512, 3)

Next, for every pixel we create a “pseudo-pan” value by combining the values of the spectral values at each pixel of the (upsampled) multispectral image.

So:

However the above assumes that the spectral bands contributed equally to the panchromatic image. As you can see in the illustration of the Pleiades spectral resolution, the green and red bands share a larger overlap with the pancharomatic band than the blue or NIR bands, and should therefore contribute more. In fact, the four spectral bands of the multispectral sensor do not cover the entire spectral range covered by the panchromatic sensor.

Therefore we need a set of weights that tells us how each band contributes to the panchromatic image. In case of satellite sensors, these can be found via calibration and are usually available from the manufacturer.

In the case of our toy example, we know the weights used to create our “panchromatic” image. The weights are 0.299, 0.587 and 0.114 for red green and blue respectively, and a description of why these weights are used can be found here.

Our equation is now:

and we are now technically doing a weighted Brovey transform.

By doing this for every pixel we can get an approximation of what the expected panchromatic image should look like:

channel_dependant_luminance_perception = np.array([0.299, 0.587, 0.114])
psuedo_pan_array = np.true_divide((toy_upscaled_multispectral_array*channel_dependant_luminance_perception).sum(axis=2), channel_dependant_luminance_perception.sum())
psuedo_pan_image = Image.fromarray(psuedo_pan_array.astype(np.uint8))
psuedo_pan_image

By using the ratio between real panchromatic image and the psuedo-pan image we can get a map that tells us which pixels to lighten and which to darken:

ratio = toy_panchromatic_array/psuedo_pan_array
Image.fromarray((ratio*255/ratio.max()).astype(np.uint8))

Applying this ratio to each of the bands of our upscaled multispectral image, we get our pansharpened image:

new_red = toy_upscaled_multispectral_array[:, :, 0] * ratio
new_green = toy_upscaled_multispectral_array[:, :, 1] * ratio
new_blue = toy_upscaled_multispectral_array[:, :, 2] * ratio
pansharpened_array = np.stack([new_red, new_green, new_blue], axis=2)
pansharpened_image = Image.fromarray(pansharpened_array.astype(np.uint8))
pansharpened_image

This looks pretty good! The image is in colour, and has crisp sharp edges.

Putting these side-by-side, with the original image on the left and the pansharpened image on the right, we can see a few small differences:

side_by_side = np.concatenate((original_array, pansharpened_array), axis=1)
Image.fromarray(side_by_side.astype(np.uint8))

The most prominent artefacting occurs in the white lines in the “satellites”, i.e. in areas where there is a sudden change from one colour to another.

In order to more clearly see the differences between the two cases, let us create a difference map:

difference_array = np.abs((pansharpened_array - original_array))
difference_array = difference_array/difference_array.max()*255
Image.fromarray((difference_array).astype(np.uint8), mode="RGB")

As might be expected, the largest differences occur where colours change quickly (i.e. high gradient). This is because the lower resolution sensor will naturally average the colour information in a pixel, so the sharpness that we add to it doesn’t perfectly correspond to the averaged out colour information.

Conclusions

The above is an intuitive introduction to pansharpening, and the Brovey pansharpening algorithm in particular.

Applying the concept to remote sensing images is not much more complex, one just needs to account for more spectral bands and different ratios between bands.

There are more complex pansharpening algorithms out there, but Brovey is one of the most widely used, and gives good enough results for most applications.

Citation

BibTeX citation:

@online{schenck2024,
  author = {Schenck, Ferdinand},
  title = {An Intuitive Introduction to Pansharpening},
  date = {2024-01-02},
  url = {https://fnands.com/blog/2024/pansharpening-intro-python/},
  langid = {en}
}

For attribution, please cite this work as:

Schenck, Ferdinand. 2024. “An Intuitive Introduction to Pansharpening.” January 2. https://fnands.com/blog/2024/pansharpening-intro-python/.

Mojo 0.6.0, now with traits and better Python like string wrangling.

Ferdinand Schenck — Sun, 10 Dec 2023 00:00:00 GMT

Another month, another Mojo release.

I am busy doing the 2023 edition of the Advent of Code (AoC) in Mojo, and had a few complaints 😅.
If you not familiar with the AoC, it’s basically a coding advent calendar that gives you a new coding challenge every day for the first 25 days of December.

In a bit of foreshadowing, I used an AoC 2022 puzzle in my first post on Mojo, which was using Mojo 0.2.1, and it is encouraging to see how far the language has come.

The AoC puzzles are often pretty heavy in string wrangling, a task that Python is pretty strong in, and that Mojo is still somewhat lacking in. One of the features that I found was lacking in Mojo 0.5.0 was the ability to easily split a string as one does in Python.

In the case of the first day, I found myself needing to split a string by newlines, something which you can do trivially in Python by calling my_string.split('\n'). In Mojo 0.5.0 this did not exist and I had to write a struct to implement this functionality. I ended up generalizing it a bit and putting it in a library as it was super useful for the following days as well.

And then on the fourth of December Mojo 0.6.0 was released, which now includes the ability to call .split() on a string, as well as a bunch of useful Python methods (rfind(), len(), str(), int()). These will definitely help going forward with the AoC challenges.

I’ll write a rundown of my experience with the AoC in Mojo when I complete all the puzzles, so now on the the spotlighted feature from 0.6.0: traits

Traits in Mojo

Traits are a fairly common concept in programming languages, and allow you to add required functionality to a struct if it conforms to this trait.

As an example, take the len() function that we know and love from Python, and that is also now a part of Mojo. The trait associated with len() in Mojo is Sized, meaning that any struct conforming to the Sized trait is required to have a __len__() method that returns an integer size. When the function len() is applied to a struct that conforms to Sized, the __len__() function is called.

An example of a struct that conforms to the Sized trait is the builtin String:

example_string = "This is a String and it conforms to the Sized trait."

print(len(example_string))

Additionally, we can then write our own struct that conforms to Sized, and as long as it has a method named __len__() it will conform to the Sized trait (the compiler will let you know if it doesn’t):

@value
struct MySizedStruct(Sized):
    var size: Int

    fn __len__(self) -> Int:
        return self.size

If we now call len() on an instance of this struct it will return the size value:

sized_struct = MySizedStruct(10)
print(len(sized_struct))

As a side note, I used the @value decorator above which hides a bit of boilerplate code for us. The above initialization is equivalent to:

struct MySizedStruct(Sized):
    var size: Int

    fn __len__(self) -> Int:
        return self.size

    fn __init__(inout self, size: Int):
        self.size = size

    fn __copyinit__(inout self, existing: Self):
        self.size = existing.size

    fn __moveinit__(inout self, owned existing: Self):
        self.size = existing.size

So @value is a pretty useful way to save us a few lines of boilerplate code.

I’m still getting used to decorators in Mojo (maybe a good idea to do a post on them in the future).

One question I had about traits is how difficult it is to chain them? E.g. what if I have a struct that I want to conform to both Sized and Stringable, which allows the function str() to apply to the struct, and makes it printable?

It turns out this is easy; just pass them during as a comma

@value
struct MySizedAndStingableStruct(Sized, Stringable):
    var size: Int

    fn __len__(self) -> Int:
        return self.size

    fn __str__(self) -> String:
        return str(self.size)

sized_and_stringable = MySizedAndStingableStruct(11)
print(sized_and_stringable)

So it is very simple to add multiple traits.

To create our own trait, we only need to define it with a method that conforming structs need to inherit:

trait Jazzable:
    fn jazz(self): ...

Here the ... indicates that nothing is specified yet (needs to be done per struct). It is not possible yet to define a default method, but is apparently coming in the future.

Let’s create a struct that conforms to Jazzable:

@value
struct JazzX(Jazzable):
    var jazz_level: Int
    
    fn jazz(self):
        for i in range(self.jazz_level):
            print("Jazzing at level", i + 1)

ten_jazz = JazzX(10)
ten_jazz.jazz()

Jazzing at level 1
Jazzing at level 2
Jazzing at level 3
Jazzing at level 4
Jazzing at level 5
Jazzing at level 6
Jazzing at level 7
Jazzing at level 8
Jazzing at level 9
Jazzing at level 10

We can also define a function that calls a specific method. An example of this is the len() function that calls __len__(), we can create our own function that will call jazz():

fn make_it_jazz[T: Jazzable](jazz_struct: T):
    jazz_struct.jazz()


make_it_jazz(ten_jazz)

Jazzing at level 1
Jazzing at level 2
Jazzing at level 3
Jazzing at level 4
Jazzing at level 5
Jazzing at level 6
Jazzing at level 7
Jazzing at level 8
Jazzing at level 9
Jazzing at level 10

Additionally, traits can inherit from other traits, and keep the functionality of the parent trait:

trait SuperJazzable(Jazzable):
    fn super_jazz(self): ...


@value 
struct SuperJazz(SuperJazzable):
    var jazz_level: Int

    fn jazz(self):
        for i in range(self.jazz_level):
            print("Jazzing at level", i + 1)

    fn super_jazz(self):
        for i in range(self.jazz_level):
            print("Super Jazzing at level", (i + 1)*10)

This new struct will have all the methods of Jazzable, so make_it_jazz() will work:

super_jazz_5 = SuperJazz(5)
make_it_jazz(super_jazz_5)

Jazzing at level 1
Jazzing at level 2
Jazzing at level 3
Jazzing at level 4
Jazzing at level 5

And we can define additional functions that will activate the new methods as well:

fn make_it_super_jazz[T: SuperJazzable](superjazz_struct: T):
    superjazz_struct.super_jazz()

make_it_super_jazz(super_jazz_5)

Super Jazzing at level 10
Super Jazzing at level 20
Super Jazzing at level 30
Super Jazzing at level 40
Super Jazzing at level 50

Conclusions

Traits provide a convenient way of adding functionality to structs, and as you can see they are pretty simple to use.

I’ve never used traits in any other language before, but it does work similarly to generic classes, and feels really familiar, except for the fact that you can’t have default behaviour (yet).

From what I’ve seen from Mojo so far, writing structs seems to be a pretty core part of how Mojo is supposed to be used, so I guess I better get used to it.

Citation

BibTeX citation:

@online{schenck2023,
  author = {Schenck, Ferdinand},
  title = {Mojo 0.6.0, Now with Traits and Better {Python} Like String
    Wrangling.},
  date = {2023-12-10},
  url = {https://fnands.com/blog/2023/mojo-traits/},
  langid = {en}
}

For attribution, please cite this work as:

Schenck, Ferdinand. 2023. “Mojo 0.6.0, Now with Traits and Better Python Like String Wrangling.” December 10. https://fnands.com/blog/2023/mojo-traits/.

Mojo 0.5.0 and SIMD

Ferdinand Schenck — Fri, 17 Nov 2023 00:00:00 GMT

Another month, another Mojo release!
It feel like every time I run into a missing feature in Mojo it gets added in the next version:
In my first post I complained about a lack of file handling, which was then added soon after in version 0.4.0.
For version 0.4.0 I ran into the issue that you can’t print Tensors, which has now been added in the 0.5.0 release. So this means Mojo has now unlocked everyone’s favourite method of debugging: printing to stdout.
In addition to that, Tensors can now also be written to and read from files with fromfile() and tofile().

There have also been a couple of updates to the SIMD type, which lead me to ask: How does the SIMD type work in Mojo?

For a bit of background, you might have noticed that CPUs clock speeds haven’t really increased by much in the last decade or so, but computers have definitely gotten faster. One of the factors that have increased processing speed has been a focus on vectorization through SIMD, which stands for Single Instruction, Multiple Data, i.e. applying the same operation to multiple pieces of data.
Modern CPUs come with SIMD registers that allow the CPU to apply the same operation over all the data in that register, resulting in large speedups, especially in cases where you are applying the same operation to multiple pieces of data, e.g. in image processing where you might apply the same operation to the millions of pixels in an image.

The SIMD type in Mojo

One of the main goals of the Mojo language is to leverage the ability of modern hardware, both CPUs and GPUs, to execute SIMD operations.

There is no native SIMD support in Python, however Numpy does make this possible.

Note: SIMD is not the same as concurrency, where you have several different threads running different instructions. SIMD is doing the same operation on different data.

Generally, SIMD objects are initialized as SIMD[datatype, size](values), so to create a SIMD object consisting of four 8-bit unsigned integers we would do:

simd_4x_uint8 = SIMD[DType.uint8, 4](1, 2, 4, 8)
print(simd_4x_uint8)

[1, 2, 4, 8]

And actually, SIMD is so central to Mojo that the builtin Float32 type is actually just an alias for SIMD[DType.float32, 1]:

let x: Float32 = 10.
print(x.type)

let y = SIMD[DType.float32, 1](10.)
print(y.type)

float32
float32

Modern CPUs have SIMD registers, so lets use the sys.info package in Mojo to see what the register width on my computer is:

from sys.info import simdbitwidth, simdwidthof
print(simdbitwidth())

This means we can pack 256 bits of data into this register and efficiently vectorize an operation over it. Some CPUs support AVX-512, with as the name suggests 512 bit SIMD registers. Most modern CPUs will apply the same operation to all values in their register in one step, allowing for significant speedup for functions that can exploit SIMD vectorization.

In my case, we’ll have to live with 256 bits. This means in this register we can either put 4 64-bit values, 8 32-bit values 16 16-bit values, or even 32 8-bit values.

We can use the utility function simdwidthof to tell us how many 32-bit floating point numbers will fit in our register:

print(simdwidthof[DType.float32]())

One of the new features in Mojo 0.5.0 is that SIMD types will default to the width of the architecture, meaning if we call:

simd_float_32 = SIMD[DType.float32]()
print(len(simd_float_32))

simd_int_8 = SIMD[DType.int8]()
print(len(simd_int_8))

Mojo will automatically pack 8 32-bit values, or 32 8-bit values int the register.

This is equivalent to calling:

SIMD[DType.float32, simdwidthof[DType.float32]()]

SIMD operations

Operations over SIMD types are quite intuitive. Let’s try adding two SIMD objects together:

simd_a = SIMD[DType.float32, 8](1.0)
simd_b = SIMD[DType.float32, 8](2.0)

print(simd_a + simd_b)

[3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0]

Additionally, since the version 0.5.0, we can also concatenate SIMD objects with .join():

print(simd_a.join(simd_b))

[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]

Operations applied to a SIMD object will be applied element-wise to the data in it, if the function is set up to handle this:

from math import sqrt

simd_c = SIMD[DType.int32, 4](4, 16, 49, 144)
simd_c_sqrt = sqrt(simd_c)
print(simd_c_sqrt)

[2, 4, 7, 12]

As far as I can tell, this doesn’t just work automatically. If I define a function as:

fn single_add_two(a: Float32) -> Float32:
    return a + 2.0

Then applying it to a single floating point number works as expected:

print(single_add_two(2.0))

4.0

But trying this on a SIMD object does not:

print(single_add_two(simd_a))

error: Expression [85]:1:21: invalid call to 'single_add_two': argument #0 cannot be converted from 'SIMD[f32, 8]' to 'SIMD[f32, 1]'

print(single_add_two(simd_a))

      ~~~~~~~~~~~~~~^~~~~~~~



Expression [83]:1:1: function declared here

fn single_add_two(a: Float32) -> Float32:

^



expression failed to parse (no further compiler diagnostics)

However, if I define a version of the function to take a SIMD object:

fn simd_add_two[simd_width: Int](a: SIMD[DType.float32, simd_width]) -> SIMD[DType.float32, simd_width]:
    return a + 2.0

Then (with the additional specification of the parameter simd_width), it will apply the function to all the values:

simd_ones = SIMD[DType.float32, 8](1.0)
print(simd_add_two(simd_ones))

[3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0]

While still working on single floating point values, as they are just SIMD objects of width one under the hood:

let one: Float32 = 1.0
print(simd_add_two(one))

3.0

I do miss the flexibility of Julia a bit, where you can define one function and then vectorize it with a dot, i.e. if you have a function myfun(scalar) that operates on scalar values, then calling myfun.(vector) will apply it element-wise to all values of that vector, and return a vector of the same shape.
But for the most part, defining functions just to apply to SIMD values in Mojo doesn’t lose you much generality anyway.

Conclusions:

To be honest, I was a little bit daunted when I first saw the SIMD datatype in Mojo. I vaguely remember playing around with SIMD in C++, where it can be quite complicate to implement SIMD operations.
But in Mojo, it really is transparent and relatively straightforward to get going with SIMD.

It is clear that exploiting vectorization is a top priority for the Modular team, and a lot of through has clearly gone into making it easy to exploit the SIMD capabilities of modern hardware.

I might take a look at vectorization vs parallelization in Mojo in the future, and maybe even try my hand at a bit of benchmarking.

Citation

BibTeX citation:

@online{schenck2023,
  author = {Schenck, Ferdinand},
  title = {Mojo 0.5.0 and {SIMD}},
  date = {2023-11-17},
  url = {https://fnands.com/blog/2023/mojo-simd/},
  langid = {en}
}

For attribution, please cite this work as:

Schenck, Ferdinand. 2023. “Mojo 0.5.0 and SIMD.” November 17. https://fnands.com/blog/2023/mojo-simd/.

Parameters in Mojo

Ferdinand Schenck — Wed, 25 Oct 2023 00:00:00 GMT

As Mojo is an extremely new language I want to keep track of the development, and try to learn the language as it evolves.
I had a first look at version 0.2.1 of the language in a blog post a few weeks back, and while on vacation I decided I should probably try and write a blog post every time the Modular team releases a new Mojo update (at least every minor update, not patches).

To my surprise, in the two weeks I was offline the Modular team managed to do two minor releases, so I’ll jump straight to version 0.4.0.

Quick overview

One of the issues I ran into in my last post was due to the fact that Mojo had no native file handling, and I had to invoke Python just to open a simple file. But this has been fixed now, and Mojo now has a very Pythonic way of opening files:

with open("example.txt", "r") as f:
    print(f.read())

This file
has two lines

Scanning the changelog, version 0.3.0 seems to bring changes mostly related to supporting keyword arguments, while 0.4.0 seems to mostly be related to how parameters are handled, e.g. adding default parameters.

This brings up an interesting point: What exactly are parameters in Mojo?

Parameters in Mojo

When searching for the definition of parameters vs. arguments in Python, I get: > A parameter is the variable listed inside the parentheses in the function definition.
An argument is the value that are sent to the function when it is called.

I.e. in Python it is the difference between the name in the function/method definition, vs. the actual data passed.
To be honest, I have never heard anyone really making this distinction.

In Mojo however takes a different route here and makes a stronger distinction between parameter and argument, as the following lines from the Mojo documentation shows:

Python developers use the words “arguments” and “parameters” fairly interchangeably for “things that are passed into functions.” We decided to reclaim “parameter” and “parameter expression” to represent a compile-time value in Mojo, and continue to use “argument” and “expression” to refer to runtime values. This allows us to align around words like “parameterized” and “parametric” for compile-time metaprogramming.

In Mojo, arguments are denoted by round brackets like in Python, parameter values are denoted by square brackets.

So now we can define:

fn foo[parameter: Int](argument: Int) -> Int:
    return parameter + argument

and call it as:

print(foo[10](5))

So the obvious question is: why?
Why separate arguments and parameters when they seem to do the same thing?

The above statement worked because we are calling it with a fixed parameter value that is known at compile time.
Let’s try and pass a variable instead:

for i in range(5):
    print(foo[i](5))

error: Expression [66]:2:15: cannot use a dynamic value in call parameter

    print(foo[i](5))

              ^



expression failed to parse (no further compiler diagnostics)

We get an error.
Parameter values must be known at compile time, while argument values can be passed at runtime.
So passing the variable to the argument will work:

for i in range(5):
    print(foo[10](i))

This is another tool one can use in Mojo for optimization: in Python, all arguments are evaluated at runtime, while Mojo adds the option of adding values that are known at compile time as parameters.

This might be useful in cases where you have some generic version of a function that you might use several versions of.

As an example, let’s write a sliding window summation function that can take different window sizes as parameters:

from random import rand
from tensor import Tensor, TensorSpec, TensorShape

fn sliding_sum[window_size: Int](data: Tensor[DType.float32]) -> Tensor[DType.float32]:
    
    let contracted_size: Int = data.shape()[0] - window_size

    let out_spec = TensorSpec(DType.float32, contracted_size)
    var out_tensor = Tensor[DType.float32](out_spec)

    for i in range(contracted_size):
        out_tensor[i] = 0
        for j in range(window_size):
            out_tensor[i] += data[i + j]

    return out_tensor

The above function takes a window size as a parameter, and a tensor as an argument.
The function will iterate over the tensor, and add all values in a sliding window of size window_size.

As is often the case when creating a convolutional neural network, the size of the kernel is known at compile time, but the data is unknown, allowing the compiler to optimize parts of the operation that are known ahead of time. This means we can define a generic version of the function once, and then still have compiled version of the specific functions we want.

Let’s define a function that calls two versions of the summation:

fn sum_my_tensor(data: Tensor[DType.float32]) -> Tensor[DType.float32]:

    let summed_by_3 = sliding_sum[3](data)
    let summed_by_5 = sliding_sum[5](summed_by_3)

    return summed_by_5

And let’s test it:

let random_tensor = rand[DType.float32](10)
print(sum_my_tensor(random_tensor).shape()[0])

I haven’t yet found an elegant way to print tensors in Mojo. In Python/Pytorch, you can just call print(tensor) to see all the values in a tensor. In Mojo this results in an error.

for i in range(random_tensor.shape()[0]):
    print(random_tensor[i])

0.014950568787753582
0.64297890663146973
0.91784769296646118
0.97008669376373291
0.84397536516189575
0.45575159788131714
0.022083699703216553
0.70694994926452637
0.43663844466209412
0.75171011686325073

summed_by_3 = sliding_sum[3](random_tensor)
for i in range(summed_by_3.shape()[0]):
    print(summed_by_3[i])

1.5757771730422974
2.5309133529663086
2.7319097518920898
2.2698137760162354
1.3218107223510742
1.1847852468490601
1.1656720638275146

summed_by_5 = sliding_sum[5](summed_by_3)
for i in range(summed_by_5.shape()[0]):
    print(summed_by_5[i])

10.430224418640137
10.03923225402832

Conclusions

I read a tweet by Mark Tenenholtz recently that rang true:

Writing a lot of Mojo code and it feels like writing C on hard mode.

No dict-like structures, file IO is super rudimentary, docs are very sparse, etc.

It’s actually very fulfilling, though.
— Mark Tenenholtz (@marktenenholtz) October 23, 2023

As pointed out in the replies, even asking ChatGPT doesn’t get you very far as the language was released after the current training cut-off date (coupled with the fact that there is no corpus of StackOverflow posts to train on anyway).

Citation

BibTeX citation:

@online{schenck2023,
  author = {Schenck, Ferdinand},
  title = {Parameters in {Mojo}},
  date = {2023-10-25},
  url = {https://fnands.com/blog/2023/mojo-parameters/},
  langid = {en}
}

For attribution, please cite this work as:

Schenck, Ferdinand. 2023. “Parameters in Mojo.” October 25. https://fnands.com/blog/2023/mojo-parameters/.

Stereo vision and disparity maps (in Julia)

Ferdinand Schenck — Mon, 18 Sep 2023 00:00:00 GMT

I’ve been working a lot recently with stereo vision and wanted to go through the basics of how disparity is calculated. I’m partially doing this as an excuse to get better at Julia (v1.9.3 used here).

You can view the notebook for this blog post on Github:

Introduction

In much the same way that we as humans can have depth perception by sensing the difference in the images we see between our left and right eyes, we can calculate depth from a pair of images taken from different locations, called a stereo pair.

If we know the positions of out cameras, then we can use matching points in our two images to estimate how far away from the camera those points are.

Taking a look at the image below (from OpenCV):

https://docs.opencv.org/4.x/dd/d53/tutorial_py_depthmap.html

If we have two identical cameras, at points and at a distance from each other, with focal length , we can calculate the distance () to object by using the disparity between where the object appears in the left image () and where it appears in the right image ().

In this simple case, the relation between disparity and distance is simply:

If we know an , then we can rearrange this to give us distance as a function of disparity:

You might notice that in case the disparity is zero, you will have an undefined result. This is just due to the fact that in this case the cameras are pointing in parallel, so in principle a disparity of zero should not be possible.

The general case is more complicated, but we will focus on this simple setup for now.

We can define the function as:

function distance_from_disparity(B, f, disparity)
    B*f/disparity
end

distance_from_disparity (generic function with 1 method)

Where and are measured in pixels, and is measured in centimeters.

There is an inverse relation between distance and disparity:

using Plots

disparities = range(1, 50, length=50)
distances = distance_from_disparity.(2000, 0.1, disparities)

plot(disparities, distances, label="Distance [cm]")
xlabel!("Disparity")
ylabel!("Distance [cm]")

So once we have a disparity, it’s relatively straightforward to get a distance. But how do we find disparities?

Disparity maps

We usually represent the disparities for a given pair of images as a disparity map, which is an array with the same dimensions as (one of) your images, but with disparity values for each pixel.

In principle, this is a two-dimensional problem, as an object might be matched to a point that has both a horizontal and vertical shift, but luckily, you can always find a transformation to turn this into a one dimensional problem.

The cartoon below illustrates what a disparity map might look like:

Above, we calculate the disparity with respect to the right image (you can do it with respect to the left image as well), and as you can see the disparity map tells us how many pixels to the right each object shifted in the left image vs the right image.

For a set of images (taken from the Middlebury Stereo Datasets):

The corresponding disparity map can be visualized as follows:

https://vision.middlebury.edu/stereo/eval/newEval/tsukuba/

With darker pixels having lower disparity values, and brighter pixels having higher disparity values, meaning the dark objects are far away from the cameras, while the bright ones are close.

The ground truth disparity as shown above is usually calculated from LiDAR or some other accurate method, and our goal is to get as close as possible to those values using only the images above.

A naive approach

So let’s try and calculate disparity for the images above.
There are many, many approaches to calculating disparity, but let us begin with the most simple approach we can think of.
As a start, let us go through each pixel in the right image, and for that pixel, try and find the most similar pixel in the left image.

So let us try and take the squared difference between pixels values as our similarity metric. As we are going to be doing the same thing for every row of pixels, we are just going to define a function that does the basic logic, and then apply the same function to every case.

function smallest_diff(pixel, row, metric)

    disparity_candidates = metric.(pixel, row)
    # Minus one as Julia counts from 1
    argmin(disparity_candidates) - 1
end

smallest_diff (generic function with 1 method)

Let’s define a distance metric as the squared distance:

squared_difference = (x, y) -> sum((x-y).^2)

#9 (generic function with 1 method)

And as a test case let’s create the cartoon image we had above:

left_image = ones(Float64, (8, 8))
right_image = ones(Float64, (8, 8))
disparity = ones(UInt8, (8, 8))

right_image[2, 1] = 0
right_image[2, 2] = 0
right_image[2, 3] = 0
right_image[2, 6] = 0
right_image[2, 7] = 0
right_image[2, 8] = 0

left_image[7, 1] = 0
left_image[7, 2] = 0
left_image[7, 3] = 0
left_image[5, 6] = 0
left_image[5, 7] = 0
left_image[5, 8] = 0

Gray.(left_image')

Gray.(right_image')

Now we can try and match pixels in the right image to pixels in the left image.

# NOTE: Julia is column major, not row major like Python or C/C++, 
# so this looping is not optimal. But I'm going for clarity, not speed. 
for i in range(1, size(right_image)[1])
    for j in range(1, size(right_image)[2])
        disparity[j, i] = max(smallest_diff(right_image[j, i], left_image[j:end, i], squared_difference), 1)
    end
end

So how did we do?

Gray.(disparity'/5)

So the toy example works! The top line, which moved more pixels, shows up brighter (i.e. larger disparity values), and the lower line is dimmer.

So let’s move on to real images. We’ll start with the example case above, but for simplicity we’ll stick to grayscale at first:

using Images, FileIO

tsukuba_left_gray = Gray.(load("im3.png"))
tsukuba_right_gray = Gray.(load("im4.png"))


tsukuba_right_gray

"""
Iterrate over the rows of the right image. 
For each row, find the most similar pixel in the left image and return the difference in location for them (disparity). 
The similarity is defined by the provided metric. 
"""
function pixel_match(left_image, right_image, metric, max_disp)

    num_rows = size(right_image)[2]
    num_cols = size(right_image)[3]

    disparity = zeros((num_rows, num_cols))
    for i in range(1, num_rows)
        for j in range(1, num_cols)
            terminator = min(j+max_disp, num_cols)
            pixel = right_image[:, i, j]
            row = left_image[:, i, j:terminator]
            
            disparity[i, j] = max(smallest_diff(pixel, row, metric), 1)
        end
    end

    disparity
end

pixel_match

Redefining smallest_diff slightly…

function smallest_diff(pixel, row, metric)
    disparity_candidates = metric.(pixel, row)
    # Minus one as Julia counts from 1
    argmin(disparity_candidates)[2] - 1
end

smallest_diff (generic function with 1 method)

# Transform images to Float64 to avoid integer overflow when calculating metric. 
# Reshape to 3 dim array with 1 channel. Julia is channel first
array_size = (1, size(tsukuba_left_gray)...,)
tsukuba_left_gray_float = reshape(Float64.(tsukuba_left_gray), array_size)
tsukuba_right_gray_float = reshape(Float64.(tsukuba_right_gray), array_size)


tsukuba_disparity_gray = pixel_match(tsukuba_left_gray_float, tsukuba_right_gray_float,  squared_difference, 20);

So let’s see how we did?

max_plot_disp = 15
Gray.(tsukuba_disparity_gray / max_plot_disp)

Looking at the predicted disparity, we can see there is some vague resemblance to the input image, but we’re still pretty far from the target:

https://vision.middlebury.edu/stereo/eval/newEval/tsukuba/

A significant problem seems to be erroneous matches, especially in the background.

As you can imagine, we are only comparing single channel pixels values, and it’s very likely that we might just find a better match by chance. In grayscale we are only matching pixel intensity, and we have no idea whether something is bright green, or bright red.

So let’s try and improve the odds of a good match by adding colour.

tsukuba_left_rgb = load("im3.png")
tsukuba_right_rgb = load("im4.png")
 

tsukuba_left_rgb

tsukuba_left_rgb_float = Float64.(channelview(tsukuba_left_rgb))
tsukuba_right_rgb_float = Float64.(channelview(tsukuba_right_rgb))
tsukuba_disparity_rgb = pixel_match(tsukuba_left_rgb_float, tsukuba_right_rgb_float,  squared_difference, 20);

Gray.(tsukuba_disparity_rgb / max_plot_disp)

So, a slight improvement! There seem to be fewer random matches in the background, but still not that close to the desired outcome.

Is there more we can do?

Block matching

The obvious downside of the naive approach above is that it only ever looks at one pixel (in each image) at a time.
That’s not a lot of information, and also not how we intuitively match objects.

Look at the image below. Can you guess the best match for the pixel in the row of pixels below it?

Own work

Given only this information, it’s impossible for us to guess whether the green pixel matches with the pixels at location 3, 5 or 7.

If however I was to give you more context, i.e. a block of say 3x3 pixels, would this make things simpler?

Own work

In this case, there is an unambiguous answer, which is the principle behind block-matching.

To confirm our idea that more context results in better matches, we can take a quick look at a row of pixels:

max_disp = 100
block_size = 0
i = 100
j = 220
test_row = tsukuba_left_rgb[i-block_size:i+block_size, j-block_size:j+max_disp];
test_block = tsukuba_right_rgb[i-block_size:i+block_size, j-block_size:j+block_size];
test_block

Given the pixel above, where in the row below do you think this pixel matches?
You would guess somewhere in the orange part on the left right? But which pixel exactly is almost impossible to say.

test_row

If we now take a block with more context:

block_size = 5
test_row = tsukuba_left_rgb[i-block_size:i+block_size, j-block_size:j+max_disp];
test_block = tsukuba_right_rgb[i-block_size:i+block_size, j-block_size:j+block_size];
test_block

And compare it to the row below, the location of the match becomes more obvious:

test_row

Calculating the difference metric for each point with different block sizes, we can clearly see that for low block sizes, the lowest metric value is ambiguous, while for larger block sizes it becomes more clear where exactly the best match is:

sum_squared_difference_rgb_block = (x, y) -> sum((x - y).^2)

function smallest_block_diff(block, row, metric, block_size)

    view_length = size(row)[3]
    disparity_candidates = zeros(view_length - 2*block_size)

    for i in range(block_size + 1, view_length - block_size )
        
        disparity_candidates[i-block_size] = metric(block, row[:, :, i-block_size:i+block_size])
    end
    disparity_candidates 

end

p1 = plot()
for b in [0, 2, 4]
    block_size = b
    i = 100
    j = 220
    max_disp = 100
    test_row = tsukuba_left_rgb[i-block_size:i+block_size, j-block_size:j+max_disp];
    test_block = tsukuba_right_rgb[i-block_size:i+block_size, j-block_size:j+block_size];
    x = Float64.(channelview(test_block))
    y = Float64.(channelview(test_row))
    diffs = smallest_block_diff(x, y, sum_squared_difference_rgb_block, block_size)

    norm_diffs = (diffs .- minimum(diffs))./maximum(diffs)

    plot!(p1, norm_diffs, label="Blocksize $block_size", linewidth=2)

end

p1

And now we are ready to define our block matching algorithm, much in the way we did our pixel matching algorithm:

function block_match(left_image, right_image, block_size, metric, max_disp)

    
    num_rows = size(left_image)[2]
    num_cols = size(left_image)[3]

    disparity = zeros((num_rows, num_cols))
    for i in range(block_size + 1, num_rows - block_size)
        for j in range(block_size + 1, num_cols - block_size)
            terminator = min(j+max_disp, num_cols)

            row_block_index = i-block_size:i+block_size
            col_block_index = j-block_size:j+block_size
            
            block = right_image[:, row_block_index, col_block_index]
            row = left_image[:, row_block_index, j-block_size:terminator]

            disparity_candidates = smallest_block_diff(block, row, metric, block_size)

            disparity[i, j] = max(argmin(disparity_candidates) -1, 1)
        end
    end

    disparity
end

block_match (generic function with 1 method)

Let’s see how this does on the full image in comparison to the pixel matching:

tsukuba_disparity_rgb = block_match(tsukuba_left_rgb_float, tsukuba_right_rgb_float, 2,  sum_squared_difference_rgb_block, 20);
Gray.(tsukuba_disparity_rgb / max_plot_disp)

Now we are getting somewhere! Compared to the earlier results we can now start making out the depth of the separate objects like the lamp, bust and camera.

There are still a few things we could do to improve our simple algorithm (like only accepting matches that have below a certain score for the metric), but I will leave those as an exercise to the reader.

Conclusions

Above we went through a basic introduction to stereo vision and disparity, and built a bare-bones block matching algorithm from scratch.

The above is pretty far away from the state of the art, and there are many more advanced methods for calculating disparity, ranging from relatively simple methods like block matching to Deep Learning methods.

Citation

BibTeX citation:

@online{schenck2023,
  author = {Schenck, Ferdinand},
  title = {Stereo Vision and Disparity Maps (in {Julia)}},
  date = {2023-09-18},
  url = {https://fnands.com/blog/2023/stereo-vision-julia/},
  langid = {en}
}

For attribution, please cite this work as:

Schenck, Ferdinand. 2023. “Stereo Vision and Disparity Maps (in Julia).” September 18. https://fnands.com/blog/2023/stereo-vision-julia/.

A first look at Mojo 🔥

Ferdinand Schenck — Mon, 11 Sep 2023 00:00:00 GMT

The Mojo programming language was officially released in May, but could only be used through some notebooks in a sandbox.

Last week, the SDK (version 0.2.1) got released, so I decided to give it a look.

Mojo’s goal is to “combine the usability of Python with the performance of C”, and bills itself as “the programming language for all AI developers”.

It’s clear that Python is the dominant language when it comes to ML/AI, with great libraries like Pytorch and a few others being the main drivers of that. The problem comes with depth: all the fast libraries in Python are written in a performant language, usually C or C++, which means that if you want to dig into the internals of the tools you are using you have to switch languages, which greatly raises the barrier of entry for doing so.

There are other languages that try to go for the usability of Python while retaining performance, and the first language that comes to mind for me in this respect is Julia. Julia is a pretty neat language, and writing math-heavy, fast code in it feels very elegant, while retaining a very Python like syntax. Julia is about twenty years younger than Python, and to me seems like they took the best aspects of Python and Fortran and rolled them into one language, allowing you to have performant and elegant code that is Julia all the way down. Given all this, in vacuum, Julia would seem like the obvious language to choose when it comes to ML/AI programming.

The one major downside of Julia is that it doesn’t have the robust ecosystem of libraries that Python has, and unless something major changes, it seems that Python will keep winning.

Enter Mojo, a language that then (aspires to) keep interoperability with Python, while itself being very performant and allowing you to write code that is Mojo all the way down. Basically if Mojo achieves its goals then we get to have our cake and eat it: we can keep the great ecosystem of packages that Python brings with it, while getting to write new performant code in a single. My guess is if this works out that all the major packages will eventually get rewritten in Mojo, but we can have a transition period where we still get to keep the C/C++ version of them until this can be done.

The people behind Mojo (mostly Chris Lattner) seem to know what they are doing, so I wish them all the best.

A quick look

I wanted to start with something basic, so I thought I would have a look at the first puzzle from the 2022 advent of code. Basically you are given a text file with a several lists of numbers representing the amount of calories some elves are carrying (go read up on the advent of code if you are unfamiliar, it will make sense then), and have to find which elves are carrying the most calories.

So effectively a little bit of file parsing, with some basic arithmetic, i.e. a little puzzle to ease into Mojo. I won’t share the input because the creator of the AoC has explicitly asked people not to, but you can download your own and try the code below.

At first glance, a lot of Python code will “just work”:

for _ in range(10):
    print("Blah")

Blah
Blah
Blah
Blah
Blah
Blah
Blah
Blah
Blah
Blah

However, it’s clear a lot is still missing, e.g. lambda functions don’t work yet:

lambda x: x + 2

error: Expression [5]:1:1: Mojo doesn't support lambda expressions yet

lambda x: x + 2

^

expression failed to parse (no further compiler diagnostics)

This is likely coming, but for now we have to live without it.

So for the first step, let’s parse some text files. The first thing I found was that Mojo doesn’t have a native way to parse text yet. But luckily, you can just get Python to do it for you! In this case, you have to import Python as a module and call the builtin Python open function.

It’s standard practice in Python to open text files with the with open(filename) as f incantation, but this doesn’t work in Mojo, so have to open and close files manually.

from python import Python
from math import max
from utils.vector import DynamicVector
from algorithm.sort import sort

def read_file(file_name: String) -> DynamicVector[Int]:
    """ 
    There seems to be no native mojo way to read a file yet.

    There is an issue in the mojo github page that suggests the code below.
    https://github.com/modularml/mojo/issues/130

    I've tried to keep the Python code in this function, in order to keep the others
    as "pure" mojo
    """
    builtins = Python.import_module("builtins")
    in_file = builtins.open(file_name)

    file_contents = in_file.read()
    in_file.close()

    content_list = file_contents.split("\n")
    # There seems to be no way yet to go straight from python int to mojo Int,
    # so you have to go via float64: https://github.com/modularml/mojo/issues/657
    let list_len: Int = content_list.__len__().to_float64().to_int()

    item_list = DynamicVector[Int](list_len)
    
    for item in content_list:
        str_item = item.to_string()
        if str_item != "":
            item_list.push_back(atol(str_item))
        else:
            item_list.push_back(0)
        
    return item_list

All in all, it’s relatively standard Python, with a couple of caveats.

One of the big things is that there is a distinction between Python types and Mojo types, i.e. the Python int is not the same as Mojo’s Int, so if you want to get the most out of Mojo, you need to cast from the one to the other. Right now, there seems to be no direct way to go from int to Int, so I had to take a detour via float64.

I tried to keep the Python imports in the read_file function, so that the other functions can be in “pure” Mojo.

The my first impulse was to create a Python-esque list, but the builtin list in Mojo is immutable, so I had to go for a DynamicVector, which had a strong C++ flavour to it.

Once that was done I was done with Python for this program and could go forth in pure Mojo.

Below you can see I declare functions with fn while above I used def. Both work in Mojo, but fn functions forces you to be strongly typed and enfoces some memory safe behaviour.

fn part_one(calorie_list: DynamicVector[Int]) -> Int:
    """ 
    Loop over a vector of Ints, and find the grouping (split by 0) with the highest sum. 
    """
    var max_calories: Int = 0
    var this_calories: Int = 0

    var entry = 0
    for index in range(len(calorie_list)):
        entry = calorie_list.data[index]
        if entry != 0:
            this_calories += entry
        else:
            max_calories = max(this_calories, max_calories)
            this_calories = 0

    return max_calories

You can see here the values are all declared as mutable (var). You can also declare immutables with let. This is enforced in fn functions.

Other than that, a relatively standard loop over a container.

fn part_two(calorie_list: DynamicVector[Int]) -> Int:
    """
    Initialize a vector to keep track of the current top 3 elf calories. 
    Add a value to the container if is larger than the smallerst value, and sort. 
    """
    let k: Int = 3
    var max_k_calories: DynamicVector[Int] = DynamicVector[Int](k)
    for i in range(k):
        max_k_calories.push_back(0)

    var this_calories: Int = 0

    var entry = 0
    for index in range(len(calorie_list)):
        entry = calorie_list.data[index]
        if entry != 0:
            this_calories += entry
        else:
            if this_calories > max_k_calories[0]:
                max_k_calories[0] = this_calories
                sort(max_k_calories)

            this_calories = 0
            

    var max_calories: Int = 0
    for index in range(len(max_k_calories)):
        max_calories += max_k_calories.data[index]

    return max_calories

Again, relatively straightforward.

I’m definitely missing Python niceties like being able to easily sum over a container (can’t call sum(max_k_calories) in Mojo 😢).

To put it all together we create a main fn, and notice that we need to indicate that it might raise errors as we are calling the unsafe read_file.

fn main() raises:
    let file_contents: DynamicVector[Int] = read_file(
        "input/input_1.txt"
    )

    let answer_1: Int = part_one(file_contents)
    print("The elf carrying the most calories is carrying:", answer_1, "calories")

    let answer_2: Int = part_two(file_contents)
    print("The three elves carrying the most calories are carrying:", answer_2, "calories")


main()

The elf carrying the most calories is carrying: 72511 calories
The three elves carrying the most calories are carrying: 212117 calories

Conclusions

Mojo feels relatively familiar, but I will also say that when writing “pure” Mojo it feels like writing C with Python syntax. This makes sense given the goals of the language, but caught me a little off guard; I was expecting something a little closer to Julia, which still feels a lot like Python in most cases.

This was not the greatest example to show Mojo off, as Mojo really shines in high performance environments, so the language didn’t really get to stretch its legs here. You can find some more performance oriented examples on the official Mojo website.

I will probably give Mojo another look and try out something a bit more suited for the language in the future, maybe when the 0.3 version of the language drops.

I think I’ve been spoiled by mostly writing in two well supported languages (Python and C++) for which there are countless reference examples or StackOverflow posts on how to do things. Due to the fact that Mojo is brand new, there are very few examples to look to about how to do even relatively basic things.

For now if you want to get started, I recommend starting with the exercises on mojodojo.dev.

Citation

BibTeX citation:

@online{schenck2023,
  author = {Schenck, Ferdinand},
  title = {A First Look at {Mojo} 🔥},
  date = {2023-09-11},
  url = {https://fnands.com/blog/2023/mojo-advent-of-code/},
  langid = {en}
}

For attribution, please cite this work as:

Schenck, Ferdinand. 2023. “A First Look at Mojo 🔥.” September 11. https://fnands.com/blog/2023/mojo-advent-of-code/.