{
  "device": {
    "jax": "0.11.2",
    "backend": "cpu",
    "devices": [
      "cpu:0"
    ],
    "python": "3.14.8",
    "device_kind": "cpu",
    "gpu": "NVIDIA GeForce RTX 3090",
    "driver": "615.78.08",
    "gpu_memory_gb": 24.0,
    "torch": "2.14.1+cu130",
    "torch_cuda": "13.0"
  },
  "config": {
    "runs": [
      {
        "device": "cpu",
        "output": "docs/benchmarks/kernel_nn_cpu.json",
        "suites": [
          "transformer"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          128,
          512
        ],
        "ssm_tokens": [
          256,
          1024
        ],
        "unet_res": [
          32
        ],
        "moe_tokens": [
          1024
        ],
        "batch": 4,
        "model_dim": 256,
        "attention": null,
        "models": [
          "transformer",
          "diffusion"
        ],
        "layers": [
          12,
          24
        ],
        "scan_tokens": 512,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": false
      },
      {
        "device": "cpu",
        "output": "docs/benchmarks/kernel_nn_cpu.json",
        "suites": [
          "diffusion"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          128,
          512
        ],
        "ssm_tokens": [
          256,
          1024
        ],
        "unet_res": [
          32
        ],
        "moe_tokens": [
          1024
        ],
        "batch": 4,
        "model_dim": 256,
        "attention": null,
        "models": [
          "transformer",
          "diffusion"
        ],
        "layers": [
          12,
          24
        ],
        "scan_tokens": 512,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": false
      },
      {
        "device": "cpu",
        "output": "docs/benchmarks/kernel_nn_cpu.json",
        "suites": [
          "ssm",
          "unet"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          128,
          512
        ],
        "ssm_tokens": [
          256,
          1024
        ],
        "unet_res": [
          32
        ],
        "moe_tokens": [
          1024
        ],
        "batch": 4,
        "model_dim": 256,
        "attention": null,
        "models": [
          "transformer",
          "diffusion"
        ],
        "layers": [
          12,
          24
        ],
        "scan_tokens": 512,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": false
      },
      {
        "device": "cpu",
        "output": "docs/benchmarks/kernel_nn_cpu.json",
        "suites": [
          "train_loop",
          "remat"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          128,
          512
        ],
        "ssm_tokens": [
          256,
          1024
        ],
        "unet_res": [
          32
        ],
        "moe_tokens": [
          1024
        ],
        "batch": 4,
        "model_dim": 256,
        "attention": null,
        "models": [
          "transformer",
          "diffusion"
        ],
        "layers": [
          12,
          24
        ],
        "scan_tokens": 512,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": false
      },
      {
        "device": "cpu",
        "output": "docs/benchmarks/kernel_nn_cpu.json",
        "suites": [
          "scan_remat"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          128,
          512
        ],
        "ssm_tokens": [
          256,
          1024
        ],
        "unet_res": [
          32
        ],
        "moe_tokens": [
          1024
        ],
        "batch": 4,
        "model_dim": 256,
        "attention": null,
        "models": [
          "transformer"
        ],
        "layers": [
          12
        ],
        "scan_tokens": 512,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": false
      },
      {
        "device": "cpu",
        "output": "docs/benchmarks/kernel_nn_cpu.json",
        "suites": [
          "scan_remat"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          128,
          512
        ],
        "ssm_tokens": [
          256,
          1024
        ],
        "unet_res": [
          32
        ],
        "moe_tokens": [
          1024
        ],
        "batch": 4,
        "model_dim": 256,
        "attention": null,
        "models": [
          "transformer"
        ],
        "layers": [
          24
        ],
        "scan_tokens": 512,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": false
      }
    ]
  },
  "results": [
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.042302236004616134,
      "compile_s": 0.12686770500295097,
      "xla_temp_mb": 6.515625,
      "xla_arg_mb": 12.55078125,
      "xla_out_mb": 0.5,
      "xla_alias_mb": 0.0,
      "inner": 4,
      "median_ms": 10.265692750181188,
      "q1_ms": 10.095603499394201,
      "q3_ms": 10.306017999937467,
      "first_call_s": 0.04581898000469664
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.24866047900286503,
      "compile_s": 0.6083436830012943,
      "xla_temp_mb": 64.50201797485352,
      "xla_arg_mb": 36.652347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 52.88562999339774,
      "q1_ms": 52.28546949729207,
      "q3_ms": 53.54956600058358,
      "first_call_s": 0.15373318499769084
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.04106284199951915,
      "compile_s": 0.12959267100086436,
      "xla_temp_mb": 6.515625,
      "xla_arg_mb": 12.55078125,
      "xla_out_mb": 0.5,
      "xla_alias_mb": 0.0,
      "inner": 4,
      "median_ms": 10.354598749472643,
      "q1_ms": 9.974672249882133,
      "q3_ms": 10.521713375055697,
      "first_call_s": 0.044706535001751035
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.24454239699844038,
      "compile_s": 0.6030947120016208,
      "xla_temp_mb": 64.50201797485352,
      "xla_arg_mb": 36.652347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 53.15650300326524,
      "q1_ms": 52.71412600268377,
      "q3_ms": 53.5252900008345,
      "first_call_s": 0.14824721599870827
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.027292498001770582,
      "compile_s": 0.0859706470000674,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 10.54296875,
      "xla_out_mb": 0.5,
      "xla_alias_mb": 0.0,
      "inner": 6,
      "median_ms": 6.656125001124262,
      "q1_ms": 6.588386999889432,
      "q3_ms": 7.064226333871678,
      "first_call_s": 0.022364333999576047
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.20406308999372413,
      "compile_s": 0.4893276729999343,
      "xla_temp_mb": 33.007877349853516,
      "xla_arg_mb": 36.652347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 34.463755997421686,
      "q1_ms": 34.04353099904256,
      "q3_ms": 35.03796750010224,
      "first_call_s": 0.08616782099852571
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.040508372003387194,
      "compile_s": 0.1256850170029793,
      "xla_temp_mb": 98.0625,
      "xla_arg_mb": 14.05078125,
      "xla_out_mb": 2.0,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 80.0148589987657,
      "q1_ms": 79.18807000169181,
      "q3_ms": 81.53739949921146,
      "first_call_s": 0.1091556139945169
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.24517964199912967,
      "compile_s": 0.6113993740000296,
      "xla_temp_mb": 614.0078773498535,
      "xla_arg_mb": 38.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 341.3698549993569,
      "q1_ms": 337.34402850313927,
      "q3_ms": 342.7844269972411,
      "first_call_s": 0.3744469270022819
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.041415929001232143,
      "compile_s": 0.12968791100138333,
      "xla_temp_mb": 98.0625,
      "xla_arg_mb": 14.05078125,
      "xla_out_mb": 2.0,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 79.08410799427656,
      "q1_ms": 77.96905349459848,
      "q3_ms": 79.86439950036583,
      "first_call_s": 0.1076200780007639
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.24935322999954224,
      "compile_s": 0.613948808000714,
      "xla_temp_mb": 614.0078773498535,
      "xla_arg_mb": 38.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 347.40182800305774,
      "q1_ms": 343.40954949948355,
      "q3_ms": 348.4011840009771,
      "first_call_s": 0.36980213900096714
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.027743150996684562,
      "compile_s": 0.1038510669968673,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.04296875,
      "xla_out_mb": 2.0,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 18.752289499389008,
      "q1_ms": 18.585078500109375,
      "q3_ms": 19.57334550024825,
      "first_call_s": 0.0404436099997838
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.20941197000502143,
      "compile_s": 0.5339064660001895,
      "xla_temp_mb": 174.00787734985352,
      "xla_arg_mb": 38.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 107.85368100187043,
      "q1_ms": 107.17130899865879,
      "q3_ms": 110.03262799567892,
      "first_call_s": 0.14888327800144907
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.044820932002039626,
      "compile_s": 0.1708248420036398,
      "xla_temp_mb": 10.03173828125,
      "xla_arg_mb": 12.30078125,
      "xla_out_mb": 0.25,
      "xla_alias_mb": 0.0,
      "inner": 4,
      "median_ms": 10.742903250502422,
      "q1_ms": 10.690445125874248,
      "q3_ms": 10.952463249850553,
      "first_call_s": 0.04679507699620444
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.2792044669986353,
      "compile_s": 0.7269164129975252,
      "xla_temp_mb": 43.183658599853516,
      "xla_arg_mb": 36.402347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 91.39861700532492,
      "q1_ms": 87.62794799986295,
      "q3_ms": 92.91638849754236,
      "first_call_s": 0.17267757800436812
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.044465268998465035,
      "compile_s": 0.16514253099740017,
      "xla_temp_mb": 10.03173828125,
      "xla_arg_mb": 12.30078125,
      "xla_out_mb": 0.25,
      "xla_alias_mb": 0.0,
      "inner": 4,
      "median_ms": 10.881926000365638,
      "q1_ms": 10.787724125293607,
      "q3_ms": 11.129400624668051,
      "first_call_s": 0.04716143100085901
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.26964933999988716,
      "compile_s": 0.7061014620048809,
      "xla_temp_mb": 43.183658599853516,
      "xla_arg_mb": 36.402347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 92.61014999356121,
      "q1_ms": 91.41851000094903,
      "q3_ms": 93.0778090005333,
      "first_call_s": 0.1729841279957327
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.029523128003347665,
      "compile_s": 0.12139742499857675,
      "xla_temp_mb": 6.7626953125,
      "xla_arg_mb": 10.29296875,
      "xla_out_mb": 0.25,
      "xla_alias_mb": 0.0,
      "inner": 6,
      "median_ms": 6.721552334056469,
      "q1_ms": 6.540884583349301,
      "q3_ms": 6.792333666453487,
      "first_call_s": 0.02409052900475217
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.25531098100327654,
      "compile_s": 0.6204273050025222,
      "xla_temp_mb": 17.281314849853516,
      "xla_arg_mb": 36.402347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 73.17482799408026,
      "q1_ms": 72.43420649683685,
      "q3_ms": 73.23840800381731,
      "first_call_s": 0.12126311100291787
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.04631115200027125,
      "compile_s": 0.17395952399965608,
      "xla_temp_mb": 71.57861328125,
      "xla_arg_mb": 13.05078125,
      "xla_out_mb": 1.0,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 64.6655680029653,
      "q1_ms": 64.12904249737039,
      "q3_ms": 113.60412050271407,
      "first_call_s": 0.09432120699784718
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.5111106569966068,
      "compile_s": 1.7276251259972923,
      "xla_temp_mb": 522.0938148498535,
      "xla_arg_mb": 37.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 485.0231590025942,
      "q1_ms": 469.89709499757737,
      "q3_ms": 768.8846029996057,
      "first_call_s": 1.341870973999903
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.04669512699911138,
      "compile_s": 0.17855966599745443,
      "xla_temp_mb": 71.57861328125,
      "xla_arg_mb": 13.05078125,
      "xla_out_mb": 1.0,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 63.30018500011647,
      "q1_ms": 62.24471750101657,
      "q3_ms": 63.411315000848845,
      "first_call_s": 0.09493952599586919
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.2889255400004913,
      "compile_s": 0.7536816530046053,
      "xla_temp_mb": 522.0938148498535,
      "xla_arg_mb": 37.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 464.2962539946893,
      "q1_ms": 456.3550890015904,
      "q3_ms": 468.1445400019584,
      "first_call_s": 0.4965866280035698
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.029376079997746274,
      "compile_s": 0.12280594000185374,
      "xla_temp_mb": 14.2626953125,
      "xla_arg_mb": 11.04296875,
      "xla_out_mb": 1.0,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 19.72243999989587,
      "q1_ms": 19.6195352509676,
      "q3_ms": 20.185420999041526,
      "first_call_s": 0.03814393900393043
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.2303185049968306,
      "compile_s": 0.582545000994287,
      "xla_temp_mb": 151.10943984985352,
      "xla_arg_mb": 37.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 265.83769799617585,
      "q1_ms": 263.0691594968084,
      "q3_ms": 266.34737699714606,
      "first_call_s": 0.2906312170016463
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.07643944199662656,
      "compile_s": 0.21098566099681193,
      "xla_temp_mb": 7.1796875,
      "xla_arg_mb": 22.209152221679688,
      "xla_out_mb": 0.015625,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 12.13440633243105,
      "q1_ms": 12.030037334625376,
      "q3_ms": 12.24068733426975,
      "first_call_s": 0.05820738900365541
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4744995569999446,
      "compile_s": 1.1456431920014438,
      "xla_temp_mb": 61.131900787353516,
      "xla_arg_mb": 66.5961799621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 60.50128799688537,
      "q1_ms": 58.8984364985663,
      "q3_ms": 62.080736497591715,
      "first_call_s": 0.20405342600133736
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.08146663200022886,
      "compile_s": 0.21008341400010977,
      "xla_temp_mb": 7.1796875,
      "xla_arg_mb": 22.209152221679688,
      "xla_out_mb": 0.015625,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 11.846788998809643,
      "q1_ms": 11.829781666165218,
      "q3_ms": 11.949601665643666,
      "first_call_s": 0.055246998999791685
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.426264624002215,
      "compile_s": 1.0686417169999913,
      "xla_temp_mb": 61.131900787353516,
      "xla_arg_mb": 66.5961799621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 56.66854800074361,
      "q1_ms": 56.58159350059577,
      "q3_ms": 57.401935002417304,
      "first_call_s": 0.16623165399505524
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.058670083999459166,
      "compile_s": 0.16102324000530643,
      "xla_temp_mb": 4.546875,
      "xla_arg_mb": 20.201339721679688,
      "xla_out_mb": 0.015625,
      "xla_alias_mb": 0.0,
      "inner": 5,
      "median_ms": 8.267302399326582,
      "q1_ms": 8.209971000178484,
      "q3_ms": 8.394502100418322,
      "first_call_s": 0.031567039994115476
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.39246790500328643,
      "compile_s": 0.9427427960035857,
      "xla_temp_mb": 30.006900787353516,
      "xla_arg_mb": 66.5961799621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 34.73559400299564,
      "q1_ms": 34.36364050139673,
      "q3_ms": 34.806984000169905,
      "first_call_s": 0.10418693500105292
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.07318860100349411,
      "compile_s": 0.20819287699850975,
      "xla_temp_mb": 100.5546875,
      "xla_arg_mb": 22.256027221679688,
      "xla_out_mb": 0.0625,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 84.86040199932177,
      "q1_ms": 83.95191300223814,
      "q3_ms": 85.51355750023504,
      "first_call_s": 0.11975584799802164
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.42868050399556523,
      "compile_s": 1.0567375259997789,
      "xla_temp_mb": 636.0069007873535,
      "xla_arg_mb": 66.6430549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 347.7450170030352,
      "q1_ms": 346.25396100091166,
      "q3_ms": 348.5767060010403,
      "first_call_s": 0.4177741710009286
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.07517652199749136,
      "compile_s": 0.21128707099705935,
      "xla_temp_mb": 100.5546875,
      "xla_arg_mb": 22.256027221679688,
      "xla_out_mb": 0.0625,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 83.66079199913656,
      "q1_ms": 82.96835799774271,
      "q3_ms": 85.16725299705286,
      "first_call_s": 0.12161820199980866
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.42267630400601774,
      "compile_s": 1.0174058060001698,
      "xla_temp_mb": 636.0069007873535,
      "xla_arg_mb": 66.6430549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 348.7291190031101,
      "q1_ms": 344.29816649935674,
      "q3_ms": 351.1137099994812,
      "first_call_s": 0.4061615409955266
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.05766779799887445,
      "compile_s": 0.16366819399991073,
      "xla_temp_mb": 18.046875,
      "xla_arg_mb": 20.248214721679688,
      "xla_out_mb": 0.0625,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 20.218515001033666,
      "q1_ms": 20.089060251848423,
      "q3_ms": 20.59154850030609,
      "first_call_s": 0.04471764799382072
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.3774886270039133,
      "compile_s": 0.8986049149971223,
      "xla_temp_mb": 196.00690078735352,
      "xla_arg_mb": 66.6430549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 115.8516419964144,
      "q1_ms": 114.76994149779784,
      "q3_ms": 119.09502349953982,
      "first_call_s": 0.16830275100073777
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.07639294300315669,
      "compile_s": 0.2518782730039675,
      "xla_temp_mb": 10.520034790039062,
      "xla_arg_mb": 22.201332092285156,
      "xla_out_mb": 0.0078125,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 12.68161133339163,
      "q1_ms": 12.56819549962529,
      "q3_ms": 12.814440999742754,
      "first_call_s": 0.0558341889991425
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4421478130025207,
      "compile_s": 1.23945002900291,
      "xla_temp_mb": 34.138248443603516,
      "xla_arg_mb": 66.58835983276367,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 89.06790200126125,
      "q1_ms": 87.40599749580724,
      "q3_ms": 89.27927650074707,
      "first_call_s": 0.18351798599906033
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.07511015600175597,
      "compile_s": 0.2449603950008168,
      "xla_temp_mb": 10.520034790039062,
      "xla_arg_mb": 22.201332092285156,
      "xla_out_mb": 0.0078125,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 12.72494000052878,
      "q1_ms": 12.560081666985448,
      "q3_ms": 12.796846332397157,
      "first_call_s": 0.057696282005053945
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4464339809928788,
      "compile_s": 1.262631543999305,
      "xla_temp_mb": 34.138248443603516,
      "xla_arg_mb": 66.58835983276367,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 88.26380100072129,
      "q1_ms": 87.86025700101163,
      "q3_ms": 89.30421900367946,
      "first_call_s": 0.1847968049987685
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.060840771002403926,
      "compile_s": 0.21358558900101343,
      "xla_temp_mb": 6.8056793212890625,
      "xla_arg_mb": 20.193519592285156,
      "xla_out_mb": 0.0078125,
      "xla_alias_mb": 0.0,
      "inner": 6,
      "median_ms": 7.626356166535213,
      "q1_ms": 7.593473083034041,
      "q3_ms": 7.816938332931992,
      "first_call_s": 0.032670327003870625
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 128,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4052089800024987,
      "compile_s": 1.1133024279988604,
      "xla_temp_mb": 10.036685943603516,
      "xla_arg_mb": 66.58835983276367,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 74.93349200376542,
      "q1_ms": 74.35452649951912,
      "q3_ms": 75.83238750157761,
      "first_call_s": 0.13821099200140452
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.08447313999931794,
      "compile_s": 0.2697407360028592,
      "xla_temp_mb": 73.00440979003906,
      "xla_arg_mb": 22.224769592285156,
      "xla_out_mb": 0.03125,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 67.96568899881095,
      "q1_ms": 65.41758400271647,
      "q3_ms": 71.1426805028168,
      "first_call_s": 0.11029509100626456
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.45062222400156315,
      "compile_s": 1.2743807910010219,
      "xla_temp_mb": 532.6069984436035,
      "xla_arg_mb": 66.61179733276367,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 462.80278799531516,
      "q1_ms": 460.5610920007166,
      "q3_ms": 464.3491435017495,
      "first_call_s": 0.506574967999768
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.07678294100333005,
      "compile_s": 0.25575590900552925,
      "xla_temp_mb": 73.00440979003906,
      "xla_arg_mb": 22.224769592285156,
      "xla_out_mb": 0.03125,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 69.40130599832628,
      "q1_ms": 65.77737949919538,
      "q3_ms": 71.15118149886257,
      "first_call_s": 0.10675466199609218
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4494637829993735,
      "compile_s": 1.2662489380018087,
      "xla_temp_mb": 532.6069984436035,
      "xla_arg_mb": 66.61179733276367,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 458.0212720029522,
      "q1_ms": 455.31494149690843,
      "q3_ms": 459.98816300198087,
      "first_call_s": 0.5001043369993567
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.06067334399995161,
      "compile_s": 0.20402173599723028,
      "xla_temp_mb": 14.305679321289062,
      "xla_arg_mb": 20.216957092285156,
      "xla_out_mb": 0.03125,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 20.523387000139337,
      "q1_ms": 20.273536500099,
      "q3_ms": 20.90567024970369,
      "first_call_s": 0.04675630600104341
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4028703360017971,
      "compile_s": 1.100929560001532,
      "xla_temp_mb": 164.08356094360352,
      "xla_arg_mb": 66.61179733276367,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 274.9419859974296,
      "q1_ms": 273.36660050059436,
      "q3_ms": 275.73181499974453,
      "first_call_s": 0.3070171159997699
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 256,
      "batch": 4,
      "model_dim": 128,
      "op": "forward",
      "params": 930704,
      "finite": true,
      "trace_s": 0.09911006000038469,
      "compile_s": 0.4163090719957836,
      "xla_temp_mb": 309.703125,
      "xla_arg_mb": 3.61285400390625,
      "xla_out_mb": 0.0625,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 355.20447100134334,
      "q1_ms": 353.50768649732345,
      "q3_ms": 360.14509149754304,
      "first_call_s": 0.38096095099899685
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 256,
      "batch": 4,
      "model_dim": 128,
      "op": "train_step",
      "params": 930704,
      "finite": true,
      "trace_s": 0.6328132380003808,
      "compile_s": 1.8014982650056481,
      "xla_temp_mb": 2956.0157470703125,
      "xla_arg_mb": 10.713565826416016,
      "xla_out_mb": 10.653236389160156,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 1439.8029559961287,
      "q1_ms": 1436.7153354978655,
      "q3_ms": 1446.6426679973665,
      "first_call_s": 1.29624570300075
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 256,
      "batch": 4,
      "model_dim": 128,
      "op": "forward",
      "params": 865172,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 256,
      "batch": 4,
      "model_dim": 128,
      "op": "train_step",
      "params": 865172,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 256,
      "batch": 4,
      "model_dim": 128,
      "op": "forward",
      "params": 995728,
      "finite": true,
      "trace_s": 0.10344651200284716,
      "compile_s": 0.3325873700014199,
      "xla_temp_mb": 4.8837890625,
      "xla_arg_mb": 3.86090087890625,
      "xla_out_mb": 0.0625,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 12.395032332278788,
      "q1_ms": 12.141382499976316,
      "q3_ms": 12.638372664999526,
      "first_call_s": 0.043588108994299546
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 256,
      "batch": 4,
      "model_dim": 128,
      "op": "train_step",
      "params": 995728,
      "finite": true,
      "trace_s": 0.4989620139967883,
      "compile_s": 1.2915428769993014,
      "xla_temp_mb": 73.03229141235352,
      "xla_arg_mb": 11.457706451416016,
      "xla_out_mb": 11.397377014160156,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 53.2435519999126,
      "q1_ms": 52.060223497392144,
      "q3_ms": 54.238879503827775,
      "first_call_s": 0.12905087399849435
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 4,
      "model_dim": 128,
      "op": "forward",
      "params": 930704,
      "finite": true,
      "trace_s": 0.11018283400335349,
      "compile_s": 0.37266717200691346,
      "xla_temp_mb": 1320.734375,
      "xla_arg_mb": 3.80035400390625,
      "xla_out_mb": 0.25,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 1678.517312000622,
      "q1_ms": 1677.8247764996195,
      "q3_ms": 1688.8623400009237,
      "first_call_s": 1.6071617909983615
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 4,
      "model_dim": 128,
      "op": "train_step",
      "params": 930704,
      "finite": true,
      "trace_s": 0.7126191530041979,
      "compile_s": 2.0666067870042752,
      "xla_temp_mb": 12016.656314849854,
      "xla_arg_mb": 10.901065826416016,
      "xla_out_mb": 10.653236389160156,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 6658.037622997654,
      "q1_ms": 6639.457021999988,
      "q3_ms": 6687.913098001445,
      "first_call_s": 6.068756108004891
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 4,
      "model_dim": 128,
      "op": "forward",
      "params": 865172,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 4,
      "model_dim": 128,
      "op": "train_step",
      "params": 865172,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 4,
      "model_dim": 128,
      "op": "forward",
      "params": 995728,
      "finite": true,
      "trace_s": 0.11299680700176395,
      "compile_s": 0.3824293080033385,
      "xla_temp_mb": 60.4306640625,
      "xla_arg_mb": 4.04840087890625,
      "xla_out_mb": 0.25,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 47.60231200634735,
      "q1_ms": 47.3654980014544,
      "q3_ms": 48.11246149984072,
      "first_call_s": 0.08672893100447254
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 4,
      "model_dim": 128,
      "op": "train_step",
      "params": 995728,
      "finite": true,
      "trace_s": 0.5430057240009774,
      "compile_s": 1.4140389429958304,
      "xla_temp_mb": 326.1729164123535,
      "xla_arg_mb": 11.645206451416016,
      "xla_out_mb": 11.397377014160156,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 186.09004600148182,
      "q1_ms": 185.24433050333755,
      "q3_ms": 186.36759049695684,
      "first_call_s": 0.26971192899509333
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 32,
      "batch": 4,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "trace_s": 0.07266195299598621,
      "compile_s": 0.1990394800013746,
      "xla_temp_mb": 8.00048828125,
      "xla_arg_mb": 10.7431640625,
      "xla_out_mb": 0.046875,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 21.674668500054395,
      "q1_ms": 21.511058750547818,
      "q3_ms": 21.896332500546123,
      "first_call_s": 0.06084577199362684
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 32,
      "batch": 4,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "trace_s": 0.42927463699743384,
      "compile_s": 1.041989037003077,
      "xla_temp_mb": 94.05426406860352,
      "xla_arg_mb": 32.135746002197266,
      "xla_out_mb": 32.09076690673828,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 103.03344199928688,
      "q1_ms": 101.7368435022945,
      "q3_ms": 104.43764750016271,
      "first_call_s": 0.20248946599895135
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 32,
      "batch": 4,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "trace_s": 0.0845930209979997,
      "compile_s": 0.32913971899688477,
      "xla_temp_mb": 14.9415283203125,
      "xla_arg_mb": 10.7197265625,
      "xla_out_mb": 0.0234375,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 23.133442999096587,
      "q1_ms": 23.030543001368642,
      "q3_ms": 23.442419498678646,
      "first_call_s": 0.06300407399976393
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 32,
      "batch": 4,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "trace_s": 0.49452454000129364,
      "compile_s": 1.3026041299963254,
      "xla_temp_mb": 76.80145645141602,
      "xla_arg_mb": 32.112308502197266,
      "xla_out_mb": 32.09076690673828,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 97.90638300182763,
      "q1_ms": 96.82056949895923,
      "q3_ms": 101.73616750034853,
      "first_call_s": 0.20063811400177656
    },
    {
      "suite": "train_loop",
      "model": "MLP",
      "width": 64,
      "batch": 256,
      "op": "bare_jit",
      "per_step_ms": 0.1420983610808401,
      "trace_s": 0.04833957899973029,
      "compile_s": 0.1719937599991681,
      "xla_temp_mb": 0.4375648498535156,
      "xla_arg_mb": 0.08758926391601562,
      "xla_out_mb": 0.07212066650390625,
      "xla_alias_mb": 0.0,
      "inner": 108,
      "median_ms": 0.1420983610808401,
      "q1_ms": 0.14179918054179763,
      "q3_ms": 0.14540020371593226,
      "first_call_s": 0.01381899400439579
    },
    {
      "suite": "train_loop",
      "model": "MLP",
      "width": 64,
      "batch": 256,
      "op": "fit_array",
      "first_call_s": 0.3967388809978729,
      "second_call_s": 0.294713754003169,
      "per_step_ms": 0.1303612225001416
    },
    {
      "suite": "train_loop",
      "model": "MLP",
      "width": 64,
      "batch": 256,
      "op": "fit_stream",
      "first_call_s": 0.39906332699320046,
      "second_call_s": 0.2532259840008919,
      "per_step_ms": 0.22672854499433015
    },
    {
      "suite": "train_loop",
      "model": "MLP",
      "width": 512,
      "batch": 4096,
      "op": "bare_jit",
      "per_step_ms": 23.351061001449125,
      "trace_s": 0.04463308900449192,
      "compile_s": 0.1695163650001632,
      "xla_temp_mb": 56.000064849853516,
      "xla_arg_mb": 3.4494056701660156,
      "xla_out_mb": 3.1995620727539062,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 23.351061001449125,
      "q1_ms": 22.951345001274603,
      "q3_ms": 23.401819500577403,
      "first_call_s": 0.024254308998933993
    },
    {
      "suite": "train_loop",
      "model": "MLP",
      "width": 512,
      "batch": 4096,
      "op": "fit_array",
      "first_call_s": 3.09065154199925,
      "second_call_s": 2.9687805130015477,
      "per_step_ms": 13.395016582489916
    },
    {
      "suite": "train_loop",
      "model": "MLP",
      "width": 512,
      "batch": 4096,
      "op": "fit_stream",
      "first_call_s": 3.9087354990042513,
      "second_call_s": 3.7428961870027706,
      "per_step_ms": 17.744689147502868
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "variant": "plain",
      "trace_s": 0.3603025939955842,
      "compile_s": 0.727881298000284,
      "xla_temp_mb": 614.0078773498535,
      "xla_arg_mb": 38.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 334.94049000000814,
      "q1_ms": 330.75007050138083,
      "q3_ms": 337.78020549652865,
      "first_call_s": 0.3731752389940084
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "variant": "remat",
      "trace_s": 0.25560839099489385,
      "compile_s": 0.6508433509952738,
      "xla_temp_mb": 614.0078773498535,
      "xla_arg_mb": 38.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 390.69438300066395,
      "q1_ms": 386.73378199746367,
      "q3_ms": 393.2805040021776,
      "first_call_s": 0.4387152929994045
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 4,
      "model_dim": 256,
      "op": "train_step",
      "variant": "donate",
      "compile_s": 0.60290725199593,
      "xla_temp_mb": 619.5361976623535,
      "xla_arg_mb": 38.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 36.152347564697266,
      "inner": 1,
      "median_ms": 333.6537450013566,
      "q1_ms": 330.87735900335247,
      "q3_ms": 334.9303394970775
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 32,
      "batch": 4,
      "op": "train_step",
      "variant": "plain",
      "trace_s": 0.4247122300002957,
      "compile_s": 1.0756253029976506,
      "xla_temp_mb": 94.05426406860352,
      "xla_arg_mb": 32.135746002197266,
      "xla_out_mb": 32.09076690673828,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 101.4090099997702,
      "q1_ms": 101.29805549877346,
      "q3_ms": 102.7739820019633,
      "first_call_s": 0.19108472000516485
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 32,
      "batch": 4,
      "op": "train_step",
      "variant": "remat",
      "trace_s": 0.43453220200171927,
      "compile_s": 1.1384285590029322,
      "xla_temp_mb": 94.05426406860352,
      "xla_arg_mb": 32.135746002197266,
      "xla_out_mb": 32.09076690673828,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 114.4309580049594,
      "q1_ms": 112.051257998246,
      "q3_ms": 123.21525649895193,
      "first_call_s": 0.23289640699658776
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 32,
      "batch": 4,
      "op": "train_step",
      "variant": "donate",
      "compile_s": 1.0578125739993993,
      "xla_temp_mb": 106.82208633422852,
      "xla_arg_mb": 32.135746002197266,
      "xla_out_mb": 32.09076690673828,
      "xla_alias_mb": 32.088871002197266,
      "inner": 1,
      "median_ms": 104.41329200693872,
      "q1_ms": 101.45704100068542,
      "q3_ms": 105.65098850202048
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 4,
      "op": "train_step",
      "variant": "plain",
      "trace_s": 0.5411387610001839,
      "compile_s": 1.406590080994647,
      "xla_temp_mb": 326.1729164123535,
      "xla_arg_mb": 11.645206451416016,
      "xla_out_mb": 11.397377014160156,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 186.16145900159609,
      "q1_ms": 184.5967200024461,
      "q3_ms": 186.2370175003889,
      "first_call_s": 0.2685736019993783
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 4,
      "op": "train_step",
      "variant": "remat",
      "trace_s": 0.6314840190025279,
      "compile_s": 1.7384499740001047,
      "xla_temp_mb": 326.1729164123535,
      "xla_arg_mb": 11.645206451416016,
      "xla_out_mb": 11.397377014160156,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 212.8206580018741,
      "q1_ms": 212.33994499925757,
      "q3_ms": 213.90271500058589,
      "first_call_s": 0.3233974169997964
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 4,
      "op": "train_step",
      "variant": "donate",
      "compile_s": 1.3988324499950977,
      "xla_temp_mb": 327.83069229125977,
      "xla_arg_mb": 11.645206451416016,
      "xla_out_mb": 11.397377014160156,
      "xla_alias_mb": 11.395206451416016,
      "inner": 1,
      "median_ms": 184.97424400266027,
      "q1_ms": 181.72409150065505,
      "q3_ms": 185.84213000212912
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": false,
      "remat": "none",
      "op": "train_step",
      "params": 37828608,
      "finite": true,
      "trace_s": 0.685416789005103,
      "compile_s": 1.5632780230007484,
      "xla_temp_mb": 1034.0000648498535,
      "xla_arg_mb": 434.91406631469727,
      "xla_out_mb": 432.9184799194336,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 881.2876849988243,
      "q1_ms": 879.373564501293,
      "q3_ms": 884.8604414997681,
      "first_call_s": 1.0946938630004297
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": false,
      "remat": "full",
      "op": "train_step",
      "params": 37828608,
      "finite": true,
      "trace_s": 0.39591930000460707,
      "compile_s": 1.3986464600020554,
      "xla_temp_mb": 124.00006484985352,
      "xla_arg_mb": 434.91406631469727,
      "xla_out_mb": 432.9184799194336,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 761.3476749975234,
      "q1_ms": 758.5396345020854,
      "q3_ms": 764.7879064970766,
      "first_call_s": 1.1364693329960573
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": false,
      "remat": "dots_saveable",
      "op": "train_step",
      "params": 37828608,
      "finite": true,
      "trace_s": 0.4010849300029804,
      "compile_s": 1.4178080900019268,
      "xla_temp_mb": 382.0000648498535,
      "xla_arg_mb": 434.91406631469727,
      "xla_out_mb": 432.9184799194336,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 785.2346820000093,
      "q1_ms": 783.0609505035682,
      "q3_ms": 790.220301500085,
      "first_call_s": 1.0328704990024562
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": true,
      "remat": "none",
      "op": "train_step",
      "params": 37828608,
      "finite": true,
      "trace_s": 0.18588755399832735,
      "compile_s": 0.4850826030015014,
      "xla_temp_mb": 1101.0947189331055,
      "xla_arg_mb": 434.91406631469727,
      "xla_out_mb": 432.9144515991211,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 1034.5973549992777,
      "q1_ms": 1027.155274499819,
      "q3_ms": 1037.1863059990574,
      "first_call_s": 0.9589102099998854
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": true,
      "remat": "full",
      "op": "train_step",
      "params": 37828608,
      "finite": true,
      "trace_s": 0.15382445300201653,
      "compile_s": 0.43680883399792947,
      "xla_temp_mb": 40.00056457519531,
      "xla_arg_mb": 434.91406631469727,
      "xla_out_mb": 432.9144515991211,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 810.897848001332,
      "q1_ms": 801.1108379978396,
      "q3_ms": 814.1204784988076,
      "first_call_s": 0.8270636119996198
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "none",
      "op": "train_step",
      "params": 75657216,
      "finite": true,
      "trace_s": 1.21997123699839,
      "compile_s": 3.1947149629995693,
      "xla_temp_mb": 2042.0000648498535,
      "xla_arg_mb": 867.8281288146973,
      "xla_out_mb": 865.8369369506836,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 1755.3573669938487,
      "q1_ms": 1754.1555850002624,
      "q3_ms": 1769.829637996736,
      "first_call_s": 2.183044990997587
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "full",
      "op": "train_step",
      "params": 75657216,
      "finite": true,
      "trace_s": 0.6868047039970406,
      "compile_s": 3.07472985899949,
      "xla_temp_mb": 124.00006484985352,
      "xla_arg_mb": 867.8281288146973,
      "xla_out_mb": 865.8369369506836,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 1505.2953319973312,
      "q1_ms": 1496.6381285012176,
      "q3_ms": 1507.5610224994307,
      "first_call_s": 2.27545901099802
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "dots_saveable",
      "op": "train_step",
      "params": 75657216,
      "finite": true,
      "trace_s": 0.6869732439954532,
      "compile_s": 3.0610618999999133,
      "xla_temp_mb": 670.0000648498535,
      "xla_arg_mb": 867.8281288146973,
      "xla_out_mb": 865.8369369506836,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 1541.4966190000996,
      "q1_ms": 1534.767021501466,
      "q3_ms": 1543.692496496078,
      "first_call_s": 2.082432675997552
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "none",
      "op": "train_step",
      "params": 75657216,
      "finite": true,
      "trace_s": 0.18645886499871267,
      "compile_s": 0.5165025469977991,
      "xla_temp_mb": 2130.0634689331055,
      "xla_arg_mb": 867.8281288146973,
      "xla_out_mb": 865.8285140991211,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 2109.7919029998593,
      "q1_ms": 2104.209705004905,
      "q3_ms": 2115.4232640001283,
      "first_call_s": 1.9566624910003156
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 512,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "full",
      "op": "train_step",
      "params": 75657216,
      "finite": true,
      "trace_s": 0.15469887200015364,
      "compile_s": 0.44752499199967133,
      "xla_temp_mb": 26.000564575195312,
      "xla_arg_mb": 867.8281288146973,
      "xla_out_mb": 865.8285140991211,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 1611.4211630047066,
      "q1_ms": 1602.5111355011177,
      "q3_ms": 1627.2153999998409,
      "first_call_s": 1.5830461850055144
    }
  ]
}