{
  "device": {
    "jax": "0.11.2",
    "backend": "gpu",
    "devices": [
      "cuda:0"
    ],
    "python": "3.14.8",
    "device_kind": "NVIDIA GeForce RTX 3090",
    "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": "gpu",
        "output": "docs/benchmarks/kernel_nn_gpu.json",
        "suites": [
          "transformer"
        ],
        "dtypes": [
          "float32"
        ],
        "tokens": [
          256,
          1024,
          4096
        ],
        "ssm_tokens": [
          1024,
          4096
        ],
        "unet_res": [
          32,
          64
        ],
        "moe_tokens": [
          4096,
          32768
        ],
        "batch": 8,
        "model_dim": 256,
        "attention": null,
        "models": [
          "transformer",
          "diffusion"
        ],
        "layers": [
          12,
          24
        ],
        "scan_tokens": 4096,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": false
      },
      {
        "device": "gpu",
        "output": "docs/benchmarks/kernel_nn_gpu.json",
        "suites": [
          "transformer"
        ],
        "dtypes": [
          "bfloat16"
        ],
        "tokens": [
          256,
          1024,
          4096
        ],
        "ssm_tokens": [
          1024,
          4096
        ],
        "unet_res": [
          32,
          64
        ],
        "moe_tokens": [
          4096,
          32768
        ],
        "batch": 8,
        "model_dim": 256,
        "attention": null,
        "models": [
          "transformer",
          "diffusion"
        ],
        "layers": [
          12,
          24
        ],
        "scan_tokens": 4096,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": false
      },
      {
        "device": "gpu",
        "output": "docs/benchmarks/kernel_nn_gpu.json",
        "suites": [
          "scan_remat"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          256,
          1024,
          4096
        ],
        "ssm_tokens": [
          1024,
          4096
        ],
        "unet_res": [
          32,
          64
        ],
        "moe_tokens": [
          4096,
          32768
        ],
        "batch": 8,
        "model_dim": 256,
        "attention": [
          "probjax_default"
        ],
        "models": [
          "transformer"
        ],
        "layers": [
          24
        ],
        "scan_tokens": 4096,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": true
      },
      {
        "device": "gpu",
        "output": "docs/benchmarks/kernel_nn_gpu.json",
        "suites": [
          "scan_remat"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          256,
          1024,
          4096
        ],
        "ssm_tokens": [
          1024,
          4096
        ],
        "unet_res": [
          32,
          64
        ],
        "moe_tokens": [
          4096,
          32768
        ],
        "batch": 8,
        "model_dim": 256,
        "attention": [
          "probjax_flex"
        ],
        "models": [
          "transformer"
        ],
        "layers": [
          24
        ],
        "scan_tokens": 4096,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": true
      },
      {
        "device": "gpu",
        "output": "docs/benchmarks/kernel_nn_gpu.json",
        "suites": [
          "diffusion"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          256,
          1024,
          4096
        ],
        "ssm_tokens": [
          1024,
          4096
        ],
        "unet_res": [
          32,
          64
        ],
        "moe_tokens": [
          4096,
          32768
        ],
        "batch": 8,
        "model_dim": 256,
        "attention": [
          "probjax_default",
          "probjax_flex",
          "jax_cudnn",
          "no_attention"
        ],
        "models": [
          "transformer",
          "diffusion"
        ],
        "layers": [
          12,
          24
        ],
        "scan_tokens": 4096,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": false
      },
      {
        "device": "gpu",
        "output": "docs/benchmarks/kernel_nn_gpu.json",
        "suites": [
          "ssm",
          "unet"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          256,
          1024,
          4096
        ],
        "ssm_tokens": [
          1024,
          4096
        ],
        "unet_res": [
          32,
          64
        ],
        "moe_tokens": [
          4096,
          32768
        ],
        "batch": 8,
        "model_dim": 256,
        "attention": null,
        "models": [
          "transformer",
          "diffusion"
        ],
        "layers": [
          12,
          24
        ],
        "scan_tokens": 4096,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": false
      },
      {
        "device": "gpu",
        "output": "docs/benchmarks/kernel_nn_gpu.json",
        "suites": [
          "scan_remat"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          256,
          1024,
          4096
        ],
        "ssm_tokens": [
          1024,
          4096
        ],
        "unet_res": [
          32,
          64
        ],
        "moe_tokens": [
          4096,
          32768
        ],
        "batch": 8,
        "model_dim": 256,
        "attention": [
          "probjax_default"
        ],
        "models": [
          "transformer"
        ],
        "layers": [
          12
        ],
        "scan_tokens": 4096,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": true
      },
      {
        "device": "gpu",
        "output": "docs/benchmarks/kernel_nn_gpu.json",
        "suites": [
          "train_loop",
          "remat",
          "peaks"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          256,
          1024,
          4096
        ],
        "ssm_tokens": [
          1024,
          4096
        ],
        "unet_res": [
          32,
          64
        ],
        "moe_tokens": [
          4096,
          32768
        ],
        "batch": 8,
        "model_dim": 256,
        "attention": null,
        "models": [
          "transformer",
          "diffusion"
        ],
        "layers": [
          12,
          24
        ],
        "scan_tokens": 4096,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": false
      },
      {
        "device": "gpu",
        "output": "docs/benchmarks/kernel_nn_gpu.json",
        "suites": [
          "scan_remat"
        ],
        "dtypes": [
          "float32",
          "bfloat16"
        ],
        "tokens": [
          256,
          1024,
          4096
        ],
        "ssm_tokens": [
          1024,
          4096
        ],
        "unet_res": [
          32,
          64
        ],
        "moe_tokens": [
          4096,
          32768
        ],
        "batch": 8,
        "model_dim": 256,
        "attention": [
          "probjax_flex"
        ],
        "models": [
          "diffusion"
        ],
        "layers": [
          24
        ],
        "scan_tokens": 4096,
        "scan_batch": 2,
        "cells": [
          "mamba",
          "ssd",
          "lru"
        ],
        "repeats": 7,
        "append": true,
        "peak": true
      }
    ]
  },
  "results": [
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.046376865000638645,
      "compile_s": 4.251169338000182,
      "xla_temp_mb": 34.001487731933594,
      "xla_arg_mb": 14.05078125,
      "xla_out_mb": 2.0,
      "xla_alias_mb": 0.0,
      "inner": 36,
      "median_ms": 1.1840126666761028,
      "q1_ms": 1.1831276528432175,
      "q3_ms": 1.1864102777205212,
      "first_call_s": 0.0041486689951852895
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.25094626300415257,
      "compile_s": 7.057888120005373,
      "xla_temp_mb": 292.0237121582031,
      "xla_arg_mb": 38.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 9,
      "median_ms": 4.549813888540181,
      "q1_ms": 4.416288944614482,
      "q3_ms": 4.56216838897995,
      "first_call_s": 0.022518985999340657
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.063991563001764,
      "compile_s": 1.321260679003899,
      "xla_temp_mb": 12.001487731933594,
      "xla_arg_mb": 14.05078125,
      "xla_out_mb": 2.0,
      "xla_alias_mb": 0.0,
      "inner": 44,
      "median_ms": 0.8944463863206857,
      "q1_ms": 0.8802757273025981,
      "q3_ms": 0.9132485569022141,
      "first_call_s": 0.0037225229971227236
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.3029752830043435,
      "compile_s": 4.001623094001843,
      "xla_temp_mb": 164.02566528320312,
      "xla_arg_mb": 38.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 8,
      "median_ms": 4.109051874365832,
      "q1_ms": 4.040832813188899,
      "q3_ms": 4.150287375523476,
      "first_call_s": 0.014846400001260918
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.043450669007143006,
      "compile_s": 0.5083155080064898,
      "xla_temp_mb": 34.001487731933594,
      "xla_arg_mb": 14.05078125,
      "xla_out_mb": 2.0,
      "xla_alias_mb": 0.0,
      "inner": 37,
      "median_ms": 1.179501621762127,
      "q1_ms": 1.1551179324325513,
      "q3_ms": 1.1813294053403345,
      "first_call_s": 0.004102340993995313
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.3333736159984255,
      "compile_s": 1.8461822159952135,
      "xla_temp_mb": 292.0237121582031,
      "xla_arg_mb": 38.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 9,
      "median_ms": 4.380801888247434,
      "q1_ms": 4.3187193886195825,
      "q3_ms": 4.483088722028898,
      "first_call_s": 0.014309363003121689
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.02828093800053466,
      "compile_s": 0.7846381540002767,
      "xla_temp_mb": 10.001487731933594,
      "xla_arg_mb": 12.04296875,
      "xla_out_mb": 2.0,
      "xla_alias_mb": 0.0,
      "inner": 60,
      "median_ms": 0.6259977333684219,
      "q1_ms": 0.6212753833703271,
      "q3_ms": 0.6386833917228312,
      "first_call_s": 0.0025640509993536398
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.21819336500630016,
      "compile_s": 2.1214750740036834,
      "xla_temp_mb": 140.0234603881836,
      "xla_arg_mb": 38.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 15,
      "median_ms": 2.6671589332787944,
      "q1_ms": 2.6046857664672034,
      "q3_ms": 2.673685633393082,
      "first_call_s": 0.010160443001950625
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.047096165995753836,
      "compile_s": 3.5444228560008924,
      "xla_temp_mb": 520.0014877319336,
      "xla_arg_mb": 20.05078125,
      "xla_out_mb": 8.0,
      "xla_alias_mb": 0.0,
      "inner": 5,
      "median_ms": 8.294056399608962,
      "q1_ms": 8.285819799493765,
      "q3_ms": 8.30484489997616,
      "first_call_s": 0.034477481996873394
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.27403804299683543,
      "compile_s": 6.812458134001645,
      "xla_temp_mb": 2840.0271224975586,
      "xla_arg_mb": 44.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 27.518836002855096,
      "q1_ms": 26.69885949944728,
      "q3_ms": 29.892811999161495,
      "first_call_s": 0.09375654900213704
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.06444508100685198,
      "compile_s": 1.3782589210022707,
      "xla_temp_mb": 48.001487731933594,
      "xla_arg_mb": 20.05078125,
      "xla_out_mb": 8.0,
      "xla_alias_mb": 0.0,
      "inner": 11,
      "median_ms": 4.201908181791872,
      "q1_ms": 4.196615045326656,
      "q3_ms": 4.234518590899164,
      "first_call_s": 0.007056141002976801
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.31151646199577954,
      "compile_s": 3.5057898619998014,
      "xla_temp_mb": 644.0915756225586,
      "xla_arg_mb": 44.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 16.41792300142697,
      "q1_ms": 15.849028750380967,
      "q3_ms": 16.72794699879887,
      "first_call_s": 0.029867429999285378
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.049294436001218855,
      "compile_s": 0.6371047600041493,
      "xla_temp_mb": 520.0014877319336,
      "xla_arg_mb": 20.05078125,
      "xla_out_mb": 8.0,
      "xla_alias_mb": 0.0,
      "inner": 5,
      "median_ms": 8.264057400810998,
      "q1_ms": 8.235334500204772,
      "q3_ms": 8.316927700070664,
      "first_call_s": 0.010839638001925778
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.2583567339970614,
      "compile_s": 1.8992643789970316,
      "xla_temp_mb": 2840.0271224975586,
      "xla_arg_mb": 44.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 25.8513309963746,
      "q1_ms": 25.669457001640694,
      "q3_ms": 26.227212001685984,
      "first_call_s": 0.036313415999757126
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.027134367999678943,
      "compile_s": 0.8695743639982538,
      "xla_temp_mb": 40.001487731933594,
      "xla_arg_mb": 18.04296875,
      "xla_out_mb": 8.0,
      "xla_alias_mb": 0.0,
      "inner": 20,
      "median_ms": 2.32604890006769,
      "q1_ms": 2.308684800118499,
      "q3_ms": 2.3791378500391147,
      "first_call_s": 0.004554607003228739
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.2126552360059577,
      "compile_s": 2.479022628001985,
      "xla_temp_mb": 548.0896224975586,
      "xla_arg_mb": 44.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 6,
      "median_ms": 7.293992000389456,
      "q1_ms": 7.290307416648526,
      "q3_ms": 7.343994083688206,
      "first_call_s": 0.015480959998967592
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 8.03GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_forward']"
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 24.72GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.05805857800442027,
      "compile_s": 1.5625011490046745,
      "xla_temp_mb": 160.0034408569336,
      "xla_arg_mb": 44.05078125,
      "xla_out_mb": 32.0,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 28.972887004783843,
      "q1_ms": 28.73831299439189,
      "q3_ms": 29.046129999187542,
      "first_call_s": 0.03181165699788835
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.2826658170015435,
      "compile_s": 2.9313242270000046,
      "xla_temp_mb": 2569.2790756225586,
      "xla_arg_mb": 68.15234756469727,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 133.80351699743187,
      "q1_ms": 133.62933350072126,
      "q3_ms": 133.91278499693726,
      "first_call_s": 0.14086676300212275
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 8.03GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_forward']"
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 24.72GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.027723402999981772,
      "compile_s": 0.26841140699980315,
      "xla_temp_mb": 160.0034408569336,
      "xla_arg_mb": 42.04296875,
      "xla_out_mb": 32.0,
      "xla_alias_mb": 0.0,
      "inner": 5,
      "median_ms": 8.68432080023922,
      "q1_ms": 8.675428199785529,
      "q3_ms": 8.694590099912602,
      "first_call_s": 0.010599052999168634
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.20534870200208388,
      "compile_s": 2.2111422149973805,
      "xla_temp_mb": 2184.0271224975586,
      "xla_arg_mb": 68.15234756469727,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 26.77960700384574,
      "q1_ms": 26.722585498646367,
      "q3_ms": 26.844594998692628,
      "first_call_s": 0.03394257999752881
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.04753115399944363,
      "compile_s": 2.320398366005975,
      "xla_temp_mb": 28.002708435058594,
      "xla_arg_mb": 13.05078125,
      "xla_out_mb": 1.0,
      "xla_alias_mb": 0.0,
      "inner": 55,
      "median_ms": 0.736396799964661,
      "q1_ms": 0.728124981833389,
      "q3_ms": 0.7367510272491595,
      "first_call_s": 0.0040297739978996105
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.2760252660009428,
      "compile_s": 5.049283810003544,
      "xla_temp_mb": 160.02273559570312,
      "xla_arg_mb": 37.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 12,
      "median_ms": 3.067599249940637,
      "q1_ms": 3.0551561249012593,
      "q3_ms": 3.120436666904425,
      "first_call_s": 0.01957187699736096
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.06534654500137549,
      "compile_s": 0.8201553309991141,
      "xla_temp_mb": 9.002708435058594,
      "xla_arg_mb": 13.05078125,
      "xla_out_mb": 1.0,
      "xla_alias_mb": 0.0,
      "inner": 69,
      "median_ms": 0.5279566376238886,
      "q1_ms": 0.5245631811861962,
      "q3_ms": 0.5285126811730256,
      "first_call_s": 0.0030986320052761585
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.31425092199788196,
      "compile_s": 2.079451565994532,
      "xla_temp_mb": 61.024688720703125,
      "xla_arg_mb": 37.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 14,
      "median_ms": 2.443509143111961,
      "q1_ms": 2.428621678648467,
      "q3_ms": 2.6249967857958314,
      "first_call_s": 0.012070035001670476
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.046341770997969434,
      "compile_s": 0.5953125949963578,
      "xla_temp_mb": 28.002708435058594,
      "xla_arg_mb": 13.05078125,
      "xla_out_mb": 1.0,
      "xla_alias_mb": 0.0,
      "inner": 54,
      "median_ms": 0.7358979444823625,
      "q1_ms": 0.7234735092896262,
      "q3_ms": 0.7373551480769593,
      "first_call_s": 0.0035278340001241304
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.2697521969967056,
      "compile_s": 1.6599106819994631,
      "xla_temp_mb": 160.02273559570312,
      "xla_arg_mb": 37.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 13,
      "median_ms": 3.0573918461763805,
      "q1_ms": 3.053842346092167,
      "q3_ms": 3.0630891540987744,
      "first_call_s": 0.013354944996535778
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.03469580900127767,
      "compile_s": 0.554636302003928,
      "xla_temp_mb": 9.003684997558594,
      "xla_arg_mb": 13.05078125,
      "xla_out_mb": 1.0,
      "xla_alias_mb": 0.0,
      "inner": 65,
      "median_ms": 0.5766748154061273,
      "q1_ms": 0.5510960461553902,
      "q3_ms": 0.5788271230882787,
      "first_call_s": 0.010650499003531877
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.23555437000322854,
      "compile_s": 1.675628447002964,
      "xla_temp_mb": 59.087432861328125,
      "xla_arg_mb": 37.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 13,
      "median_ms": 2.9591104617933386,
      "q1_ms": 2.8579101535097626,
      "q3_ms": 2.9708353845657816,
      "first_call_s": 0.03171812599612167
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.02915900399966631,
      "compile_s": 0.3443606139990152,
      "xla_temp_mb": 9.002708435058594,
      "xla_arg_mb": 11.04296875,
      "xla_out_mb": 1.0,
      "xla_alias_mb": 0.0,
      "inner": 85,
      "median_ms": 0.39680645885014465,
      "q1_ms": 0.37429846467121564,
      "q3_ms": 0.3974879764855894,
      "first_call_s": 0.002505550000932999
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.2323964659954072,
      "compile_s": 1.835212425001373,
      "xla_temp_mb": 50.069358825683594,
      "xla_arg_mb": 37.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 22,
      "median_ms": 1.501590772849423,
      "q1_ms": 1.4688337955290494,
      "q3_ms": 1.57057013645747,
      "first_call_s": 0.010028652999608312
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.04471589500462869,
      "compile_s": 2.295278313999006,
      "xla_temp_mb": 401.0046615600586,
      "xla_arg_mb": 16.05078125,
      "xla_out_mb": 4.0,
      "xla_alias_mb": 0.0,
      "inner": 8,
      "median_ms": 5.563719374549692,
      "q1_ms": 5.546887625314412,
      "q3_ms": 5.572862937697209,
      "first_call_s": 0.008527452999260277
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.2714356130018132,
      "compile_s": 5.448487155001203,
      "xla_temp_mb": 2040.5271224975586,
      "xla_arg_mb": 40.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 18.384906001301715,
      "q1_ms": 18.197594250523252,
      "q3_ms": 18.85379074883531,
      "first_call_s": 0.07192040100198938
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.06415841699345037,
      "compile_s": 0.842685163996066,
      "xla_temp_mb": 48.004661560058594,
      "xla_arg_mb": 16.05078125,
      "xla_out_mb": 4.0,
      "xla_alias_mb": 0.0,
      "inner": 19,
      "median_ms": 2.238920894390168,
      "q1_ms": 2.2102161052125187,
      "q3_ms": 2.27264973670428,
      "first_call_s": 0.005282794001686852
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.32696682900132146,
      "compile_s": 2.1312256140008685,
      "xla_temp_mb": 288.5290756225586,
      "xla_arg_mb": 40.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 5,
      "median_ms": 7.9837477998808035,
      "q1_ms": 7.9581194993807,
      "q3_ms": 8.04052439998486,
      "first_call_s": 0.017876389996672515
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.04638032400544034,
      "compile_s": 0.6322670370063861,
      "xla_temp_mb": 401.0046615600586,
      "xla_arg_mb": 16.05078125,
      "xla_out_mb": 4.0,
      "xla_alias_mb": 0.0,
      "inner": 8,
      "median_ms": 5.644651249895105,
      "q1_ms": 5.607704437807115,
      "q3_ms": 5.759847937497398,
      "first_call_s": 0.009666707999713253
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.28531792599824257,
      "compile_s": 1.7900159299970255,
      "xla_temp_mb": 2040.5271224975586,
      "xla_arg_mb": 40.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 17.90116100164596,
      "q1_ms": 17.83089299897256,
      "q3_ms": 18.0157707491162,
      "first_call_s": 0.029080485001031775
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.03560162299982039,
      "compile_s": 0.42359460199804744,
      "xla_temp_mb": 48.005638122558594,
      "xla_arg_mb": 16.05078125,
      "xla_out_mb": 4.0,
      "xla_alias_mb": 0.0,
      "inner": 15,
      "median_ms": 2.777298200332249,
      "q1_ms": 2.7417225666188942,
      "q3_ms": 2.938075700027791,
      "first_call_s": 0.012857795001764316
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.24116107699956046,
      "compile_s": 1.3656374839993077,
      "xla_temp_mb": 288.5290832519531,
      "xla_arg_mb": 40.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 4,
      "median_ms": 9.830354749283288,
      "q1_ms": 9.801903625884734,
      "q3_ms": 9.852887500528595,
      "first_call_s": 0.04002884799410822
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.029567297002358828,
      "compile_s": 0.3676825970032951,
      "xla_temp_mb": 52.003684997558594,
      "xla_arg_mb": 14.04296875,
      "xla_out_mb": 4.0,
      "xla_alias_mb": 0.0,
      "inner": 30,
      "median_ms": 1.3829042332266301,
      "q1_ms": 1.3823307999700774,
      "q3_ms": 1.383474983231281,
      "first_call_s": 0.004484623997996096
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.22837342500133673,
      "compile_s": 1.699419085001864,
      "xla_temp_mb": 240.6833724975586,
      "xla_arg_mb": 40.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 10,
      "median_ms": 3.9777441998012364,
      "q1_ms": 3.959345949988346,
      "q3_ms": 4.006272750484641,
      "first_call_s": 0.012627670002984814
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.04505556800722843,
      "compile_s": 2.7122578020062065,
      "xla_temp_mb": 6210.007591247559,
      "xla_arg_mb": 28.05078125,
      "xla_out_mb": 16.0,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 66.38350599678233,
      "q1_ms": 66.26614999913727,
      "q3_ms": 66.42623299921979,
      "first_call_s": 0.06879844500508625
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.72GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.06861093099723803,
      "compile_s": 0.8503747719951207,
      "xla_temp_mb": 192.0085678100586,
      "xla_arg_mb": 28.05078125,
      "xla_out_mb": 16.0,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 14.844733666298756,
      "q1_ms": 14.676649333220363,
      "q3_ms": 14.95264299956034,
      "first_call_s": 0.017617288998735603
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.3360892659984529,
      "compile_s": 2.2925914309962536,
      "xla_temp_mb": 1150.1560287475586,
      "xla_arg_mb": 52.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 60.98118699446786,
      "q1_ms": 59.41270899711526,
      "q3_ms": 65.34928550172481,
      "first_call_s": 0.06918931499967584
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.045707811004831456,
      "compile_s": 0.5807964120031102,
      "xla_temp_mb": 6210.007591247559,
      "xla_arg_mb": 28.05078125,
      "xla_out_mb": 16.0,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 66.41866399877472,
      "q1_ms": 66.21679099771427,
      "q3_ms": 66.52460800250992,
      "first_call_s": 0.06922843900247244
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.72GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.037110345998371486,
      "compile_s": 0.23737879200052703,
      "xla_temp_mb": 192.0095443725586,
      "xla_arg_mb": 28.05078125,
      "xla_out_mb": 16.0,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 22.516903500218177,
      "q1_ms": 22.451351998824975,
      "q3_ms": 23.91041974988184,
      "first_call_s": 0.03406481500132941
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.23646805000316817,
      "compile_s": 1.2977697199967224,
      "xla_temp_mb": 1150.1560363769531,
      "xla_arg_mb": 52.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 87.74101099697873,
      "q1_ms": 87.03256200169562,
      "q3_ms": 87.75369099748787,
      "first_call_s": 0.11440213999594562
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.030182179994881153,
      "compile_s": 0.17437345600046683,
      "xla_temp_mb": 208.0085678100586,
      "xla_arg_mb": 26.04296875,
      "xla_out_mb": 16.0,
      "xla_alias_mb": 0.0,
      "inner": 9,
      "median_ms": 4.897783555659569,
      "q1_ms": 4.870067277806811,
      "q3_ms": 4.901163889169564,
      "first_call_s": 0.007314387999940664
    },
    {
      "suite": "transformer",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "trace_s": 0.23699415500595933,
      "compile_s": 1.7651038800031529,
      "xla_temp_mb": 957.0290756225586,
      "xla_arg_mb": 52.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 13.738791666886149,
      "q1_ms": 13.703710166737437,
      "q3_ms": 13.81368800017905,
      "first_call_s": 0.02161462200456299
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "none",
      "op": "train_step",
      "params": 75657216,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.03GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "none",
      "op": "train_step",
      "measure": "device_peak",
      "peak_error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.03GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "full",
      "op": "train_step",
      "params": 75657216,
      "finite": true,
      "trace_s": 0.7170142469985876,
      "compile_s": 8.6947597699982,
      "xla_temp_mb": 2808.181770324707,
      "xla_arg_mb": 873.8281288146973,
      "xla_out_mb": 865.8369369506836,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 491.87353599700145,
      "q1_ms": 491.6374089989404,
      "q3_ms": 492.23035799877835,
      "first_call_s": 0.6284079480028595
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "full",
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 4560.0126953125
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "dots_saveable",
      "op": "train_step",
      "params": 75657216,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 27.34GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "dots_saveable",
      "op": "train_step",
      "measure": "device_peak",
      "peak_error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 27.34GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "none",
      "op": "train_step",
      "params": 75657216,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 45.14GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "none",
      "op": "train_step",
      "measure": "device_peak",
      "peak_error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 45.14GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "full",
      "op": "train_step",
      "params": 75657216,
      "finite": true,
      "trace_s": 0.1659792650025338,
      "compile_s": 1.2200785770037328,
      "xla_temp_mb": 2769.316909790039,
      "xla_arg_mb": 873.8281288146973,
      "xla_out_mb": 865.8285140991211,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 488.7213760011946,
      "q1_ms": 488.2922845026769,
      "q3_ms": 488.9795960007177,
      "first_call_s": 0.4920758799999021
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "full",
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 4660.97802734375
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "none",
      "op": "train_step",
      "params": 75657216,
      "finite": true,
      "trace_s": 1.3431954469997436,
      "compile_s": 9.462703562996467,
      "xla_temp_mb": 3459.7380599975586,
      "xla_arg_mb": 873.8281288146973,
      "xla_out_mb": 865.8369369506836,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 211.7474080005195,
      "q1_ms": 210.54090299730888,
      "q3_ms": 212.56981850092416,
      "first_call_s": 0.3252296220016433
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "none",
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 5211.57470703125
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "full",
      "op": "train_step",
      "params": 75657216,
      "finite": true,
      "trace_s": 0.775807801001065,
      "compile_s": 7.315280126000289,
      "xla_temp_mb": 320.19373321533203,
      "xla_arg_mb": 873.8281288146973,
      "xla_out_mb": 865.8369369506836,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 259.9855590015068,
      "q1_ms": 259.4470360018022,
      "q3_ms": 260.3319600057148,
      "first_call_s": 0.3349842509996961
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "full",
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 2072.0302734375
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "dots_saveable",
      "op": "train_step",
      "params": 75657216,
      "finite": true,
      "trace_s": 0.7801792689933791,
      "compile_s": 7.131762135999452,
      "xla_temp_mb": 1808.1879119873047,
      "xla_arg_mb": 873.8281288146973,
      "xla_out_mb": 865.8369369506836,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 248.43002499983413,
      "q1_ms": 247.95101850031642,
      "q3_ms": 248.6412085017946,
      "first_call_s": 0.40993636500206776
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "dots_saveable",
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 3560.0244140625
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "none",
      "op": "train_step",
      "params": 75657216,
      "finite": true,
      "trace_s": 0.25408775600226363,
      "compile_s": 1.8288422190016718,
      "xla_temp_mb": 7372.533843994141,
      "xla_arg_mb": 873.8281288146973,
      "xla_out_mb": 865.8285140991211,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 248.17660600092495,
      "q1_ms": 247.97982750169467,
      "q3_ms": 248.87935950027895,
      "first_call_s": 0.3555212119972566
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "none",
      "op": "train_step",
      "measure": "device_peak",
      "peak_error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 7.20GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "full",
      "op": "train_step",
      "params": 75657216,
      "finite": true,
      "trace_s": 0.21477210900047794,
      "compile_s": 1.066051095003786,
      "xla_temp_mb": 234.52784729003906,
      "xla_arg_mb": 873.8281288146973,
      "xla_out_mb": 865.8285140991211,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 257.27963099780027,
      "q1_ms": 256.8760360009037,
      "q3_ms": 258.1301055033691,
      "first_call_s": 0.2643064710064209
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "full",
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 2126.189208984375
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.0719166939961724,
      "compile_s": 3.324162302997138,
      "xla_temp_mb": 36.345481872558594,
      "xla_arg_mb": 22.25604248046875,
      "xla_out_mb": 0.0625,
      "xla_alias_mb": 0.0,
      "inner": 34,
      "median_ms": 1.2672944411700453,
      "q1_ms": 1.2417384558134739,
      "q3_ms": 1.2712420588393476,
      "first_call_s": 0.0057476410001982
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4084665109985508,
      "compile_s": 6.655990824998298,
      "xla_temp_mb": 294.0605773925781,
      "xla_arg_mb": 66.64307022094727,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 8,
      "median_ms": 4.841258500164258,
      "q1_ms": 4.80284799959918,
      "q3_ms": 4.849567375003971,
      "first_call_s": 0.029194925002229866
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.08674265399895376,
      "compile_s": 1.036791202997847,
      "xla_temp_mb": 14.345481872558594,
      "xla_arg_mb": 22.25604248046875,
      "xla_out_mb": 0.0625,
      "xla_alias_mb": 0.0,
      "inner": 40,
      "median_ms": 0.9591162001015618,
      "q1_ms": 0.9585219499058439,
      "q3_ms": 0.9895934625092195,
      "first_call_s": 0.004831006001040805
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4379387719964143,
      "compile_s": 4.108626718996675,
      "xla_temp_mb": 166.06253051757812,
      "xla_arg_mb": 66.64307022094727,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 8,
      "median_ms": 4.1579076250855,
      "q1_ms": 4.088318937647273,
      "q3_ms": 4.379325624540797,
      "first_call_s": 0.018564411002444103
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.054595746005361434,
      "compile_s": 0.386138824003865,
      "xla_temp_mb": 12.095481872558594,
      "xla_arg_mb": 20.24822998046875,
      "xla_out_mb": 0.0625,
      "xla_alias_mb": 0.0,
      "inner": 51,
      "median_ms": 0.7038971568416694,
      "q1_ms": 0.693169401903801,
      "q3_ms": 0.7295506569361725,
      "first_call_s": 0.004088702000444755
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.3640629510045983,
      "compile_s": 2.7481416820010054,
      "xla_temp_mb": 142.06155395507812,
      "xla_arg_mb": 66.64307022094727,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 11,
      "median_ms": 3.113921000145968,
      "q1_ms": 2.961804318064506,
      "q3_ms": 3.1524977272387003,
      "first_call_s": 0.016104026995890308
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.07038962500519119,
      "compile_s": 3.095853579005052,
      "xla_temp_mb": 529.1113510131836,
      "xla_arg_mb": 22.44354248046875,
      "xla_out_mb": 0.25,
      "xla_alias_mb": 0.0,
      "inner": 5,
      "median_ms": 8.279308599594515,
      "q1_ms": 8.275472300010733,
      "q3_ms": 8.285098399210256,
      "first_call_s": 0.025680609003757127
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.3972327950032195,
      "compile_s": 6.295738321001409,
      "xla_temp_mb": 2904.0442123413086,
      "xla_arg_mb": 66.83057022094727,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 25.712098002259154,
      "q1_ms": 25.587258500308963,
      "q3_ms": 25.722133497765753,
      "first_call_s": 0.092048368001997
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.08446610900136875,
      "compile_s": 1.0507016120027401,
      "xla_temp_mb": 57.111351013183594,
      "xla_arg_mb": 22.44354248046875,
      "xla_out_mb": 0.25,
      "xla_alias_mb": 0.0,
      "inner": 10,
      "median_ms": 4.3036694994952995,
      "q1_ms": 4.290388249864918,
      "q3_ms": 4.304675299863447,
      "first_call_s": 0.00825344899931224
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4393652689977898,
      "compile_s": 3.8733854640013305,
      "xla_temp_mb": 716.0735092163086,
      "xla_arg_mb": 66.83057022094727,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 16.710073501599254,
      "q1_ms": 16.644757251924602,
      "q3_ms": 16.755386501245084,
      "first_call_s": 0.031245865000528283
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.054683672002283856,
      "compile_s": 0.3977318390025175,
      "xla_temp_mb": 48.095726013183594,
      "xla_arg_mb": 20.43572998046875,
      "xla_out_mb": 0.25,
      "xla_alias_mb": 0.0,
      "inner": 18,
      "median_ms": 2.3939910000384166,
      "q1_ms": 2.3908389446458185,
      "q3_ms": 2.4036980000043213,
      "first_call_s": 0.0059713129958254285
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.3633395940050832,
      "compile_s": 2.814181601999735,
      "xla_temp_mb": 620.0542221069336,
      "xla_arg_mb": 66.83057022094727,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 5,
      "median_ms": 8.128980199398939,
      "q1_ms": 8.06748390023131,
      "q3_ms": 8.201450100023067,
      "first_call_s": 0.021925146000285167
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 8.07GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_forward']"
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 24.78GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.08495302899973467,
      "compile_s": 1.058590902001015,
      "xla_temp_mb": 196.0039291381836,
      "xla_arg_mb": 23.19354248046875,
      "xla_out_mb": 1.0,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 29.25963899906492,
      "q1_ms": 29.067447499983246,
      "q3_ms": 29.449718498653965,
      "first_call_s": 0.034078222000971437
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4395523330022115,
      "compile_s": 3.3338070240060915,
      "xla_temp_mb": 2856.0783920288086,
      "xla_arg_mb": 67.58057022094727,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 136.23484699928667,
      "q1_ms": 136.0165190017142,
      "q3_ms": 136.67375649674796,
      "first_call_s": 0.14756041699729394
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.05503736100217793,
      "compile_s": 0.38406649399985326,
      "xla_temp_mb": 192.0039291381836,
      "xla_arg_mb": 21.18572998046875,
      "xla_out_mb": 1.0,
      "xla_alias_mb": 0.0,
      "inner": 5,
      "median_ms": 8.85613540012855,
      "q1_ms": 8.849557900248328,
      "q3_ms": 8.883711499947822,
      "first_call_s": 0.01232573100423906
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.3608606849957141,
      "compile_s": 2.9576003850015695,
      "xla_temp_mb": 2472.0551986694336,
      "xla_arg_mb": 67.58057022094727,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 28.957160997379106,
      "q1_ms": 28.857780998805538,
      "q3_ms": 29.135377502825577,
      "first_call_s": 0.041785118002735544
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.0761359239986632,
      "compile_s": 3.0906760779980686,
      "xla_temp_mb": 28.065208435058594,
      "xla_arg_mb": 22.224777221679688,
      "xla_out_mb": 0.03125,
      "xla_alias_mb": 0.0,
      "inner": 47,
      "median_ms": 0.8367218296699959,
      "q1_ms": 0.8196707552842588,
      "q3_ms": 0.8401297553141474,
      "first_call_s": 0.005293291003908962
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4393347010045545,
      "compile_s": 6.558155342005193,
      "xla_temp_mb": 166.07596588134766,
      "xla_arg_mb": 66.6118049621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 9,
      "median_ms": 3.6837146665978557,
      "q1_ms": 3.681587333428777,
      "q3_ms": 3.709535222393849,
      "first_call_s": 0.022047256999940146
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.0918821029990795,
      "compile_s": 0.9674681649994454,
      "xla_temp_mb": 10.065208435058594,
      "xla_arg_mb": 22.224777221679688,
      "xla_out_mb": 0.03125,
      "xla_alias_mb": 0.0,
      "inner": 59,
      "median_ms": 0.6088679830269431,
      "q1_ms": 0.5787260508780384,
      "q3_ms": 0.6236250339528758,
      "first_call_s": 0.006575787003384903
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4746370470020338,
      "compile_s": 2.9027569310055696,
      "xla_temp_mb": 65.04666900634766,
      "xla_arg_mb": 66.6118049621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 9,
      "median_ms": 3.4555163333101926,
      "q1_ms": 3.3713499999268812,
      "q3_ms": 3.5053318335687638,
      "first_call_s": 0.019644893996883184
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.0650766670005396,
      "compile_s": 0.5656034889980219,
      "xla_temp_mb": 10.066184997558594,
      "xla_arg_mb": 22.224777221679688,
      "xla_out_mb": 0.03125,
      "xla_alias_mb": 0.0,
      "inner": 54,
      "median_ms": 0.673181703691457,
      "q1_ms": 0.6290621388837248,
      "q3_ms": 0.6773243426491886,
      "first_call_s": 0.01235952499700943
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4112849320008536,
      "compile_s": 2.082166001004225,
      "xla_temp_mb": 67.07791900634766,
      "xla_arg_mb": 66.6118049621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 9,
      "median_ms": 3.6666714448882782,
      "q1_ms": 3.6047531664533183,
      "q3_ms": 3.6971349998364333,
      "first_call_s": 0.04206102899479447
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.06039709800097626,
      "compile_s": 0.505979435001791,
      "xla_temp_mb": 10.065208435058594,
      "xla_arg_mb": 20.216964721679688,
      "xla_out_mb": 0.03125,
      "xla_alias_mb": 0.0,
      "inner": 66,
      "median_ms": 0.49429651511617884,
      "q1_ms": 0.49314282577225554,
      "q3_ms": 0.4944208408921964,
      "first_call_s": 0.004408847999002319
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.39755520499602426,
      "compile_s": 2.5360304119967623,
      "xla_temp_mb": 62.060340881347656,
      "xla_arg_mb": 66.6118049621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 13,
      "median_ms": 2.71431192343768,
      "q1_ms": 2.7127544613344634,
      "q3_ms": 2.721150653647223,
      "first_call_s": 0.017817326996009797
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.07533061000140151,
      "compile_s": 2.733985348000715,
      "xla_temp_mb": 401.0049057006836,
      "xla_arg_mb": 22.318527221679688,
      "xla_out_mb": 0.125,
      "xla_alias_mb": 0.0,
      "inner": 8,
      "median_ms": 5.683103375304199,
      "q1_ms": 5.673913375176198,
      "q3_ms": 5.691973688044527,
      "first_call_s": 0.010275193002598826
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.43825617199763656,
      "compile_s": 6.263244606998342,
      "xla_temp_mb": 2104.0417709350586,
      "xla_arg_mb": 66.7055549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 18.860740998206893,
      "q1_ms": 18.793490999087226,
      "q3_ms": 18.90598350291839,
      "first_call_s": 0.03650837700115517
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.09143215299991425,
      "compile_s": 0.9757446880030329,
      "xla_temp_mb": 52.004905700683594,
      "xla_arg_mb": 22.318527221679688,
      "xla_out_mb": 0.125,
      "xla_alias_mb": 0.0,
      "inner": 16,
      "median_ms": 2.4076986246654997,
      "q1_ms": 2.406328218512499,
      "q3_ms": 2.4081860624391993,
      "first_call_s": 0.008602709996921476
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.48067291299958015,
      "compile_s": 2.7808199350038194,
      "xla_temp_mb": 356.0434799194336,
      "xla_arg_mb": 66.7055549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 5,
      "median_ms": 8.719749200099614,
      "q1_ms": 8.601065799302887,
      "q3_ms": 8.76128819945734,
      "first_call_s": 0.024452386001939885
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.06479922300059116,
      "compile_s": 0.5604783759990823,
      "xla_temp_mb": 52.005882263183594,
      "xla_arg_mb": 22.318527221679688,
      "xla_out_mb": 0.125,
      "xla_alias_mb": 0.0,
      "inner": 14,
      "median_ms": 3.061818642891012,
      "q1_ms": 3.0555030711444227,
      "q3_ms": 3.0623821429409355,
      "first_call_s": 0.01517076199525036
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4042863050053711,
      "compile_s": 2.112096498000028,
      "xla_temp_mb": 356.0437316894531,
      "xla_arg_mb": 66.7055549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 11.573234667594079,
      "q1_ms": 11.387735332391458,
      "q3_ms": 11.601329500990687,
      "first_call_s": 0.04897291099769063
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.06047773000318557,
      "compile_s": 0.506953356001759,
      "xla_temp_mb": 52.003929138183594,
      "xla_arg_mb": 20.310714721679688,
      "xla_out_mb": 0.125,
      "xla_alias_mb": 0.0,
      "inner": 29,
      "median_ms": 1.4724551377906303,
      "q1_ms": 1.4256276206589096,
      "q3_ms": 1.4765261206274114,
      "first_call_s": 0.005448845004139002
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.41332730399881257,
      "compile_s": 2.5649982780014398,
      "xla_temp_mb": 308.0464096069336,
      "xla_arg_mb": 66.7055549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 7,
      "median_ms": 4.731722000412576,
      "q1_ms": 4.716943999872975,
      "q3_ms": 4.854079642557606,
      "first_call_s": 0.021449826999742072
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.07586851799715078,
      "compile_s": 2.7553889689952484,
      "xla_temp_mb": 6209.679710388184,
      "xla_arg_mb": 22.693527221679688,
      "xla_out_mb": 0.5,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 66.99329199909698,
      "q1_ms": 66.6117509972537,
      "q3_ms": 67.11121600164915,
      "first_call_s": 0.07068508399970597
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.67GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.09121997200418264,
      "compile_s": 0.9902428080022219,
      "xla_temp_mb": 208.1338119506836,
      "xla_arg_mb": 22.693527221679688,
      "xla_out_mb": 0.5,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 15.078357665818961,
      "q1_ms": 15.074810666798538,
      "q3_ms": 15.265796166204382,
      "first_call_s": 0.021839004999492317
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.4773119259989471,
      "compile_s": 2.838692017001449,
      "xla_temp_mb": 1422.0825424194336,
      "xla_arg_mb": 67.0805549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 62.83559000439709,
      "q1_ms": 61.402649502269924,
      "q3_ms": 63.05536500076414,
      "first_call_s": 0.07530918100383133
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.06613085199933266,
      "compile_s": 0.4205626979965018,
      "xla_temp_mb": 208.1347885131836,
      "xla_arg_mb": 22.693527221679688,
      "xla_out_mb": 0.5,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 22.877141997014405,
      "q1_ms": 22.77456074989459,
      "q3_ms": 24.30294850091741,
      "first_call_s": 0.0368762029975187
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.40789352700085146,
      "compile_s": 2.10508200199547,
      "xla_temp_mb": 1422.1606750488281,
      "xla_arg_mb": 67.0805549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 88.98813899577362,
      "q1_ms": 88.55222200145363,
      "q3_ms": 89.17836000182433,
      "first_call_s": 0.1324011469987454
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.06001158799335826,
      "compile_s": 0.37268571999447886,
      "xla_temp_mb": 208.1338119506836,
      "xla_arg_mb": 20.685714721679688,
      "xla_out_mb": 0.5,
      "xla_alias_mb": 0.0,
      "inner": 8,
      "median_ms": 5.1553113753470825,
      "q1_ms": 5.125464187585749,
      "q3_ms": 5.161546875115164,
      "first_call_s": 0.009217452003213111
    },
    {
      "suite": "diffusion",
      "attention": "no_attention",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "trace_s": 0.3972750850007287,
      "compile_s": 2.5842505100008566,
      "xla_temp_mb": 1230.0642318725586,
      "xla_arg_mb": 67.0805549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 15.816877666414559,
      "q1_ms": 15.794285499093045,
      "q3_ms": 15.83105983202889,
      "first_call_s": 0.031317011002101935
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 930704,
      "finite": true,
      "trace_s": 0.10697733300185064,
      "compile_s": 2.202744078997057,
      "xla_temp_mb": 468.06910705566406,
      "xla_arg_mb": 4.05035400390625,
      "xla_out_mb": 0.5,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 14.500905999739189,
      "q1_ms": 14.49413383306819,
      "q3_ms": 14.518720333095796,
      "first_call_s": 0.026630238004145212
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 930704,
      "finite": true,
      "trace_s": 0.697152549000748,
      "compile_s": 7.328764285004581,
      "xla_temp_mb": 5646.244178771973,
      "xla_arg_mb": 11.151065826416016,
      "xla_out_mb": 10.653236389160156,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 138.53852600004757,
      "q1_ms": 138.18411149986787,
      "q3_ms": 138.88392049921094,
      "first_call_s": 0.25544305599760264
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 865172,
      "error": "RuntimeError: ssd does not support this GPU Pallas configuration. backend='triton', seq_len=1024, num_groups=1, num_heads=1, dk=128, dv=128. Reference fallback on GPU is disabled."
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 865172,
      "error": "RuntimeError: ssd does not support this GPU Pallas configuration. backend='triton', seq_len=1024, num_groups=1, num_heads=1, dk=128, dv=128. Reference fallback on GPU is disabled."
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 995728,
      "finite": true,
      "trace_s": 0.10577234399534063,
      "compile_s": 1.4481309519978822,
      "xla_temp_mb": 32.00367736816406,
      "xla_arg_mb": 4.29840087890625,
      "xla_out_mb": 0.5,
      "xla_alias_mb": 0.0,
      "inner": 27,
      "median_ms": 1.61366425918122,
      "q1_ms": 1.606059462776304,
      "q3_ms": 1.6141314074586801,
      "first_call_s": 0.005840917001478374
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 995728,
      "finite": true,
      "trace_s": 0.5371652850008104,
      "compile_s": 7.308605087004253,
      "xla_temp_mb": 444.1372528076172,
      "xla_arg_mb": 11.895206451416016,
      "xla_out_mb": 11.397377014160156,
      "xla_alias_mb": 0.0,
      "inner": 7,
      "median_ms": 5.371187428993705,
      "q1_ms": 5.317043570747566,
      "q3_ms": 5.373506786002379,
      "first_call_s": 0.023124668005038984
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 930704,
      "finite": true,
      "trace_s": 0.11565162899933057,
      "compile_s": 2.3815510679996805,
      "xla_temp_mb": 2640.072036743164,
      "xla_arg_mb": 5.55035400390625,
      "xla_out_mb": 2.0,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 65.11574999603909,
      "q1_ms": 65.0189429979946,
      "q3_ms": 65.23034550264128,
      "first_call_s": 0.06972034499631263
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 930704,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 20.57GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 865172,
      "error": "RuntimeError: ssd does not support this GPU Pallas configuration. backend='triton', seq_len=4096, num_groups=1, num_heads=1, dk=128, dv=128. Reference fallback on GPU is disabled."
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 865172,
      "error": "RuntimeError: ssd does not support this GPU Pallas configuration. backend='triton', seq_len=4096, num_groups=1, num_heads=1, dk=128, dv=128. Reference fallback on GPU is disabled."
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 995728,
      "finite": true,
      "trace_s": 0.11927481199381873,
      "compile_s": 1.5353008219972253,
      "xla_temp_mb": 128.00465393066406,
      "xla_arg_mb": 5.79840087890625,
      "xla_out_mb": 2.0,
      "xla_alias_mb": 0.0,
      "inner": 7,
      "median_ms": 6.187279571479719,
      "q1_ms": 6.184982428489352,
      "q3_ms": 6.197850928271821,
      "first_call_s": 0.010775599999760743
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 995728,
      "finite": true,
      "trace_s": 0.5769222890012315,
      "compile_s": 6.745953607998672,
      "xla_temp_mb": 1862.4204635620117,
      "xla_arg_mb": 13.395206451416016,
      "xla_out_mb": 11.397377014160156,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 18.15778250238509,
      "q1_ms": 18.10756050144846,
      "q3_ms": 18.167070749768754,
      "first_call_s": 0.03401280099933501
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 32,
      "batch": 8,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "trace_s": 0.1933520189995761,
      "compile_s": 1.651257778001309,
      "xla_temp_mb": 23.511245727539062,
      "xla_arg_mb": 10.7900390625,
      "xla_out_mb": 0.09375,
      "xla_alias_mb": 0.0,
      "inner": 26,
      "median_ms": 1.6643443845028882,
      "q1_ms": 1.6635305768991553,
      "q3_ms": 1.6653490000373075,
      "first_call_s": 0.00979483500123024
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 32,
      "batch": 8,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "trace_s": 0.42479661499965005,
      "compile_s": 3.7720606360016973,
      "xla_temp_mb": 113.61966705322266,
      "xla_arg_mb": 32.182621002197266,
      "xla_out_mb": 32.09076690673828,
      "xla_alias_mb": 0.0,
      "inner": 7,
      "median_ms": 5.448039857063642,
      "q1_ms": 5.4383973568162345,
      "q3_ms": 5.45278892864839,
      "first_call_s": 0.029425604996504262
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 64,
      "batch": 8,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "trace_s": 0.07276808100141352,
      "compile_s": 1.330261718001566,
      "xla_temp_mb": 83.01124572753906,
      "xla_arg_mb": 11.0712890625,
      "xla_out_mb": 0.375,
      "xla_alias_mb": 0.0,
      "inner": 10,
      "median_ms": 4.701519200170878,
      "q1_ms": 4.696477199831861,
      "q3_ms": 4.711945100279991,
      "first_call_s": 0.010507083003176376
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 64,
      "batch": 8,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "trace_s": 0.4232518839999102,
      "compile_s": 4.125170225001057,
      "xla_temp_mb": 436.28055572509766,
      "xla_arg_mb": 32.463871002197266,
      "xla_out_mb": 32.09076690673828,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 15.203756001331689,
      "q1_ms": 15.177168666317204,
      "q3_ms": 15.241963000638254,
      "first_call_s": 0.03095527800178388
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 32,
      "batch": 8,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "trace_s": 0.0834688390023075,
      "compile_s": 1.4380813440002385,
      "xla_temp_mb": 11.761245727539062,
      "xla_arg_mb": 10.7431640625,
      "xla_out_mb": 0.046875,
      "xla_alias_mb": 0.0,
      "inner": 36,
      "median_ms": 1.1552858888737523,
      "q1_ms": 1.1495773194635857,
      "q3_ms": 1.155399583391651,
      "first_call_s": 0.010045078997791279
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 32,
      "batch": 8,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "trace_s": 0.47448503299528966,
      "compile_s": 10.793616147995635,
      "xla_temp_mb": 72.71244049072266,
      "xla_arg_mb": 32.135746002197266,
      "xla_out_mb": 32.09076690673828,
      "xla_alias_mb": 0.0,
      "inner": 9,
      "median_ms": 3.9903248883395763,
      "q1_ms": 3.9876998889667448,
      "q3_ms": 4.00003000039659,
      "first_call_s": 0.03258724499755772
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 64,
      "batch": 8,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "trace_s": 0.08268636999855516,
      "compile_s": 1.2364772569999332,
      "xla_temp_mb": 44.76222229003906,
      "xla_arg_mb": 10.8837890625,
      "xla_out_mb": 0.1875,
      "xla_alias_mb": 0.0,
      "inner": 16,
      "median_ms": 2.898571687637741,
      "q1_ms": 2.8898050311454426,
      "q3_ms": 2.9016117498485983,
      "first_call_s": 0.010071639000670984
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 64,
      "batch": 8,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "trace_s": 0.4937567430024501,
      "compile_s": 13.316599375997612,
      "xla_temp_mb": 225.49100494384766,
      "xla_arg_mb": 32.276371002197266,
      "xla_out_mb": 32.09076690673828,
      "xla_alias_mb": 0.0,
      "inner": 4,
      "median_ms": 9.496067750660586,
      "q1_ms": 9.483307499976945,
      "q3_ms": 9.498061624981347,
      "first_call_s": 0.03347545499855187
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": false,
      "remat": "none",
      "op": "train_step",
      "params": 37828608,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 15.61GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": false,
      "remat": "none",
      "op": "train_step",
      "measure": "device_peak",
      "peak_error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 15.61GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": false,
      "remat": "full",
      "op": "train_step",
      "params": 37828608,
      "finite": true,
      "trace_s": 0.41471265200380003,
      "compile_s": 4.099621022003703,
      "xla_temp_mb": 2712.088264465332,
      "xla_arg_mb": 440.91406631469727,
      "xla_out_mb": 432.9184799194336,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 245.38969199784333,
      "q1_ms": 244.09522249698057,
      "q3_ms": 246.2446879981144,
      "first_call_s": 0.3384691150058643
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": false,
      "remat": "full",
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 3599.009033203125
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": false,
      "remat": "dots_saveable",
      "op": "train_step",
      "params": 37828608,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 14.41GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": false,
      "remat": "dots_saveable",
      "op": "train_step",
      "measure": "device_peak",
      "peak_error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 14.41GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": true,
      "remat": "none",
      "op": "train_step",
      "params": 37828608,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 23.59GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": true,
      "remat": "none",
      "op": "train_step",
      "measure": "device_peak",
      "peak_error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 23.59GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": true,
      "remat": "full",
      "op": "train_step",
      "params": 37828608,
      "finite": true,
      "trace_s": 0.16723991199978627,
      "compile_s": 1.2216550140001345,
      "xla_temp_mb": 2673.480728149414,
      "xla_arg_mb": 440.91406631469727,
      "xla_out_mb": 432.9144515991211,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 244.8797969991574,
      "q1_ms": 244.75612699825433,
      "q3_ms": 245.0621774987667,
      "first_call_s": 0.24908667400450213
    },
    {
      "suite": "scan_remat",
      "model": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 12,
      "scan": true,
      "remat": "full",
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 3619.314453125
    },
    {
      "suite": "train_loop",
      "model": "MLP",
      "width": 64,
      "batch": 256,
      "op": "bare_jit",
      "per_step_ms": 0.15892503502982436,
      "trace_s": 0.0487997579984949,
      "compile_s": 1.1816286989997025,
      "xla_temp_mb": 0.25197601318359375,
      "xla_arg_mb": 0.08758926391601562,
      "xla_out_mb": 0.07212066650390625,
      "xla_alias_mb": 0.0,
      "inner": 200,
      "median_ms": 0.15892503502982436,
      "q1_ms": 0.15861915751884226,
      "q3_ms": 0.1600043525104411,
      "first_call_s": 0.001717598999675829
    },
    {
      "suite": "train_loop",
      "model": "MLP",
      "width": 64,
      "batch": 256,
      "op": "fit_array",
      "first_call_s": 0.604801157998736,
      "second_call_s": 0.3571169370043208,
      "per_step_ms": 0.20496536499194917
    },
    {
      "suite": "train_loop",
      "model": "MLP",
      "width": 64,
      "batch": 256,
      "op": "fit_stream",
      "first_call_s": 0.5565333060003468,
      "second_call_s": 0.3750505969946971,
      "per_step_ms": 0.37763262751468574
    },
    {
      "suite": "train_loop",
      "model": "MLP",
      "width": 512,
      "batch": 4096,
      "op": "bare_jit",
      "per_step_ms": 0.4594273249495018,
      "trace_s": 0.0462861839987454,
      "compile_s": 2.014602153998567,
      "xla_temp_mb": 32.001976013183594,
      "xla_arg_mb": 3.4494056701660156,
      "xla_out_mb": 3.1995620727539062,
      "xla_alias_mb": 0.0,
      "inner": 80,
      "median_ms": 0.4594273249495018,
      "q1_ms": 0.4593618749822781,
      "q3_ms": 0.4704179187228874,
      "first_call_s": 0.0019375249976292253
    },
    {
      "suite": "train_loop",
      "model": "MLP",
      "width": 512,
      "batch": 4096,
      "op": "fit_array",
      "first_call_s": 0.8489061099971877,
      "second_call_s": 0.6505573799950071,
      "per_step_ms": 0.532062987504105
    },
    {
      "suite": "train_loop",
      "model": "MLP",
      "width": 512,
      "batch": 4096,
      "op": "fit_stream",
      "first_call_s": 0.8168241109960945,
      "second_call_s": 0.6179635139997117,
      "per_step_ms": 0.6356029937523999
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "variant": "plain",
      "trace_s": 0.39239443100086646,
      "compile_s": 5.578387718000158,
      "xla_temp_mb": 2040.5271224975586,
      "xla_arg_mb": 40.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 17.668898002739297,
      "q1_ms": 17.639640249399235,
      "q3_ms": 17.673006750555942,
      "first_call_s": 0.05574469400016824
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "variant": "remat",
      "trace_s": 0.28513656700670253,
      "compile_s": 1.9427239670039853,
      "xla_temp_mb": 2045.655288696289,
      "xla_arg_mb": 40.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 23.02704249814269,
      "q1_ms": 23.008187248706236,
      "q3_ms": 23.22536274914455,
      "first_call_s": 0.03643101500347257
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "variant": "donate",
      "compile_s": 1.6315844290002133,
      "xla_temp_mb": 2046.6858139038086,
      "xla_arg_mb": 40.152347564697266,
      "xla_out_mb": 36.153831481933594,
      "xla_alias_mb": 36.152347564697266,
      "inner": 2,
      "median_ms": 17.4191804981092,
      "q1_ms": 17.37588400101231,
      "q3_ms": 17.487149249063805
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "variant": "plain",
      "trace_s": 0.47682618100225227,
      "compile_s": 5.796727686996746,
      "xla_temp_mb": 1422.0825424194336,
      "xla_arg_mb": 67.0805549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 62.31057900004089,
      "q1_ms": 61.49284999628435,
      "q3_ms": 63.18817249848507,
      "first_call_s": 0.07337273400480626
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "variant": "remat",
      "trace_s": 0.503686904994538,
      "compile_s": 2.6410391539975535,
      "xla_temp_mb": 1422.2454528808594,
      "xla_arg_mb": 67.0805549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 76.45299100113334,
      "q1_ms": 76.22347150027053,
      "q3_ms": 77.70342650474049,
      "first_call_s": 0.09368227799859596
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "variant": "donate",
      "compile_s": 2.1078961870007333,
      "xla_temp_mb": 1436.4184799194336,
      "xla_arg_mb": 67.0805549621582,
      "xla_out_mb": 66.58296203613281,
      "xla_alias_mb": 66.58053970336914,
      "inner": 1,
      "median_ms": 60.58648099860875,
      "q1_ms": 60.29520949959988,
      "q3_ms": 61.0084990003088
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 64,
      "batch": 8,
      "op": "train_step",
      "variant": "plain",
      "trace_s": 0.4784173610023572,
      "compile_s": 14.36543505499867,
      "xla_temp_mb": 225.49100494384766,
      "xla_arg_mb": 32.276371002197266,
      "xla_out_mb": 32.09076690673828,
      "xla_alias_mb": 0.0,
      "inner": 4,
      "median_ms": 9.56414675056294,
      "q1_ms": 9.532637375741615,
      "q3_ms": 9.581911624991335,
      "first_call_s": 0.029982378000568133
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 64,
      "batch": 8,
      "op": "train_step",
      "variant": "remat",
      "trace_s": 0.4853751100017689,
      "compile_s": 2.5327361629970255,
      "xla_temp_mb": 231.17422485351562,
      "xla_arg_mb": 32.276371002197266,
      "xla_out_mb": 32.09076690673828,
      "xla_alias_mb": 0.0,
      "inner": 3,
      "median_ms": 13.255032999344015,
      "q1_ms": 12.554420000621274,
      "q3_ms": 13.35738683337695,
      "first_call_s": 0.040224399002909195
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 64,
      "batch": 8,
      "op": "train_step",
      "variant": "donate",
      "compile_s": 2.268382655995083,
      "xla_temp_mb": 238.08646392822266,
      "xla_arg_mb": 32.276371002197266,
      "xla_out_mb": 32.09076690673828,
      "xla_alias_mb": 32.088871002197266,
      "inner": 4,
      "median_ms": 9.408346999407513,
      "q1_ms": 9.354379874821461,
      "q3_ms": 9.951966500011622
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "op": "train_step",
      "variant": "plain",
      "trace_s": 0.5848000190017046,
      "compile_s": 8.400441275000048,
      "xla_temp_mb": 1862.4204635620117,
      "xla_arg_mb": 13.395206451416016,
      "xla_out_mb": 11.397377014160156,
      "xla_alias_mb": 0.0,
      "inner": 2,
      "median_ms": 18.371501999354223,
      "q1_ms": 18.156412999815075,
      "q3_ms": 18.48912700006622,
      "first_call_s": 0.03461919799883617
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "op": "train_step",
      "variant": "remat",
      "trace_s": 0.6911827669973718,
      "compile_s": 6.403347073995974,
      "xla_temp_mb": 1894.426513671875,
      "xla_arg_mb": 13.395206451416016,
      "xla_out_mb": 11.397377014160156,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 24.692222999874502,
      "q1_ms": 24.443172998871887,
      "q3_ms": 24.743548998230835,
      "first_call_s": 0.044983813000726514
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "op": "train_step",
      "variant": "donate",
      "compile_s": 3.3932634340017103,
      "xla_temp_mb": 1866.3086471557617,
      "xla_arg_mb": 13.395206451416016,
      "xla_out_mb": 11.397377014160156,
      "xla_alias_mb": 11.395206451416016,
      "inner": 2,
      "median_ms": 17.87197249723249,
      "q1_ms": 17.728611250277027,
      "q3_ms": 17.905039252582355
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 653.999755859375
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 2126.154296875
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 8377.03125
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "measure": "device_peak",
      "peak_error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.72GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 174.375
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 366.68359375
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 329.40625
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 1240.310546875
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 174.375
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 366.68359375
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 329.40625
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 1240.310546875
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 682.0
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 2238.877197265625
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 8461.97509765625
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "measure": "device_peak",
      "peak_error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.67GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 190.625
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 490.8779296875
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 354.475341796875
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 1556.8935546875
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 190.625
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 490.878173828125
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 354.475341796875
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 1556.9716796875
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 523.80908203125
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 5668.3984375
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 204.242919921875
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 535.648193359375
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 64,
      "batch": 8,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 261.251220703125
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 64,
      "batch": 8,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 12662.250732421875
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 64,
      "batch": 8,
      "op": "forward",
      "measure": "device_peak",
      "peak_mb": 254.0
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 64,
      "batch": 8,
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 398.00048828125
    },
    {
      "suite": "scan_remat",
      "model": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "none",
      "op": "train_step",
      "params": 132412200,
      "finite": true,
      "trace_s": 2.095586634997744,
      "compile_s": 16.895156646998657,
      "xla_temp_mb": 3848.238311767578,
      "xla_arg_mb": 1515.4623794555664,
      "xla_out_mb": 1515.3507843017578,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 228.97155700047733,
      "q1_ms": 228.58100500161527,
      "q3_ms": 229.03771349956514,
      "first_call_s": 0.3937067529986962
    },
    {
      "suite": "scan_remat",
      "model": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "none",
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 6893.590576171875
    },
    {
      "suite": "scan_remat",
      "model": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "full",
      "op": "train_step",
      "params": 132412200,
      "finite": true,
      "trace_s": 1.205418504003319,
      "compile_s": 13.824197428002662,
      "xla_temp_mb": 328.2567901611328,
      "xla_arg_mb": 1515.4623794555664,
      "xla_out_mb": 1515.3507843017578,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 274.243783998827,
      "q1_ms": 273.0750299997453,
      "q3_ms": 274.5389175033779,
      "first_call_s": 0.3802042040042579
    },
    {
      "suite": "scan_remat",
      "model": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "full",
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 3373.60888671875
    },
    {
      "suite": "scan_remat",
      "model": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "dots_saveable",
      "op": "train_step",
      "params": 132412200,
      "finite": true,
      "trace_s": 1.2256390990005457,
      "compile_s": 12.966631349001545,
      "xla_temp_mb": 1824.2671127319336,
      "xla_arg_mb": 1515.4623794555664,
      "xla_out_mb": 1515.3507843017578,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 263.42352900246624,
      "q1_ms": 262.02435300365323,
      "q3_ms": 264.58142600313295,
      "first_call_s": 0.3599403180051013
    },
    {
      "suite": "scan_remat",
      "model": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": false,
      "remat": "dots_saveable",
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 4869.619384765625
    },
    {
      "suite": "scan_remat",
      "model": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "none",
      "op": "train_step",
      "params": 132412200,
      "finite": true,
      "trace_s": 0.3521739579955465,
      "compile_s": 2.8970100169972284,
      "xla_temp_mb": 8081.1671142578125,
      "xla_arg_mb": 1515.4623794555664,
      "xla_out_mb": 1515.338150024414,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 257.1779870049795,
      "q1_ms": 256.75797700023395,
      "q3_ms": 258.3601385013026,
      "first_call_s": 0.36277653200522764
    },
    {
      "suite": "scan_remat",
      "model": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "none",
      "op": "train_step",
      "measure": "device_peak",
      "peak_error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 7.89GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
    },
    {
      "suite": "scan_remat",
      "model": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "full",
      "op": "train_step",
      "params": 132412200,
      "finite": true,
      "trace_s": 0.31852480999805266,
      "compile_s": 1.5512422609972418,
      "xla_temp_mb": 203.78482055664062,
      "xla_arg_mb": 1515.4623794555664,
      "xla_out_mb": 1515.338150024414,
      "xla_alias_mb": 0.0,
      "inner": 1,
      "median_ms": 265.4083070010529,
      "q1_ms": 265.0628935043642,
      "q3_ms": 266.4260569981707,
      "first_call_s": 0.2746178359957412
    },
    {
      "suite": "scan_remat",
      "model": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 2,
      "model_dim": 512,
      "layers": 24,
      "scan": true,
      "remat": "full",
      "op": "train_step",
      "measure": "device_peak",
      "peak_mb": 3385.11865234375
    }
  ]
}