Reference Models

Reference Models#

As described in the Introduction and FAQ pages, the PolyBlocks compiler can compile functions written in frameworks like PyTorch, JAX, and TensorFlow. The functions can be from any domain, whether deep learning or other scientific, engineering, data analytics, or high-performance computing domains. Good optimization can be expected as long as they are written using operators on dense tensors/matrices.

While it is hard to state applicability for successful compilation and high performance in general, to provide a reference for coverage for the AI and deep learning domain, below is a list of popular deep-learning models from HuggingFace, torchvision, timm, and torchbench, which have been tested with PolyBlocks and are known to compile successfully and validate against the standard runtime of PyTorch (eager) as well as torch.compile (Torch Inductor). The JAX models listed at the end are validated against JAX’s JIT (XLA) instead.

Many of these models are also available on the PolyBlocks Playground, and they are expected to compile and execute successfully through the Docker release. Any recent regressions are marked with a red cross.

HuggingFace Models#

Model

Status

AlbertForMaskedLM

✅

AlbertForQuestionAnswering

✅

AllenaiLongformerBase

✅

BLOOM

✅

BartForCausalLM

✅

BartForConditionalGeneration

✅

BertForMaskedLM

✅

BertForQuestionAnswering

✅

BlenderbotForCausalLM

✅

BlenderbotSmallForCausalLM

✅

BlenderbotSmallForConditionalGeneration

✅

CamemBert

✅

ConvNext

✅

DETR

✅

DPT Large

✅

DebertaForMaskedLM

✅

DebertaForQuestionAnswering

✅

DebertaV2ForMaskedLM

✅

DebertaV2ForQuestionAnswering

✅

DistilBertForMaskedLM

✅

DistilBertForQuestionAnswering

✅

DistillGPT2

✅

ElectraForCausalLM

✅

ElectraForQuestionAnswering

✅

Flux

✅

GPT2ForSequenceClassification

✅

GTE Small feature extraction

✅

Gemma-2-2b

✅

Google/Deplot

✅

LayoutLMForMaskedLM

✅

LayoutLMForSequenceClassification

✅

Llama 3.1 8B

✅

Llama 3.2 1B Instruct

✅

M2M100ForConditionalGeneration

✅

MBartForCausalLM

✅

MBartForConditionalGeneration

✅

MPNet base v2

✅

MT5ForConditionalGeneration

✅

MPT-7B

✅

MegatronBertForCausalLM

✅

MegatronBertForQuestionAnswering

✅

Mini-lm

✅

MiniCPM

✅

Mistoline

✅

Mistral instruct

✅

MobileBertForMaskedLM

✅

MobileBertForQuestionAnswering

✅

Moondream

✅

nanoLLaVA

✅

OPTForCausalLM

✅

owlvit-base-patch32

✅

PLBartForCausalLM

✅

PLBartForConditionalGeneration

✅

PegasusForCausalLM

✅

PegasusForConditionalGeneration

✅

Query Wellformedness score

✅

RobertaForCausalLM

✅

RobertaForQuestionAnswering

✅

Speech2Text2ForCausalLM

✅

SqlCoder

✅

Stable diffusion 3.5

✅

Stable diffusion turbo Unet block

✅

Stable diffusion XL base 1.0

✅

Stable diffusion image to image XL refiner

✅

Starling-LM-7B-alpha

✅

T5ForConditionalGeneration

✅

T5Small

✅

TableTransformer

✅

TrOCRForCausalLM

✅

XGLMForCausalLM

✅

XLNetLMHeadModel

✅

XLM Roberta Base

✅

YituTechConvBert

✅

YoloS

✅

Zephyr 7B Alpha

✅

TorchVision Models#

Listed below are some TorchVision models that have been tested with PolyBlocks: they compile successfully and validate. Many also run significantly faster with PolyBlocks than with the Torch standard runtime or Torch Inductor.

Model

Status

AlexNet

✅

DenseNet

✅

EfficientNet

✅

GoogleNet

✅

Inception

✅

MNasNet

✅

MobileNetv3

✅

ShuffleNet

✅

SqueezeNet

✅

HRNet

✅

ResNet50

✅

UNet

✅

VGG19

✅

VIT

✅

TIMM Models#

Listed below are some TIMM models that have been tested with PolyBlocks: they compile successfully and validate. Many also run significantly faster with PolyBlocks than with the Torch standard runtime or Torch Inductor.

Model

Status

adv_inception_v3

✅

beit_base_patch16_224

✅

botnet26t_256

✅

cait_m36_384

✅

coat_lite_mini

✅

convit_base

✅

convmixer_768_32

✅

convnext_base

✅

crossvit_9_240

✅

cspdarknet53

✅

deit_base_distilled_patch16_224

✅

dla102

✅

dm_nfnet_f0

✅

dpn107

✅

eca_botnext26ts_256

✅

eca_halonext26ts

✅

ese_vovnet19b_dw

✅

fbnetc_100

✅

fbnetv3_b

✅

gernet_l

