{
  "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": {
    "epsilon": 0.5
  },
  "results": [
    {
      "suite": "accuracy",
      "impl": "dense",
      "n": 2048,
      "d": 8,
      "eps": 0.5,
      "lse_max_abs_err": 8.87399717797166e-06,
      "grad_rel_err": 1.2889438717622052e-06
    },
    {
      "suite": "accuracy",
      "impl": "pallas",
      "n": 2048,
      "d": 8,
      "eps": 0.5,
      "lse_max_abs_err": 1.283099290105838e-05,
      "grad_rel_err": 1.6526876170735592e-06
    },
    {
      "suite": "accuracy",
      "impl": "dense",
      "n": 2048,
      "d": 8,
      "eps": 0.01,
      "lse_max_abs_err": 0.0004396351951925226,
      "grad_rel_err": 3.933484188226024e-05
    },
    {
      "suite": "accuracy",
      "impl": "pallas",
      "n": 2048,
      "d": 8,
      "eps": 0.01,
      "lse_max_abs_err": 0.0006696364775962138,
      "grad_rel_err": 5.0953895602938286e-05
    },
    {
      "suite": "accuracy",
      "impl": "dense",
      "n": 2048,
      "d": 32,
      "eps": 0.05,
      "lse_max_abs_err": 0.00033174961572512984,
      "grad_rel_err": 3.92375901866659e-05
    },
    {
      "suite": "accuracy",
      "impl": "pallas",
      "n": 2048,
      "d": 32,
      "eps": 0.05,
      "lse_max_abs_err": 0.00044034209895471577,
      "grad_rel_err": 3.618726291664595e-05
    },
    {
      "suite": "accuracy",
      "impl": "dense_default_precision",
      "n": 2048,
      "d": 8,
      "eps": 0.01,
      "lse_max_abs_err": 1.152265871297459
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd",
      "n": 1024,
      "m": 1024,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 0.125,
      "xla_arg_mb": 0.01953125,
      "xla_out_mb": 0.00390625,
      "compile_s": 0.00014774900046177208,
      "inner": 200,
      "median_ms": 0.04531748501904076,
      "q1_ms": 0.04509834499913268,
      "q3_ms": 0.04585595499520423
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "n": 1024,
      "m": 1024,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 0.19580078125,
      "xla_arg_mb": 0.01953125,
      "xla_out_mb": 0.01955413818359375,
      "compile_s": 0.00030152100225677714,
      "inner": 163,
      "median_ms": 0.1753421656646697,
      "q1_ms": 0.17528543252343642,
      "q3_ms": 0.17535513191215374
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd",
      "n": 1024,
      "m": 1024,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 4.01171875,
      "xla_arg_mb": 0.01953125,
      "xla_out_mb": 0.00390625,
      "compile_s": 0.00014134799857856706,
      "inner": 200,
      "median_ms": 0.05405877502198564,
      "q1_ms": 0.05351750251065823,
      "q3_ms": 0.05472696500874008
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 1024,
      "m": 1024,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 4.0080718994140625,
      "xla_arg_mb": 0.01953125,
      "xla_out_mb": 0.01955413818359375,
      "compile_s": 0.0001594320056028664,
      "inner": 200,
      "median_ms": 0.05840549998538336,
      "q1_ms": 0.05831262249557767,
      "q3_ms": 0.058532237489998806
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd",
      "n": 4096,
      "m": 4096,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 0.5,
      "xla_arg_mb": 0.078125,
      "xla_out_mb": 0.015625,
      "compile_s": 0.00024743799440329894,
      "inner": 200,
      "median_ms": 0.15233740999974543,
      "q1_ms": 0.1522750425101549,
      "q3_ms": 0.15239962749546976
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "n": 4096,
      "m": 4096,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 0.78173828125,
      "xla_arg_mb": 0.078125,
      "xla_out_mb": 0.07814788818359375,
      "compile_s": 0.000701167002262082,
      "inner": 72,
      "median_ms": 0.5702373194556761,
      "q1_ms": 0.5702320208405985,
      "q3_ms": 0.5703224027860819
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd",
      "n": 4096,
      "m": 4096,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 68.00001525878906,
      "xla_arg_mb": 0.078125,
      "xla_out_mb": 0.015625,
      "compile_s": 0.0003462649983703159,
      "inner": 154,
      "median_ms": 0.2434758961109245,
      "q1_ms": 0.24254549675611264,
      "q3_ms": 0.24360123052837798
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 4096,
      "m": 4096,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 128.03175354003906,
      "xla_arg_mb": 0.078125,
      "xla_out_mb": 0.07814788818359375,
      "compile_s": 0.0007452299978467636,
      "inner": 67,
      "median_ms": 0.6469203282578556,
      "q1_ms": 0.6456995148771206,
      "q3_ms": 0.6469943432081768
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd",
      "n": 16384,
      "m": 16384,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 0.3125,
      "xla_out_mb": 0.0625,
      "compile_s": 0.0013583499967353418,
      "inner": 37,
      "median_ms": 1.2218310541750208,
      "q1_ms": 1.2158850675749253,
      "q3_ms": 1.2243892839202042
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "n": 16384,
      "m": 16384,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 3.12548828125,
      "xla_arg_mb": 0.3125,
      "xla_out_mb": 0.31252288818359375,
      "compile_s": 0.0065366070048185065,
      "inner": 7,
      "median_ms": 6.513827856646718,
      "q1_ms": 6.488596499756179,
      "q3_ms": 6.518726428371988
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd",
      "n": 16384,
      "m": 16384,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 1028.000015258789,
      "xla_arg_mb": 0.3125,
      "xla_out_mb": 0.0625,
      "compile_s": 0.004099454999959562,
      "inner": 12,
      "median_ms": 4.476529999616711,
      "q1_ms": 4.424239791357347,
      "q3_ms": 4.4795494162220475
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 16384,
      "m": 16384,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 2064.0627670288086,
      "xla_arg_mb": 0.3125,
      "xla_out_mb": 0.31252288818359375,
      "compile_s": 0.010387932001322042,
      "inner": 4,
      "median_ms": 10.662845750630368,
      "q1_ms": 10.651913749825326,
      "q3_ms": 10.765633875053027
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd",
      "n": 65536,
      "m": 65536,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 1.25,
      "xla_out_mb": 0.25,
      "compile_s": 0.015578241996990982,
      "inner": 3,
      "median_ms": 15.899770999870572,
      "q1_ms": 15.889559999777703,
      "q3_ms": 15.907032999772733
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "n": 65536,
      "m": 65536,
      "d": 2,
      "dtype": "float32",
      "xla_temp_mb": 12.50048828125,
      "xla_arg_mb": 1.25,
      "xla_out_mb": 1.2500228881835938,
      "compile_s": 0.08694458099489566,
      "inner": 1,
      "median_ms": 87.16526799980784,
      "q1_ms": 87.12183600073331,
      "q3_ms": 87.44443150135339
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd",
      "n": 65536,
      "m": 65536,
      "d": 2,
      "dtype": "float32",
      "mem_error": "RESOURCE_EXHAUSTED: Failed to get configs for: 3 out of 8 instructions. See logs for all failures. Example failure: ",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Failed to get configs for: 3 out of 8 instructions. See logs for all failures. Example failure: "
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 65536,
      "m": 65536,
      "d": 2,
      "dtype": "float32",
      "mem_error": "RESOURCE_EXHAUSTED: Failed to get configs for: 5 out of 16 instructions. See logs for all failures. Example failure: ",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Failed to get configs for: 5 out of 16 instructions. See logs for all failures. Example failure: "
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd",
      "n": 1024,
      "m": 1024,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 0.125,
      "xla_arg_mb": 0.06640625,
      "xla_out_mb": 0.00390625,
      "compile_s": 0.00011873499897774309,
      "inner": 200,
      "median_ms": 0.04595940001308918,
      "q1_ms": 0.04580315500788856,
      "q3_ms": 0.049436502522439696
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "n": 1024,
      "m": 1024,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 0.19580078125,
      "xla_arg_mb": 0.06640625,
      "xla_out_mb": 0.06642913818359375,
      "compile_s": 0.0002738580005825497,
      "inner": 182,
      "median_ms": 0.1755259286005889,
      "q1_ms": 0.1754762994736026,
      "q3_ms": 0.17554742309864824
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd",
      "n": 1024,
      "m": 1024,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 8.000015258789062,
      "xla_arg_mb": 0.06640625,
      "xla_out_mb": 0.00390625,
      "compile_s": 0.00010325600305804983,
      "inner": 200,
      "median_ms": 0.054572350018133875,
      "q1_ms": 0.054257177507679444,
      "q3_ms": 0.055670325018581934
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 1024,
      "m": 1024,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 8.008316040039062,
      "xla_arg_mb": 0.06640625,
      "xla_out_mb": 0.06642913818359375,
      "compile_s": 0.00014037599612493068,
      "inner": 200,
      "median_ms": 0.05697073498595273,
      "q1_ms": 0.05672446748576476,
      "q3_ms": 0.05827890749060316
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd",
      "n": 4096,
      "m": 4096,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 0.5,
      "xla_arg_mb": 0.265625,
      "xla_out_mb": 0.015625,
      "compile_s": 0.00028907699743285775,
      "inner": 187,
      "median_ms": 0.1523580320738505,
      "q1_ms": 0.15230576737231655,
      "q3_ms": 0.15236622993809418
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "n": 4096,
      "m": 4096,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 0.78173828125,
      "xla_arg_mb": 0.265625,
      "xla_out_mb": 0.26564788818359375,
      "compile_s": 0.0007159739980124868,
      "inner": 70,
      "median_ms": 0.6066477142407426,
      "q1_ms": 0.6006784928363882,
      "q3_ms": 0.6069485713981391
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd",
      "n": 4096,
      "m": 4096,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 68.00001525878906,
      "xla_arg_mb": 0.265625,
      "xla_out_mb": 0.015625,
      "compile_s": 0.0003164790032315068,
      "inner": 159,
      "median_ms": 0.24400540878674096,
      "q1_ms": 0.24334450628324383,
      "q3_ms": 0.24401980815854352
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 4096,
      "m": 4096,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 128.03175354003906,
      "xla_arg_mb": 0.265625,
      "xla_out_mb": 0.26564788818359375,
      "compile_s": 0.0007146220013964921,
      "inner": 68,
      "median_ms": 0.6498078971216112,
      "q1_ms": 0.6485942868119523,
      "q3_ms": 0.6501675515049583
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd",
      "n": 16384,
      "m": 16384,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 1.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.001354974003334064,
      "inner": 36,
      "median_ms": 1.2408581387717277,
      "q1_ms": 1.229968236051112,
      "q3_ms": 1.2420910277594683
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "n": 16384,
      "m": 16384,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 3.12548828125,
      "xla_arg_mb": 1.0625,
      "xla_out_mb": 1.0625228881835938,
      "compile_s": 0.006546494994836394,
      "inner": 7,
      "median_ms": 6.924822571038801,
      "q1_ms": 6.781871356776849,
      "q3_ms": 7.019407928705082
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd",
      "n": 16384,
      "m": 16384,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 1028.000015258789,
      "xla_arg_mb": 1.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.0040282109985128045,
      "inner": 12,
      "median_ms": 4.378555499594465,
      "q1_ms": 4.315096291369022,
      "q3_ms": 4.397210374615194
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 16384,
      "m": 16384,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 2064.0627670288086,
      "xla_arg_mb": 1.0625,
      "xla_out_mb": 1.0625228881835938,
      "compile_s": 0.010283715004334226,
      "inner": 4,
      "median_ms": 10.628157500832458,
      "q1_ms": 10.614595626066148,
      "q3_ms": 10.717339875554899
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd",
      "n": 65536,
      "m": 65536,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 0.25,
      "compile_s": 0.01771580600325251,
      "inner": 2,
      "median_ms": 18.254406000778545,
      "q1_ms": 17.928010751347756,
      "q3_ms": 18.413329000395606
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "n": 65536,
      "m": 65536,
      "d": 8,
      "dtype": "float32",
      "xla_temp_mb": 12.50048828125,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.250022888183594,
      "compile_s": 0.09011533900047652,
      "inner": 1,
      "median_ms": 90.09351800341392,
      "q1_ms": 90.01102250113036,
      "q3_ms": 90.31318350025686
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd",
      "n": 65536,
      "m": 65536,
      "d": 8,
      "dtype": "float32",
      "mem_error": "RESOURCE_EXHAUSTED: Failed to get configs for: 3 out of 8 instructions. See logs for all failures. Example failure: ",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Failed to get configs for: 3 out of 8 instructions. See logs for all failures. Example failure: "
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 65536,
      "m": 65536,
      "d": 8,
      "dtype": "float32",
      "mem_error": "NOT_FOUND: Failed to get configs for: 5 out of 16 instructions. See logs for all failures. Example failure: ",
      "error": "JaxRuntimeError: NOT_FOUND: Failed to get configs for: 5 out of 16 instructions. See logs for all failures. Example failure: "
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd",
      "n": 1024,
      "m": 1024,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 0.25390625,
      "xla_out_mb": 0.00390625,
      "compile_s": 0.00014440299855777994,
      "inner": 200,
      "median_ms": 0.0725564700042014,
      "q1_ms": 0.07156374251280795,
      "q3_ms": 0.07259506750415312
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "n": 1024,
      "m": 1024,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 0.01220703125,
      "xla_arg_mb": 0.25390625,
      "xla_out_mb": 0.25392913818359375,
      "compile_s": 0.0004195439978502691,
      "inner": 119,
      "median_ms": 0.3248620083743409,
      "q1_ms": 0.32178778568196637,
      "q3_ms": 0.3337124033641274
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd",
      "n": 1024,
      "m": 1024,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 8.000015258789062,
      "xla_arg_mb": 0.25390625,
      "xla_out_mb": 0.00390625,
      "compile_s": 0.00011097999959019944,
      "inner": 200,
      "median_ms": 0.06269095501920674,
      "q1_ms": 0.06035045249518589,
      "q3_ms": 0.06661816001724219
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 1024,
      "m": 1024,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 8.008316040039062,
      "xla_arg_mb": 0.25390625,
      "xla_out_mb": 0.25392913818359375,
      "compile_s": 0.0002473380009178072,
      "inner": 199,
      "median_ms": 0.0653907989901163,
      "q1_ms": 0.0653202386909918,
      "q3_ms": 0.06792305276718053
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd",
      "n": 4096,
      "m": 4096,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.015625,
      "xla_out_mb": 0.015625,
      "compile_s": 0.0003564040016499348,
      "inner": 132,
      "median_ms": 0.27308938633447694,
      "q1_ms": 0.2704054469645736,
      "q3_ms": 0.2741680946752295
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "n": 4096,
      "m": 4096,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 0.04736328125,
      "xla_arg_mb": 1.015625,
      "xla_out_mb": 1.0156478881835938,
      "compile_s": 0.0012559759998111986,
      "inner": 37,
      "median_ms": 1.121592486506907,
      "q1_ms": 1.1212986891395476,
      "q3_ms": 1.1239312432833783
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd",
      "n": 4096,
      "m": 4096,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 68.00001525878906,
      "xla_arg_mb": 1.015625,
      "xla_out_mb": 0.015625,
      "compile_s": 0.0003336709996801801,
      "inner": 149,
      "median_ms": 0.2756188254926085,
      "q1_ms": 0.26707965100786474,
      "q3_ms": 0.2796520939655055
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 4096,
      "m": 4096,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 128.03175354003906,
      "xla_arg_mb": 1.015625,
      "xla_out_mb": 1.0156478881835938,
      "compile_s": 0.0007738640051684342,
      "inner": 65,
      "median_ms": 0.6983720153444806,
      "q1_ms": 0.6914658768577143,
      "q3_ms": 0.6987891922807858
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd",
      "n": 16384,
      "m": 16384,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 4.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.002620618004584685,
      "inner": 19,
      "median_ms": 2.5287130529728814,
      "q1_ms": 2.5250144476319147,
      "q3_ms": 2.5427068949762828
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "n": 16384,
      "m": 16384,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 0.18798828125,
      "xla_arg_mb": 4.0625,
      "xla_out_mb": 4.062522888183594,
      "compile_s": 0.014920478002750315,
      "inner": 3,
      "median_ms": 15.033285000148075,
      "q1_ms": 14.968789499107515,
      "q3_ms": 15.111633166573787
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd",
      "n": 16384,
      "m": 16384,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 1028.000015258789,
      "xla_arg_mb": 4.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.004384785999718588,
      "inner": 11,
      "median_ms": 4.88606672694335,
      "q1_ms": 4.770549545355607,
      "q3_ms": 4.887886090727989
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 16384,
      "m": 16384,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 2064.0627670288086,
      "xla_arg_mb": 4.0625,
      "xla_out_mb": 4.062522888183594,
      "compile_s": 0.011050656998122577,
      "inner": 4,
      "median_ms": 12.10914950024744,
      "q1_ms": 11.713747625435644,
      "q3_ms": 12.682155999755196
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd",
      "n": 65536,
      "m": 65536,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 16.25,
      "xla_out_mb": 0.25,
      "compile_s": 0.03568909000023268,
      "inner": 1,
      "median_ms": 35.870163002982736,
      "q1_ms": 35.79457050363999,
      "q3_ms": 36.19500750210136
    },
    {
      "suite": "lse",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "n": 65536,
      "m": 65536,
      "d": 32,
      "dtype": "float32",
      "xla_temp_mb": 0.75048828125,
      "xla_arg_mb": 16.25,
      "xla_out_mb": 16.250022888183594,
      "compile_s": 0.23745588400197448,
      "inner": 1,
      "median_ms": 238.76559199561598,
      "q1_ms": 238.08356649897178,
      "q3_ms": 244.38884099799907
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd",
      "n": 65536,
      "m": 65536,
      "d": 32,
      "dtype": "float32",
      "mem_error": "RESOURCE_EXHAUSTED: Failed to get configs for: 3 out of 8 instructions. See logs for all failures. Example failure: ",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Failed to get configs for: 3 out of 8 instructions. See logs for all failures. Example failure: "
    },
    {
      "suite": "lse",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 65536,
      "m": 65536,
      "d": 32,
      "dtype": "float32",
      "mem_error": "RESOURCE_EXHAUSTED: Failed to get configs for: 5 out of 16 instructions. See logs for all failures. Example failure: ",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Failed to get configs for: 5 out of 16 instructions. See logs for all failures. Example failure: "
    },
    {
      "suite": "sinkhorn",
      "impl": "kernel",
      "mode": "fwd",
      "n": 4096,
      "m": 4096,
      "d": 8,
      "iters": 50,
      "eps": 0.8975364685058594,
      "xla_temp_mb": 1.3286514282226562,
      "xla_arg_mb": 0.25,
      "xla_out_mb": 3.814697265625e-06,
      "compile_s": 0.013414269000350032,
      "inner": 3,
      "median_ms": 13.828763333246266,
      "q1_ms": 13.635270666175831,
      "q3_ms": 13.83131299977928
    },
    {
      "suite": "sinkhorn",
      "impl": "dense",
      "mode": "fwd",
      "n": 4096,
      "m": 4096,
      "d": 8,
      "iters": 50,
      "eps": 0.8975364685058594,
      "xla_temp_mb": 128.03201293945312,
      "xla_arg_mb": 0.25,
      "xla_out_mb": 3.814697265625e-06,
      "compile_s": 0.016233452995948028,
      "inner": 3,
      "median_ms": 16.27617300012692,
      "q1_ms": 16.23369616694011,
      "q3_ms": 16.27913216725574
    },
    {
      "suite": "sinkhorn",
      "impl": "kernel",
      "mode": "fwd_bwd_implicit",
      "n": 4096,
      "m": 4096,
      "d": 8,
      "iters": 50,
      "eps": 0.8975364685058594,
      "xla_temp_mb": 1.3129920959472656,
      "xla_arg_mb": 0.25,
      "xla_out_mb": 0.125,
      "compile_s": 0.013607263994344976,
      "inner": 3,
      "median_ms": 13.471530667933015,
      "q1_ms": 13.468310000462225,
      "q3_ms": 13.47286616752778
    },
    {
      "suite": "sinkhorn",
      "impl": "agreement",
      "n": 4096,
      "d": 8,
      "cost_kernel": 14.728631973266602,
      "cost_dense": 14.728631973266602,
      "abs_diff": 0.0
    },
    {
      "suite": "sinkhorn",
      "impl": "kernel",
      "mode": "fwd",
      "n": 16384,
      "m": 16384,
      "d": 8,
      "iters": 50,
      "eps": 0.8991072654724122,
      "xla_temp_mb": 5.313026428222656,
      "xla_arg_mb": 1.0,
      "xla_out_mb": 3.814697265625e-06,
      "compile_s": 0.12678120200143894,
      "inner": 1,
      "median_ms": 127.33585300156847,
      "q1_ms": 127.04284350184025,
      "q3_ms": 127.43647249953938
    },
    {
      "suite": "sinkhorn",
      "impl": "dense",
      "mode": "fwd",
      "n": 16384,
      "m": 16384,
      "d": 8,
      "iters": 50,
      "eps": 0.8991072654724122,
      "xla_temp_mb": 2048.125762939453,
      "xla_arg_mb": 1.0,
      "xla_out_mb": 3.814697265625e-06,
      "compile_s": 0.24999482800194528,
      "inner": 1,
      "median_ms": 250.16787899949122,
      "q1_ms": 250.0920464990486,
      "q3_ms": 250.31712649797555
    },
    {
      "suite": "sinkhorn",
      "impl": "kernel",
      "mode": "fwd_bwd_implicit",
      "n": 16384,
      "m": 16384,
      "d": 8,
      "iters": 50,
      "eps": 0.8991072654724122,
      "xla_temp_mb": 5.250492095947266,
      "xla_arg_mb": 1.0,
      "xla_out_mb": 0.5,
      "compile_s": 0.1286964259998058,
      "inner": 1,
      "median_ms": 129.55856699409196,
      "q1_ms": 129.54499649640638,
      "q3_ms": 129.7159704954538
    },
    {
      "suite": "sinkhorn",
      "impl": "agreement",
      "n": 16384,
      "d": 8,
      "cost_kernel": 14.778955459594727,
      "cost_dense": 14.77895450592041,
      "abs_diff": 9.5367431640625e-07
    },
    {
      "suite": "sinkhorn",
      "impl": "kernel",
      "mode": "fwd",
      "n": 65536,
      "m": 65536,
      "d": 8,
      "iters": 50,
      "eps": 0.898257350921631,
      "xla_temp_mb": 21.250526428222656,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.814697265625e-06,
      "compile_s": 1.6487746470011189,
      "inner": 1,
      "median_ms": 1656.4462309979717,
      "q1_ms": 1656.2574744966696,
      "q3_ms": 1656.861082996329
    },
    {
      "suite": "sinkhorn",
      "impl": "dense",
      "mode": "fwd",
      "n": 65536,
      "m": 65536,
      "d": 8,
      "iters": 50,
      "eps": 0.898257350921631,
      "mem_error": "RESOURCE_EXHAUSTED: Failed to get configs for: 5 out of 22 instructions. See logs for all failures. Example failure: ",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Failed to get configs for: 5 out of 22 instructions. See logs for all failures. Example failure: "
    },
    {
      "suite": "sinkhorn",
      "impl": "kernel",
      "mode": "fwd_bwd_implicit",
      "n": 65536,
      "m": 65536,
      "d": 8,
      "iters": 50,
      "eps": 0.898257350921631,
      "xla_temp_mb": 21.000492095947266,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 2.0,
      "compile_s": 1.6878770629991777,
      "inner": 1,
      "median_ms": 1687.96111499978,
      "q1_ms": 1687.9181885014987,
      "q3_ms": 1688.0122869988554
    },
    {
      "suite": "scale",
      "impl": "kernel",
      "mode": "fwd",
      "n": 131072,
      "m": 131072,
      "d": 2,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.06117820699728327,
      "inner": 1,
      "median_ms": 61.098144004063215,
      "q1_ms": 61.098144004063215,
      "q3_ms": 61.098144004063215
    },
    {
      "suite": "scale",
      "impl": "kernel",
      "mode": "fwd_bwd",
      "n": 131072,
      "m": 131072,
      "d": 2,
      "xla_temp_mb": 25.00048828125,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 2.5000228881835938,
      "compile_s": 0.33617783700174186,
      "inner": 1,
      "median_ms": 336.69053099583834,
      "q1_ms": 336.69053099583834,
      "q3_ms": 336.69053099583834
    },
    {
      "suite": "scale",
      "impl": "kernel",
      "mode": "fwd",
      "n": 262144,
      "m": 262144,
      "d": 2,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.24335626600077376,
      "inner": 1,
      "median_ms": 245.6712709972635,
      "q1_ms": 245.6712709972635,
      "q3_ms": 245.6712709972635
    },
    {
      "suite": "scale",
      "impl": "kernel",
      "mode": "fwd_bwd",
      "n": 262144,
      "m": 262144,
      "d": 2,
      "xla_temp_mb": 50.00048828125,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 5.000022888183594,
      "compile_s": 1.3397905680030817,
      "inner": 1,
      "median_ms": 1340.528164000716,
      "q1_ms": 1340.528164000716,
      "q3_ms": 1340.528164000716
    },
    {
      "suite": "scale",
      "impl": "kernel",
      "mode": "fwd",
      "n": 1048576,
      "m": 1048576,
      "d": 2,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 4.0,
      "compile_s": 3.974350686003163,
      "inner": 1,
      "median_ms": 3976.628623000579,
      "q1_ms": 3976.628623000579,
      "q3_ms": 3976.628623000579
    },
    {
      "suite": "scale",
      "impl": "kernel",
      "mode": "fwd_bwd",
      "n": 1048576,
      "m": 1048576,
      "d": 2,
      "xla_temp_mb": 200.00048828125,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 20.000022888183594,
      "compile_s": 21.6220034470025,
      "inner": 1,
      "median_ms": 21638.0784819994,
      "q1_ms": 21638.0784819994,
      "q3_ms": 21638.0784819994
    },
    {
      "suite": "scale",
      "impl": "dense",
      "mode": "fwd",
      "n": 32768,
      "m": 32768,
      "d": 2,
      "xla_temp_mb": 4112.250015258789,
      "xla_arg_mb": 0.625,
      "xla_out_mb": 0.125,
      "compile_s": 0.014836210000794381,
      "inner": 3,
      "median_ms": 14.965450332965702,
      "q1_ms": 14.965450332965702,
      "q3_ms": 14.965450332965702
    },
    {
      "suite": "scale",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 32768,
      "m": 32768,
      "d": 2,
      "xla_temp_mb": 8224.125267028809,
      "xla_arg_mb": 0.625,
      "xla_out_mb": 0.6250228881835938,
      "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__lambda']"
    },
    {
      "suite": "scale",
      "impl": "dense",
      "mode": "fwd",
      "n": 49152,
      "m": 49152,
      "d": 2,
      "mem_error": "RESOURCE_EXHAUSTED: Failed to get configs for: 3 out of 8 instructions. See logs for all failures. Example failure: ",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Failed to get configs for: 3 out of 8 instructions. See logs for all failures. Example failure: "
    },
    {
      "suite": "scale",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 49152,
      "m": 49152,
      "d": 2,
      "mem_error": "RESOURCE_EXHAUSTED: Failed to get configs for: 5 out of 16 instructions. See logs for all failures. Example failure: ",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Failed to get configs for: 5 out of 16 instructions. See logs for all failures. Example failure: "
    },
    {
      "suite": "scale",
      "impl": "dense",
      "mode": "fwd",
      "n": 65536,
      "m": 65536,
      "d": 2,
      "mem_error": "RESOURCE_EXHAUSTED: Failed to get configs for: 3 out of 8 instructions. See logs for all failures. Example failure: ",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Failed to get configs for: 3 out of 8 instructions. See logs for all failures. Example failure: "
    },
    {
      "suite": "scale",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 65536,
      "m": 65536,
      "d": 2,
      "mem_error": "RESOURCE_EXHAUSTED: Failed to get configs for: 5 out of 16 instructions. See logs for all failures. Example failure: ",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Failed to get configs for: 5 out of 16 instructions. See logs for all failures. Example failure: "
    },
    {
      "suite": "peak",
      "impl": "kernel",
      "mode": "fwd",
      "n": 65536,
      "m": 65536,
      "d": 8,
      "peak_mb": 12.500732421875,
      "inputs_peak_mb": 4.5009765625
    },
    {
      "suite": "peak",
      "impl": "kernel",
      "mode": "fwd_bwd",
      "n": 65536,
      "m": 65536,
      "d": 8,
      "peak_mb": 24.501220703125,
      "inputs_peak_mb": 4.5009765625
    },
    {
      "suite": "peak",
      "impl": "dense",
      "mode": "fwd",
      "n": 32768,
      "m": 32768,
      "d": 8,
      "peak_mb": 4259.37451171875,
      "inputs_peak_mb": 3.25
    },
    {
      "suite": "peak",
      "impl": "dense",
      "mode": "fwd_bwd",
      "n": 32768,
      "m": 32768,
      "d": 8,
      "error": "jax.errors.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__lambda']"
    },
    {
      "suite": "peak",
      "impl": "kernel",
      "mode": "fwd_bwd",
      "n": 1048576,
      "m": 1048576,
      "d": 2,
      "peak_mb": 296.000732421875,
      "inputs_peak_mb": 24.0009765625
    },
    {
      "suite": "cfm",
      "coupling": "independent",
      "seed": 0,
      "steps": 3000,
      "batch": 256,
      "train_s": 11.058893421002722,
      "ms_per_step": 3.686297807000907,
      "final_loss": 6.683660507202148,
      "sdiv_floor": 0.01149708777666092,
      "sdiv_nfe1": 13.585737228393555,
      "sdiv_nfe2": 11.172250747680664,
      "sdiv_nfe4": 0.41724079847335815,
      "sdiv_nfe8": 0.09921444952487946,
      "sdiv_nfe32": 0.05322100222110748,
      "straightness": 3.1146295070648193
    },
    {
      "suite": "cfm",
      "coupling": "exact",
      "seed": 0,
      "steps": 3000,
      "batch": 256,
      "train_s": 36.3240146820026,
      "ms_per_step": 12.108004894000867,
      "final_loss": 0.20252424478530884,
      "sdiv_floor": 0.01149708777666092,
      "sdiv_nfe1": 0.3216308355331421,
      "sdiv_nfe2": 2.4513533115386963,
      "sdiv_nfe4": 0.2783716917037964,
      "sdiv_nfe8": 0.09106561541557312,
      "sdiv_nfe32": 0.051405563950538635,
      "straightness": 0.10777132213115692
    },
    {
      "suite": "cfm",
      "coupling": "sinkhorn",
      "seed": 0,
      "steps": 3000,
      "batch": 256,
      "train_s": 11.888259163002658,
      "ms_per_step": 3.9627530543342195,
      "final_loss": 0.9002116918563843,
      "sdiv_floor": 0.01149708777666092,
      "sdiv_nfe1": 0.5429400205612183,
      "sdiv_nfe2": 2.8602066040039062,
      "sdiv_nfe4": 0.21758225560188293,
      "sdiv_nfe8": 0.09207914769649506,
      "sdiv_nfe32": 0.04995211958885193,
      "straightness": 0.1738063246011734
    },
    {
      "suite": "cfm",
      "coupling": "independent",
      "seed": 1,
      "steps": 3000,
      "batch": 256,
      "train_s": 11.172408806000021,
      "ms_per_step": 3.724136268666674,
      "final_loss": 6.981533050537109,
      "sdiv_floor": 0.01149708777666092,
      "sdiv_nfe1": 14.664462089538574,
      "sdiv_nfe2": 11.865185737609863,
      "sdiv_nfe4": 0.4806097149848938,
      "sdiv_nfe8": 0.10507824271917343,
      "sdiv_nfe32": 0.059169426560401917,
      "straightness": 3.785977363586426
    },
    {
      "suite": "cfm",
      "coupling": "exact",
      "seed": 1,
      "steps": 3000,
      "batch": 256,
      "train_s": 36.69143724399328,
      "ms_per_step": 12.230479081331092,
      "final_loss": 0.14435811340808868,
      "sdiv_floor": 0.01149708777666092,
      "sdiv_nfe1": 0.30231955647468567,
      "sdiv_nfe2": 2.2614312171936035,
      "sdiv_nfe4": 0.34076371788978577,
      "sdiv_nfe8": 0.12648041546344757,
      "sdiv_nfe32": 0.07177229970693588,
      "straightness": 0.10037004947662354
    },
    {
      "suite": "cfm",
      "coupling": "sinkhorn",
      "seed": 1,
      "steps": 3000,
      "batch": 256,
      "train_s": 11.983504303003429,
      "ms_per_step": 3.9945014343344765,
      "final_loss": 0.5709594488143921,
      "sdiv_floor": 0.01149708777666092,
      "sdiv_nfe1": 0.5795177221298218,
      "sdiv_nfe2": 2.792799234390259,
      "sdiv_nfe4": 0.24615135788917542,
      "sdiv_nfe8": 0.09435223042964935,
      "sdiv_nfe32": 0.04393412172794342,
      "straightness": 0.1779761165380478
    }
  ]
}