Skip to content

Remove hw from tables for clarity #87

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jenriver merged 1 commit into from
Nov 18, 2025
Merged

Remove hw from tables for clarity #87

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions bonsai/models/densenet121/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,11 @@

This directory contains a pure JAX implementation of the [DenseNet-121](https://huggingface.co/keras/densenet_121_imagenet) model, using the Flax NNX API.

## Tested on
| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) |
|:--------------------------------------------------------| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |
| [DenseNet-121](https://huggingface.co/keras/densenet_121_imagenet) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |
## Model Configuration Support Status

| Model Name | Config Support Status |
| :--- | :--- |
| [DenseNet-121](https://huggingface.co/keras/densenet_121_imagenet) | **✅ Supported** |

### Running this model

Expand Down
14 changes: 5 additions & 9 deletions bonsai/models/efficientnet/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,16 +3,12 @@
This directory contains a pure JAX implementation of the [Efficientnet](https://arxiv.org/abs/1905.11946), using the [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html) API.


## Tested on:
*(Last Updated: 2025-09-19)*



| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) |
| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |
| **Model** | | | | | | | | |
| [Efficientnet](https://arxiv.org/abs/1905.11946) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |❔ Needs check | ❔ Needs check |
## Model Configuration Support Status

| Model Name | Config Support Status |
| :--- | :--- |
| **Model** | |
| [Efficientnet](https://arxiv.org/abs/1905.11946) | **✅ Supported** |

### Running this model

Expand Down
11 changes: 5 additions & 6 deletions bonsai/models/llada_8b/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,12 @@

This directory contains a pure JAX implementation of the [LLaDA diffusion model](https://huggingface.co/GSAI-ML/LLaDA-8B-Instruct), using the [Flax NNX](flax.readthedocs.io/en/stable/index.html) API.

## Tested on:
*(Last Updated: 2025-10-20)*
## Model Configuration Support Status

| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) |
| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |
| **Dense Models** | | | | | | | | |
| [LLaDA-8B-Instruct](https://huggingface.co/GSAI-ML/LLaDA-8B-Instruct) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ✅ Runs |
| Model Name | Config Support Status |
| :--- | :--- |
| **Dense Models** | |
| [LLaDA-8B-Instruct](https://huggingface.co/GSAI-ML/LLaDA-8B-Instruct) | **✅ Supported** |

## How to contribute to this model

Expand Down
36 changes: 15 additions & 21 deletions bonsai/models/qwen3/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,27 +2,21 @@

This directory contains a pure JAX implementation of the [Qwen3 language model](https://qwenlm.github.io/blog/qwen3/), using the [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html) API.

> [!IMPORTANT]
> For large-scale high performance use case, please see the [MaxText](https://github.com/AI-Hypercomputer/maxtext?tab=readme-ov-file#getting-started) version.


## Tested on:
*(Last Updated: 2025-07-02)*



| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) |
| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |
| **Dense Models** | | | | | | | | |
| [Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) | ✅ Supported | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs |
| [Qwen3-1.7B](https://huggingface.co/Qwen/Qwen3-1.7B) | ✅ Supported | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs |
| [Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B) | ✅ Supported | ❔ Needs check | ❔ Needs check | ✅ Runs | ❔ Needs check | ❔ Needs check| ❔ Needs check | ❔ Needs check|
| [Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B) | ✅ Supported | ❔ Needs check | ❔ Needs check| ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |
| [Qwen3-14B](https://huggingface.co/Qwen/Qwen3-14B) | ✅ Supported | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |
| [Qwen3-32B](https://huggingface.co/Qwen/Qwen3-32B) | ✅ Supported | ❔ Needs check | ❔ Needs check | ⛔️ Not supported | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |
| **MoE Models** | | | | | | | | |
| [Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | 🟡 Not started | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |
| [Qwen3-235B-A22B](https://huggingface.co/Qwen/Qwen3-235B-A22B) | 🟡 Not started | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |

## Model Configuration Support Status

| Model Name | Config Support Status |
| :--- | :--- |
| **Dense Models** | |
| [Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) | **✅ Supported** |
| [Qwen3-1.7B](https://huggingface.co/Qwen/Qwen3-1.7B) | **✅ Supported** |
| [Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B) | **✅ Supported** |
| [Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B) | **✅ Supported** |
| [Qwen3-14B](https://huggingface.co/Qwen/Qwen3-14B) | **✅ Supported** |
| [Qwen3-32B](https://huggingface.co/Qwen/Qwen3-32B) | **✅ Supported** |
| **MoE Models** | |
| [Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | **🟡 Not started** |
| [Qwen3-235B-A22B](https://huggingface.co/Qwen/Qwen3-235B-A22B) | **🟡 Not started** |


### Running this model
Expand Down
8 changes: 4 additions & 4 deletions bonsai/models/resnet/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,11 @@

This directory contains a pure JAX implementation of the [ResNet-50](https://huggingface.co/microsoft/resnet-50) model, using the Flax NNX API.

## Tested on
## Model Configuration Support Status

| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) |
| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |
| [ResNet-50 v1.5](https://huggingface.co/microsoft/resnet-50) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |
| Model Name | Config Support Status |
| :--- | :--- |
| [ResNet-50 v1.5](https://huggingface.co/microsoft/resnet-50) | **✅ Supported** |

### Running this model

Expand Down
18 changes: 9 additions & 9 deletions bonsai/models/sam2/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,15 @@

A minimal, readable JAX + Flax NNX re-implementation of Meta’s [Segment Anything 2](https://github.com/facebookresearch/sam2), enabling promptable image and video segmentation.

## Tested on:

| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) |
|------------------|--------------|:--------:|:-------------:|:-------------:|:-------------:|:-------------:|:-----------:|:------------:|
| **SAM2 Variants**| | | | | | | | |
| `sam2_tiny` | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |
| `sam2_small` | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |
| `sam2_baseplus` | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |
| `sam2_large` | ✅ Supported | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |
## Model Configuration Support Status: SAM2 Variants

| Model Name | Config Support Status |
| :--- | :--- |
| **SAM2 Variants** | |
| `sam2_tiny` | **✅ Supported** |
| `sam2_small` | **✅ Supported** |
| `sam2_baseplus` | **✅ Supported** |
| `sam2_large` | **✅ Supported** |


### Running this model
Expand Down
13 changes: 5 additions & 8 deletions bonsai/models/unet/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,12 @@ Needs work on:
3. Clean up `UNet_segmentation_example.ipynb` to demonstrate proper parameter loading.


## Model Configuration Support Status

## Tested on:



| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) |
| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |
| **Model** | | | | | | | | |
| [U-Net](https://arxiv.org/abs/1505.04597) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |❔ Needs check | ❔ Needs check |
| Model Name | Config Support Status |
| :--- | :--- |
| **Model** | |
| [U-Net](https://arxiv.org/abs/1505.04597) | **✅ Supported** |


### Running this model
Expand Down
13 changes: 5 additions & 8 deletions bonsai/models/vae/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,12 @@
This directory contains a pure JAX implementation of the [VAE - Variational Autoencoder](https://arxiv.org/abs/1312.6114), using the [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html) API, including procedures to train, save and load weights, and make inference.


## Tested on:
*(Last Updated: 2025-09-19)*
## Model Configuration Support Status



| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) |
| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |
| **Model** | | | | | | | | |
| [VAE - Variational Autoencoder](https://arxiv.org/abs/1312.6114) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |❔ Needs check | ❔ Needs check |
| Model Name | Config Support Status |
| :--- | :--- |
| **Model** | |
| [VAE - Variational Autoencoder](https://arxiv.org/abs/1312.6114) | **✅ Supported** |


### Running this model
Expand Down
9 changes: 5 additions & 4 deletions bonsai/models/vgg19/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,11 @@

This directory contains a pure JAX implementation of the [VGG-19](https://huggingface.co/keras/vgg_19_imagenet) model, using the Flax NNX API.

## Tested on
| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) |
|:--------------------------------------------------------| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |
| [VGG-19](https://huggingface.co/keras/vgg_19_imagenet) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |
## Model Configuration Support Status

| Model Name | Config Support Status |
| :--- | :--- |
| [VGG-19](https://huggingface.co/keras/vgg_19_imagenet) | **✅ Supported** |

### Running this model

Expand Down
8 changes: 4 additions & 4 deletions bonsai/models/vit/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,11 @@

This directory contains a pure JAX implementation of the [ViT](https://huggingface.co/google/vit-base-patch16-224) model, using the Flax NNX API.

## Tested on
## Model Configuration Support Status

| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) |
| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |
| [ViT](https://huggingface.co/google/vit-base-patch16-224) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |
| Model Name | Config Support Status |
| :--- | :--- |
| [vit-base-patch16-224](https://huggingface.co/google/vit-base-patch16-224) | **✅ Supported** |

### Running this model

Expand Down