✅

ghostnet_100

✅

gluon_inception_v3

✅

gmixer_24_224

✅

gmlp_s16_224

✅

hrnet_w18

✅

inception_v3

✅

jx_nest_base

✅

lcnet_050

✅

levit_128

✅

mixer_b16_224

✅

mixnet_l

✅

mnasnet_100

✅

mobilenetv2_100

✅

mobilenetv3_large_100

✅

mobilevit_s

✅

nfnet_l0

✅

pit_b_224

✅

pnasnet5large

✅

poolformer_m36

✅

regnety_002

✅

repvgg_a2

✅

res2net101_26w_4s

✅

res2net50_14w_8s

✅

res2next50

✅

resmlp_12_224

✅

resnest101e

✅

rexnet_100

✅

sebotnet33ts_256

✅

selecsls42b

✅

spnasnet_100

✅

swin_base_patch4_window7_224

✅

swsl_resnext101_32x16d

✅

tf_efficientnet_b0

✅

tf_mixnet_l

✅

tinynet_a

✅

tnt_s_patch16_224

✅

twins_pcpvt_base

✅

visformer_small

✅

vit_base_patch16_224

✅

xcit_large_24_p8_224

✅

Torchbench Models#

Listed below are some Torchbench models that have been tested with PolyBlocks: they compile successfully and validate. Many also run significantly faster with PolyBlocks than with the Torch standard runtime or Torch Inductor.

Model

Status

torchrec_dlrm

✅

BERT_pytorch

✅

Background_Matting

✅

LearningToPaint

✅

Super_SloMo

✅

alexnet

✅

basic_gnn_edgecnn

✅

basic_gnn_gcn

✅

basic_gnn_gin

✅

basic_gnn_sage

✅

cm3leon_generate

✅

dcgan

✅

demucs

✅

densenet121

✅

dlrm

✅

doctr_reco_predictor

✅

drq

✅

fastNLP_Bert

✅

functorch_dp_cifar10

✅

functorch_maml_omniglot

✅

hf_Albert

✅

hf_Bart

✅

hf_Bert

✅

hf_Bert_large

✅

hf_BigBird

✅

hf_DistilBert

✅

hf_GPT2

✅

hf_GPT2_large

✅

hf_Longformer

✅

hf_Reformer

✅

hf_Roberta_base

✅

hf_T5

✅

hf_T5_base

✅

hf_T5_generate

✅

hf_T5_large

✅

hf_Whisper

✅

hf_distil_whisper

✅

lennard_jones

✅

llama_v2_7b_16h

✅

llava

✅

maml

✅

maml_omniglot

✅

microbench_unbacked_tolist_sum

✅

mnasnet1_0

✅

mobilenet_v2

✅

mobilenet_v3_large

✅

moondream

✅

nvidia_deeprecommender

✅

opacus_cifar10

✅

phlippe_densenet

✅

phlippe_resnet

✅

pyhpc_equation_of_state

✅

pyhpc_isoneutral_mixing

✅

pyhpc_turbulent_kinetic_energy

✅

pytorch_CycleGAN_and_pix2pix

✅

pytorch_stargan

✅

pytorch_unet

✅

resnet152

✅

resnet18

✅

resnet50

✅

resnext50_32x4d

✅

sam

✅

sam_fast

✅

shufflenet_v2_x1_0

✅

soft_actor_critic

✅

squeezenet1_1

✅

stable_diffusion_text_encoder

✅

stable_diffusion_unet

✅

timm_efficientdet

✅

timm_efficientnet

✅

timm_nfnet

✅

timm_regnet

✅

timm_resnest

✅

timm_vision_transformer

✅

timm_vision_transformer_large

✅

timm_vovnet

✅

torch_multimodal_clip

✅

tts_angular

✅

vgg16

✅

JAX Models#

Listed below are some JAX models that have been tested with PolyBlocks: they compile successfully and validate against JAX’s JIT (XLA). They come from Scenic, Keras Applications on Keras’s JAX backend, the Flax models of HuggingFace Transformers, MaxText, and the Gemma library of Google DeepMind, or are written in Flax after the TorchVision architecture. Inference is tested unless stated otherwise.

Model

Status

AlexNet (Flax)

✅

ConvNeXt Tiny (Keras Applications)

✅

DenseNet121 (Keras Applications)

✅

DistilBERT (HuggingFace Transformers)

✅

Gemma 4 31B (Gemma library; prefill and decode)

✅

Gemma 4 26B mixture of experts (MaxText; forward and SFT training step)*

✅

GPT-2 (HuggingFace Transformers)

✅

InceptionV3 (Keras Applications)

✅

MobileNetV3 Large (Keras Applications)

✅

ResNet (Scenic)

✅

SqueezeNet 1.0 (Flax)

✅

UNet (Scenic)

✅

VGG19 (Keras Applications)

✅

ViT (Scenic)

✅

Xception (Keras Applications)

✅

* Tested with a reduced number of its decoder and vision layers to fit a single GPU.

Contents