├── .gitignore
├── Cargo.lock
├── Cargo.toml
├── LICENSE
├── README.md
├── download.sh
├── mamba-ssm
    ├── Cargo.toml
    ├── README.md
    ├── benches
    │   ├── conv1d.rs
    │   ├── elemwise.rs
    │   └── linear.rs
    └── src
    │   ├── bin
    │       └── mamba-cli.rs
    │   ├── context.rs
    │   ├── lib.rs
    │   ├── model.rs
    │   ├── nn
    │       ├── conv1d.rs
    │       ├── linear.rs
    │       ├── mod.rs
    │       └── selective_ssm.rs
    │   ├── ops
    │       ├── elemwise.rs
    │       └── mod.rs
    │   └── utils.rs
├── models
    └── README.txt
└── prompts
    └── chat.txt


/.gitignore:
--------------------------------------------------------------------------------
1 | /target
2 | /models/**/*.safetensors
3 | /models/**/*.json
4 | trace-*.json


--------------------------------------------------------------------------------
/Cargo.lock:
--------------------------------------------------------------------------------
   1 | # This file is automatically @generated by Cargo.
   2 | # It is not intended for manual editing.
   3 | version = 3
   4 | 
   5 | [[package]]
   6 | name = "accelerate-src"
   7 | version = "0.3.2"
   8 | source = "registry+https://github.com/rust-lang/crates.io-index"
   9 | checksum = "415ed64958754dbe991900f3940677e6a7eefb4d7367afd70d642677b0c7d19d"
  10 | 
  11 | [[package]]
  12 | name = "adler"
  13 | version = "1.0.2"
  14 | source = "registry+https://github.com/rust-lang/crates.io-index"
  15 | checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
  16 | 
  17 | [[package]]
  18 | name = "aho-corasick"
  19 | version = "1.1.2"
  20 | source = "registry+https://github.com/rust-lang/crates.io-index"
  21 | checksum = "b2969dcb958b36655471fc61f7e416fa76033bdd4bfed0678d8fee1e2d07a1f0"
  22 | dependencies = [
  23 |  "memchr",
  24 | ]
  25 | 
  26 | [[package]]
  27 | name = "android-tzdata"
  28 | version = "0.1.1"
  29 | source = "registry+https://github.com/rust-lang/crates.io-index"
  30 | checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0"
  31 | 
  32 | [[package]]
  33 | name = "android_system_properties"
  34 | version = "0.1.5"
  35 | source = "registry+https://github.com/rust-lang/crates.io-index"
  36 | checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311"
  37 | dependencies = [
  38 |  "libc",
  39 | ]
  40 | 
  41 | [[package]]
  42 | name = "anes"
  43 | version = "0.1.6"
  44 | source = "registry+https://github.com/rust-lang/crates.io-index"
  45 | checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299"
  46 | 
  47 | [[package]]
  48 | name = "anstream"
  49 | version = "0.6.4"
  50 | source = "registry+https://github.com/rust-lang/crates.io-index"
  51 | checksum = "2ab91ebe16eb252986481c5b62f6098f3b698a45e34b5b98200cf20dd2484a44"
  52 | dependencies = [
  53 |  "anstyle",
  54 |  "anstyle-parse",
  55 |  "anstyle-query",
  56 |  "anstyle-wincon",
  57 |  "colorchoice",
  58 |  "utf8parse",
  59 | ]
  60 | 
  61 | [[package]]
  62 | name = "anstyle"
  63 | version = "1.0.4"
  64 | source = "registry+https://github.com/rust-lang/crates.io-index"
  65 | checksum = "7079075b41f533b8c61d2a4d073c4676e1f8b249ff94a393b0595db304e0dd87"
  66 | 
  67 | [[package]]
  68 | name = "anstyle-parse"
  69 | version = "0.2.3"
  70 | source = "registry+https://github.com/rust-lang/crates.io-index"
  71 | checksum = "c75ac65da39e5fe5ab759307499ddad880d724eed2f6ce5b5e8a26f4f387928c"
  72 | dependencies = [
  73 |  "utf8parse",
  74 | ]
  75 | 
  76 | [[package]]
  77 | name = "anstyle-query"
  78 | version = "1.0.1"
  79 | source = "registry+https://github.com/rust-lang/crates.io-index"
  80 | checksum = "a3a318f1f38d2418400f8209655bfd825785afd25aa30bb7ba6cc792e4596748"
  81 | dependencies = [
  82 |  "windows-sys 0.52.0",
  83 | ]
  84 | 
  85 | [[package]]
  86 | name = "anstyle-wincon"
  87 | version = "3.0.2"
  88 | source = "registry+https://github.com/rust-lang/crates.io-index"
  89 | checksum = "1cd54b81ec8d6180e24654d0b371ad22fc3dd083b6ff8ba325b72e00c87660a7"
  90 | dependencies = [
  91 |  "anstyle",
  92 |  "windows-sys 0.52.0",
  93 | ]
  94 | 
  95 | [[package]]
  96 | name = "anyhow"
  97 | version = "1.0.75"
  98 | source = "registry+https://github.com/rust-lang/crates.io-index"
  99 | checksum = "a4668cab20f66d8d020e1fbc0ebe47217433c1b6c8f2040faf858554e394ace6"
 100 | 
 101 | [[package]]
 102 | name = "atty"
 103 | version = "0.2.14"
 104 | source = "registry+https://github.com/rust-lang/crates.io-index"
 105 | checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8"
 106 | dependencies = [
 107 |  "hermit-abi 0.1.19",
 108 |  "libc",
 109 |  "winapi",
 110 | ]
 111 | 
 112 | [[package]]
 113 | name = "autocfg"
 114 | version = "1.1.0"
 115 | source = "registry+https://github.com/rust-lang/crates.io-index"
 116 | checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
 117 | 
 118 | [[package]]
 119 | name = "base16ct"
 120 | version = "0.1.1"
 121 | source = "registry+https://github.com/rust-lang/crates.io-index"
 122 | checksum = "349a06037c7bf932dd7e7d1f653678b2038b9ad46a74102f1fc7bd7872678cce"
 123 | 
 124 | [[package]]
 125 | name = "base64"
 126 | version = "0.13.1"
 127 | source = "registry+https://github.com/rust-lang/crates.io-index"
 128 | checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
 129 | 
 130 | [[package]]
 131 | name = "base64"
 132 | version = "0.21.5"
 133 | source = "registry+https://github.com/rust-lang/crates.io-index"
 134 | checksum = "35636a1494ede3b646cc98f74f8e62c773a38a659ebc777a2cf26b9b74171df9"
 135 | 
 136 | [[package]]
 137 | name = "bindgen_cuda"
 138 | version = "0.1.3"
 139 | source = "registry+https://github.com/rust-lang/crates.io-index"
 140 | checksum = "fcb349f499c53ad8d7b98e9657c61d6225c2812bac84504a97373e3884ca9bac"
 141 | dependencies = [
 142 |  "glob",
 143 |  "num_cpus",
 144 |  "rayon",
 145 | ]
 146 | 
 147 | [[package]]
 148 | name = "bitflags"
 149 | version = "1.3.2"
 150 | source = "registry+https://github.com/rust-lang/crates.io-index"
 151 | checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 152 | 
 153 | [[package]]
 154 | name = "bitflags"
 155 | version = "2.4.1"
 156 | source = "registry+https://github.com/rust-lang/crates.io-index"
 157 | checksum = "327762f6e5a765692301e5bb513e0d9fef63be86bbc14528052b1cd3e6f03e07"
 158 | 
 159 | [[package]]
 160 | name = "block"
 161 | version = "0.1.6"
 162 | source = "registry+https://github.com/rust-lang/crates.io-index"
 163 | checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a"
 164 | 
 165 | [[package]]
 166 | name = "block-buffer"
 167 | version = "0.10.4"
 168 | source = "registry+https://github.com/rust-lang/crates.io-index"
 169 | checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71"
 170 | dependencies = [
 171 |  "generic-array",
 172 | ]
 173 | 
 174 | [[package]]
 175 | name = "bumpalo"
 176 | version = "3.14.0"
 177 | source = "registry+https://github.com/rust-lang/crates.io-index"
 178 | checksum = "7f30e7476521f6f8af1a1c4c0b8cc94f0bee37d91763d0ca2665f299b6cd8aec"
 179 | 
 180 | [[package]]
 181 | name = "bytemuck"
 182 | version = "1.14.0"
 183 | source = "registry+https://github.com/rust-lang/crates.io-index"
 184 | checksum = "374d28ec25809ee0e23827c2ab573d729e293f281dfe393500e7ad618baa61c6"
 185 | dependencies = [
 186 |  "bytemuck_derive",
 187 | ]
 188 | 
 189 | [[package]]
 190 | name = "bytemuck_derive"
 191 | version = "1.5.0"
 192 | source = "registry+https://github.com/rust-lang/crates.io-index"
 193 | checksum = "965ab7eb5f8f97d2a083c799f3a1b994fc397b2fe2da5d1da1626ce15a39f2b1"
 194 | dependencies = [
 195 |  "proc-macro2",
 196 |  "quote",
 197 |  "syn 2.0.39",
 198 | ]
 199 | 
 200 | [[package]]
 201 | name = "byteorder"
 202 | version = "1.5.0"
 203 | source = "registry+https://github.com/rust-lang/crates.io-index"
 204 | checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
 205 | 
 206 | [[package]]
 207 | name = "candle-core"
 208 | version = "0.4.1"
 209 | source = "registry+https://github.com/rust-lang/crates.io-index"
 210 | checksum = "6f1b20174c1707e20f4cb364a355b449803c03e9b0c9193324623cf9787a4e00"
 211 | dependencies = [
 212 |  "accelerate-src",
 213 |  "byteorder",
 214 |  "candle-kernels",
 215 |  "candle-metal-kernels",
 216 |  "cudarc",
 217 |  "gemm",
 218 |  "half 2.3.1",
 219 |  "intel-mkl-src",
 220 |  "libc",
 221 |  "memmap2",
 222 |  "metal",
 223 |  "num-traits",
 224 |  "num_cpus",
 225 |  "rand",
 226 |  "rand_distr",
 227 |  "rayon",
 228 |  "safetensors",
 229 |  "thiserror",
 230 |  "yoke",
 231 |  "zip",
 232 | ]
 233 | 
 234 | [[package]]
 235 | name = "candle-kernels"
 236 | version = "0.4.1"
 237 | source = "registry+https://github.com/rust-lang/crates.io-index"
 238 | checksum = "5845911a44164ebb73b56a0e23793ba1b583bad102af7400fe4768babc5815b2"
 239 | dependencies = [
 240 |  "bindgen_cuda",
 241 | ]
 242 | 
 243 | [[package]]
 244 | name = "candle-metal-kernels"
 245 | version = "0.4.1"
 246 | source = "registry+https://github.com/rust-lang/crates.io-index"
 247 | checksum = "b20d6c0d49121e2709ed9faa958ba915ea59526036bcf27558817d1452a4ff09"
 248 | dependencies = [
 249 |  "metal",
 250 |  "once_cell",
 251 |  "thiserror",
 252 |  "tracing",
 253 | ]
 254 | 
 255 | [[package]]
 256 | name = "candle-nn"
 257 | version = "0.4.1"
 258 | source = "registry+https://github.com/rust-lang/crates.io-index"
 259 | checksum = "66a27533c8edfc915a6459f9850641ef523a829fa1a181c670766c1f752d873a"
 260 | dependencies = [
 261 |  "accelerate-src",
 262 |  "candle-core",
 263 |  "candle-metal-kernels",
 264 |  "half 2.3.1",
 265 |  "intel-mkl-src",
 266 |  "metal",
 267 |  "num-traits",
 268 |  "rayon",
 269 |  "safetensors",
 270 |  "serde",
 271 |  "thiserror",
 272 | ]
 273 | 
 274 | [[package]]
 275 | name = "candle-transformers"
 276 | version = "0.4.1"
 277 | source = "registry+https://github.com/rust-lang/crates.io-index"
 278 | checksum = "b5847699f0643da05e57fc473672566e93dc36d82c1b7eeb970c6154d3434fe1"
 279 | dependencies = [
 280 |  "accelerate-src",
 281 |  "byteorder",
 282 |  "candle-core",
 283 |  "candle-nn",
 284 |  "intel-mkl-src",
 285 |  "num-traits",
 286 |  "rand",
 287 |  "rayon",
 288 |  "serde",
 289 |  "serde_json",
 290 |  "serde_plain",
 291 |  "tracing",
 292 | ]
 293 | 
 294 | [[package]]
 295 | name = "cast"
 296 | version = "0.3.0"
 297 | source = "registry+https://github.com/rust-lang/crates.io-index"
 298 | checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
 299 | 
 300 | [[package]]
 301 | name = "cc"
 302 | version = "1.0.83"
 303 | source = "registry+https://github.com/rust-lang/crates.io-index"
 304 | checksum = "f1174fb0b6ec23863f8b971027804a42614e347eafb0a95bf0b12cdae21fc4d0"
 305 | dependencies = [
 306 |  "libc",
 307 | ]
 308 | 
 309 | [[package]]
 310 | name = "cfg-if"
 311 | version = "1.0.0"
 312 | source = "registry+https://github.com/rust-lang/crates.io-index"
 313 | checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 314 | 
 315 | [[package]]
 316 | name = "chrono"
 317 | version = "0.4.31"
 318 | source = "registry+https://github.com/rust-lang/crates.io-index"
 319 | checksum = "7f2c685bad3eb3d45a01354cedb7d5faa66194d1d58ba6e267a8de788f79db38"
 320 | dependencies = [
 321 |  "android-tzdata",
 322 |  "iana-time-zone",
 323 |  "js-sys",
 324 |  "num-traits",
 325 |  "wasm-bindgen",
 326 |  "windows-targets 0.48.5",
 327 | ]
 328 | 
 329 | [[package]]
 330 | name = "ciborium"
 331 | version = "0.2.1"
 332 | source = "registry+https://github.com/rust-lang/crates.io-index"
 333 | checksum = "effd91f6c78e5a4ace8a5d3c0b6bfaec9e2baaef55f3efc00e45fb2e477ee926"
 334 | dependencies = [
 335 |  "ciborium-io",
 336 |  "ciborium-ll",
 337 |  "serde",
 338 | ]
 339 | 
 340 | [[package]]
 341 | name = "ciborium-io"
 342 | version = "0.2.1"
 343 | source = "registry+https://github.com/rust-lang/crates.io-index"
 344 | checksum = "cdf919175532b369853f5d5e20b26b43112613fd6fe7aee757e35f7a44642656"
 345 | 
 346 | [[package]]
 347 | name = "ciborium-ll"
 348 | version = "0.2.1"
 349 | source = "registry+https://github.com/rust-lang/crates.io-index"
 350 | checksum = "defaa24ecc093c77630e6c15e17c51f5e187bf35ee514f4e2d67baaa96dae22b"
 351 | dependencies = [
 352 |  "ciborium-io",
 353 |  "half 1.8.2",
 354 | ]
 355 | 
 356 | [[package]]
 357 | name = "clap"
 358 | version = "3.2.25"
 359 | source = "registry+https://github.com/rust-lang/crates.io-index"
 360 | checksum = "4ea181bf566f71cb9a5d17a59e1871af638180a18fb0035c92ae62b705207123"
 361 | dependencies = [
 362 |  "bitflags 1.3.2",
 363 |  "clap_lex 0.2.4",
 364 |  "indexmap",
 365 |  "textwrap",
 366 | ]
 367 | 
 368 | [[package]]
 369 | name = "clap"
 370 | version = "4.4.11"
 371 | source = "registry+https://github.com/rust-lang/crates.io-index"
 372 | checksum = "bfaff671f6b22ca62406885ece523383b9b64022e341e53e009a62ebc47a45f2"
 373 | dependencies = [
 374 |  "clap_builder",
 375 |  "clap_derive",
 376 | ]
 377 | 
 378 | [[package]]
 379 | name = "clap_builder"
 380 | version = "4.4.11"
 381 | source = "registry+https://github.com/rust-lang/crates.io-index"
 382 | checksum = "a216b506622bb1d316cd51328dce24e07bdff4a6128a47c7e7fad11878d5adbb"
 383 | dependencies = [
 384 |  "anstream",
 385 |  "anstyle",
 386 |  "clap_lex 0.6.0",
 387 |  "strsim",
 388 | ]
 389 | 
 390 | [[package]]
 391 | name = "clap_derive"
 392 | version = "4.4.7"
 393 | source = "registry+https://github.com/rust-lang/crates.io-index"
 394 | checksum = "cf9804afaaf59a91e75b022a30fb7229a7901f60c755489cc61c9b423b836442"
 395 | dependencies = [
 396 |  "heck",
 397 |  "proc-macro2",
 398 |  "quote",
 399 |  "syn 2.0.39",
 400 | ]
 401 | 
 402 | [[package]]
 403 | name = "clap_lex"
 404 | version = "0.2.4"
 405 | source = "registry+https://github.com/rust-lang/crates.io-index"
 406 | checksum = "2850f2f5a82cbf437dd5af4d49848fbdfc27c157c3d010345776f952765261c5"
 407 | dependencies = [
 408 |  "os_str_bytes",
 409 | ]
 410 | 
 411 | [[package]]
 412 | name = "clap_lex"
 413 | version = "0.6.0"
 414 | source = "registry+https://github.com/rust-lang/crates.io-index"
 415 | checksum = "702fc72eb24e5a1e48ce58027a675bc24edd52096d5397d4aea7c6dd9eca0bd1"
 416 | 
 417 | [[package]]
 418 | name = "colorchoice"
 419 | version = "1.0.0"
 420 | source = "registry+https://github.com/rust-lang/crates.io-index"
 421 | checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7"
 422 | 
 423 | [[package]]
 424 | name = "console"
 425 | version = "0.15.7"
 426 | source = "registry+https://github.com/rust-lang/crates.io-index"
 427 | checksum = "c926e00cc70edefdc64d3a5ff31cc65bb97a3460097762bd23afb4d8145fccf8"
 428 | dependencies = [
 429 |  "encode_unicode",
 430 |  "lazy_static",
 431 |  "libc",
 432 |  "unicode-width",
 433 |  "windows-sys 0.45.0",
 434 | ]
 435 | 
 436 | [[package]]
 437 | name = "core-foundation"
 438 | version = "0.9.4"
 439 | source = "registry+https://github.com/rust-lang/crates.io-index"
 440 | checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f"
 441 | dependencies = [
 442 |  "core-foundation-sys",
 443 |  "libc",
 444 | ]
 445 | 
 446 | [[package]]
 447 | name = "core-foundation-sys"
 448 | version = "0.8.6"
 449 | source = "registry+https://github.com/rust-lang/crates.io-index"
 450 | checksum = "06ea2b9bc92be3c2baa9334a323ebca2d6f074ff852cd1d7b11064035cd3868f"
 451 | 
 452 | [[package]]
 453 | name = "core-graphics-types"
 454 | version = "0.1.3"
 455 | source = "registry+https://github.com/rust-lang/crates.io-index"
 456 | checksum = "45390e6114f68f718cc7a830514a96f903cccd70d02a8f6d9f643ac4ba45afaf"
 457 | dependencies = [
 458 |  "bitflags 1.3.2",
 459 |  "core-foundation",
 460 |  "libc",
 461 | ]
 462 | 
 463 | [[package]]
 464 | name = "cpufeatures"
 465 | version = "0.2.11"
 466 | source = "registry+https://github.com/rust-lang/crates.io-index"
 467 | checksum = "ce420fe07aecd3e67c5f910618fe65e94158f6dcc0adf44e00d69ce2bdfe0fd0"
 468 | dependencies = [
 469 |  "libc",
 470 | ]
 471 | 
 472 | [[package]]
 473 | name = "crc32fast"
 474 | version = "1.3.2"
 475 | source = "registry+https://github.com/rust-lang/crates.io-index"
 476 | checksum = "b540bd8bc810d3885c6ea91e2018302f68baba2129ab3e88f32389ee9370880d"
 477 | dependencies = [
 478 |  "cfg-if",
 479 | ]
 480 | 
 481 | [[package]]
 482 | name = "criterion"
 483 | version = "0.4.0"
 484 | source = "registry+https://github.com/rust-lang/crates.io-index"
 485 | checksum = "e7c76e09c1aae2bc52b3d2f29e13c6572553b30c4aa1b8a49fd70de6412654cb"
 486 | dependencies = [
 487 |  "anes",
 488 |  "atty",
 489 |  "cast",
 490 |  "ciborium",
 491 |  "clap 3.2.25",
 492 |  "criterion-plot",
 493 |  "itertools 0.10.5",
 494 |  "lazy_static",
 495 |  "num-traits",
 496 |  "oorandom",
 497 |  "plotters",
 498 |  "rayon",
 499 |  "regex",
 500 |  "serde",
 501 |  "serde_derive",
 502 |  "serde_json",
 503 |  "tinytemplate",
 504 |  "walkdir",
 505 | ]
 506 | 
 507 | [[package]]
 508 | name = "criterion-plot"
 509 | version = "0.5.0"
 510 | source = "registry+https://github.com/rust-lang/crates.io-index"
 511 | checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1"
 512 | dependencies = [
 513 |  "cast",
 514 |  "itertools 0.10.5",
 515 | ]
 516 | 
 517 | [[package]]
 518 | name = "crossbeam-deque"
 519 | version = "0.8.3"
 520 | source = "registry+https://github.com/rust-lang/crates.io-index"
 521 | checksum = "ce6fd6f855243022dcecf8702fef0c297d4338e226845fe067f6341ad9fa0cef"
 522 | dependencies = [
 523 |  "cfg-if",
 524 |  "crossbeam-epoch",
 525 |  "crossbeam-utils",
 526 | ]
 527 | 
 528 | [[package]]
 529 | name = "crossbeam-epoch"
 530 | version = "0.9.15"
 531 | source = "registry+https://github.com/rust-lang/crates.io-index"
 532 | checksum = "ae211234986c545741a7dc064309f67ee1e5ad243d0e48335adc0484d960bcc7"
 533 | dependencies = [
 534 |  "autocfg",
 535 |  "cfg-if",
 536 |  "crossbeam-utils",
 537 |  "memoffset",
 538 |  "scopeguard",
 539 | ]
 540 | 
 541 | [[package]]
 542 | name = "crossbeam-utils"
 543 | version = "0.8.16"
 544 | source = "registry+https://github.com/rust-lang/crates.io-index"
 545 | checksum = "5a22b2d63d4d1dc0b7f1b6b2747dd0088008a9be28b6ddf0b1e7d335e3037294"
 546 | dependencies = [
 547 |  "cfg-if",
 548 | ]
 549 | 
 550 | [[package]]
 551 | name = "crunchy"
 552 | version = "0.2.2"
 553 | source = "registry+https://github.com/rust-lang/crates.io-index"
 554 | checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7"
 555 | 
 556 | [[package]]
 557 | name = "crypto-common"
 558 | version = "0.1.6"
 559 | source = "registry+https://github.com/rust-lang/crates.io-index"
 560 | checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3"
 561 | dependencies = [
 562 |  "generic-array",
 563 |  "typenum",
 564 | ]
 565 | 
 566 | [[package]]
 567 | name = "cudarc"
 568 | version = "0.10.0"
 569 | source = "registry+https://github.com/rust-lang/crates.io-index"
 570 | checksum = "9395df0cab995685664e79cc35ad6302bf08fb9c5d82301875a183affe1278b1"
 571 | dependencies = [
 572 |  "half 2.3.1",
 573 | ]
 574 | 
 575 | [[package]]
 576 | name = "darling"
 577 | version = "0.14.4"
 578 | source = "registry+https://github.com/rust-lang/crates.io-index"
 579 | checksum = "7b750cb3417fd1b327431a470f388520309479ab0bf5e323505daf0290cd3850"
 580 | dependencies = [
 581 |  "darling_core",
 582 |  "darling_macro",
 583 | ]
 584 | 
 585 | [[package]]
 586 | name = "darling_core"
 587 | version = "0.14.4"
 588 | source = "registry+https://github.com/rust-lang/crates.io-index"
 589 | checksum = "109c1ca6e6b7f82cc233a97004ea8ed7ca123a9af07a8230878fcfda9b158bf0"
 590 | dependencies = [
 591 |  "fnv",
 592 |  "ident_case",
 593 |  "proc-macro2",
 594 |  "quote",
 595 |  "strsim",
 596 |  "syn 1.0.109",
 597 | ]
 598 | 
 599 | [[package]]
 600 | name = "darling_macro"
 601 | version = "0.14.4"
 602 | source = "registry+https://github.com/rust-lang/crates.io-index"
 603 | checksum = "a4aab4dbc9f7611d8b55048a3a16d2d010c2c8334e46304b40ac1cc14bf3b48e"
 604 | dependencies = [
 605 |  "darling_core",
 606 |  "quote",
 607 |  "syn 1.0.109",
 608 | ]
 609 | 
 610 | [[package]]
 611 | name = "derive_builder"
 612 | version = "0.11.2"
 613 | source = "registry+https://github.com/rust-lang/crates.io-index"
 614 | checksum = "d07adf7be193b71cc36b193d0f5fe60b918a3a9db4dad0449f57bcfd519704a3"
 615 | dependencies = [
 616 |  "derive_builder_macro 0.11.2",
 617 | ]
 618 | 
 619 | [[package]]
 620 | name = "derive_builder"
 621 | version = "0.12.0"
 622 | source = "registry+https://github.com/rust-lang/crates.io-index"
 623 | checksum = "8d67778784b508018359cbc8696edb3db78160bab2c2a28ba7f56ef6932997f8"
 624 | dependencies = [
 625 |  "derive_builder_macro 0.12.0",
 626 | ]
 627 | 
 628 | [[package]]
 629 | name = "derive_builder_core"
 630 | version = "0.11.2"
 631 | source = "registry+https://github.com/rust-lang/crates.io-index"
 632 | checksum = "1f91d4cfa921f1c05904dc3c57b4a32c38aed3340cce209f3a6fd1478babafc4"
 633 | dependencies = [
 634 |  "darling",
 635 |  "proc-macro2",
 636 |  "quote",
 637 |  "syn 1.0.109",
 638 | ]
 639 | 
 640 | [[package]]
 641 | name = "derive_builder_core"
 642 | version = "0.12.0"
 643 | source = "registry+https://github.com/rust-lang/crates.io-index"
 644 | checksum = "c11bdc11a0c47bc7d37d582b5285da6849c96681023680b906673c5707af7b0f"
 645 | dependencies = [
 646 |  "darling",
 647 |  "proc-macro2",
 648 |  "quote",
 649 |  "syn 1.0.109",
 650 | ]
 651 | 
 652 | [[package]]
 653 | name = "derive_builder_macro"
 654 | version = "0.11.2"
 655 | source = "registry+https://github.com/rust-lang/crates.io-index"
 656 | checksum = "8f0314b72bed045f3a68671b3c86328386762c93f82d98c65c3cb5e5f573dd68"
 657 | dependencies = [
 658 |  "derive_builder_core 0.11.2",
 659 |  "syn 1.0.109",
 660 | ]
 661 | 
 662 | [[package]]
 663 | name = "derive_builder_macro"
 664 | version = "0.12.0"
 665 | source = "registry+https://github.com/rust-lang/crates.io-index"
 666 | checksum = "ebcda35c7a396850a55ffeac740804b40ffec779b98fffbb1738f4033f0ee79e"
 667 | dependencies = [
 668 |  "derive_builder_core 0.12.0",
 669 |  "syn 1.0.109",
 670 | ]
 671 | 
 672 | [[package]]
 673 | name = "digest"
 674 | version = "0.10.7"
 675 | source = "registry+https://github.com/rust-lang/crates.io-index"
 676 | checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
 677 | dependencies = [
 678 |  "block-buffer",
 679 |  "crypto-common",
 680 | ]
 681 | 
 682 | [[package]]
 683 | name = "directories"
 684 | version = "4.0.1"
 685 | source = "registry+https://github.com/rust-lang/crates.io-index"
 686 | checksum = "f51c5d4ddabd36886dd3e1438cb358cdcb0d7c499cb99cb4ac2e38e18b5cb210"
 687 | dependencies = [
 688 |  "dirs-sys",
 689 | ]
 690 | 
 691 | [[package]]
 692 | name = "dirs-sys"
 693 | version = "0.3.7"
 694 | source = "registry+https://github.com/rust-lang/crates.io-index"
 695 | checksum = "1b1d1d91c932ef41c0f2663aa8b0ca0342d444d842c06914aa0a7e352d0bada6"
 696 | dependencies = [
 697 |  "libc",
 698 |  "redox_users",
 699 |  "winapi",
 700 | ]
 701 | 
 702 | [[package]]
 703 | name = "dyn-stack"
 704 | version = "0.10.0"
 705 | source = "registry+https://github.com/rust-lang/crates.io-index"
 706 | checksum = "56e53799688f5632f364f8fb387488dd05db9fe45db7011be066fc20e7027f8b"
 707 | dependencies = [
 708 |  "bytemuck",
 709 |  "reborrow",
 710 | ]
 711 | 
 712 | [[package]]
 713 | name = "either"
 714 | version = "1.9.0"
 715 | source = "registry+https://github.com/rust-lang/crates.io-index"
 716 | checksum = "a26ae43d7bcc3b814de94796a5e736d4029efb0ee900c12e2d54c993ad1a1e07"
 717 | 
 718 | [[package]]
 719 | name = "encode_unicode"
 720 | version = "0.3.6"
 721 | source = "registry+https://github.com/rust-lang/crates.io-index"
 722 | checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f"
 723 | 
 724 | [[package]]
 725 | name = "enum-as-inner"
 726 | version = "0.6.0"
 727 | source = "registry+https://github.com/rust-lang/crates.io-index"
 728 | checksum = "5ffccbb6966c05b32ef8fbac435df276c4ae4d3dc55a8cd0eb9745e6c12f546a"
 729 | dependencies = [
 730 |  "heck",
 731 |  "proc-macro2",
 732 |  "quote",
 733 |  "syn 2.0.39",
 734 | ]
 735 | 
 736 | [[package]]
 737 | name = "errno"
 738 | version = "0.3.8"
 739 | source = "registry+https://github.com/rust-lang/crates.io-index"
 740 | checksum = "a258e46cdc063eb8519c00b9fc845fc47bcfca4130e2f08e88665ceda8474245"
 741 | dependencies = [
 742 |  "libc",
 743 |  "windows-sys 0.52.0",
 744 | ]
 745 | 
 746 | [[package]]
 747 | name = "esaxx-rs"
 748 | version = "0.1.10"
 749 | source = "registry+https://github.com/rust-lang/crates.io-index"
 750 | checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6"
 751 | dependencies = [
 752 |  "cc",
 753 | ]
 754 | 
 755 | [[package]]
 756 | name = "filetime"
 757 | version = "0.2.23"
 758 | source = "registry+https://github.com/rust-lang/crates.io-index"
 759 | checksum = "1ee447700ac8aa0b2f2bd7bc4462ad686ba06baa6727ac149a2d6277f0d240fd"
 760 | dependencies = [
 761 |  "cfg-if",
 762 |  "libc",
 763 |  "redox_syscall",
 764 |  "windows-sys 0.52.0",
 765 | ]
 766 | 
 767 | [[package]]
 768 | name = "flate2"
 769 | version = "1.0.28"
 770 | source = "registry+https://github.com/rust-lang/crates.io-index"
 771 | checksum = "46303f565772937ffe1d394a4fac6f411c6013172fadde9dcdb1e147a086940e"
 772 | dependencies = [
 773 |  "crc32fast",
 774 |  "miniz_oxide",
 775 | ]
 776 | 
 777 | [[package]]
 778 | name = "fnv"
 779 | version = "1.0.7"
 780 | source = "registry+https://github.com/rust-lang/crates.io-index"
 781 | checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
 782 | 
 783 | [[package]]
 784 | name = "foreign-types"
 785 | version = "0.5.0"
 786 | source = "registry+https://github.com/rust-lang/crates.io-index"
 787 | checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965"
 788 | dependencies = [
 789 |  "foreign-types-macros",
 790 |  "foreign-types-shared",
 791 | ]
 792 | 
 793 | [[package]]
 794 | name = "foreign-types-macros"
 795 | version = "0.2.3"
 796 | source = "registry+https://github.com/rust-lang/crates.io-index"
 797 | checksum = "1a5c6c585bc94aaf2c7b51dd4c2ba22680844aba4c687be581871a6f518c5742"
 798 | dependencies = [
 799 |  "proc-macro2",
 800 |  "quote",
 801 |  "syn 2.0.39",
 802 | ]
 803 | 
 804 | [[package]]
 805 | name = "foreign-types-shared"
 806 | version = "0.3.1"
 807 | source = "registry+https://github.com/rust-lang/crates.io-index"
 808 | checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b"
 809 | 
 810 | [[package]]
 811 | name = "form_urlencoded"
 812 | version = "1.2.1"
 813 | source = "registry+https://github.com/rust-lang/crates.io-index"
 814 | checksum = "e13624c2627564efccf4934284bdd98cbaa14e79b0b5a141218e507b3a823456"
 815 | dependencies = [
 816 |  "percent-encoding",
 817 | ]
 818 | 
 819 | [[package]]
 820 | name = "gemm"
 821 | version = "0.17.0"
 822 | source = "registry+https://github.com/rust-lang/crates.io-index"
 823 | checksum = "e97d506c68f4fb12325b52a638e7d54cc87e3593a4ded0de60218b6dfd65f645"
 824 | dependencies = [
 825 |  "dyn-stack",
 826 |  "gemm-c32",
 827 |  "gemm-c64",
 828 |  "gemm-common",
 829 |  "gemm-f16",
 830 |  "gemm-f32",
 831 |  "gemm-f64",
 832 |  "num-complex",
 833 |  "num-traits",
 834 |  "paste",
 835 |  "raw-cpuid",
 836 |  "seq-macro",
 837 | ]
 838 | 
 839 | [[package]]
 840 | name = "gemm-c32"
 841 | version = "0.17.0"
 842 | source = "registry+https://github.com/rust-lang/crates.io-index"
 843 | checksum = "0dd16f26e8f34661edc906d8c9522b59ec1655c865a98a58950d0246eeaca9da"
 844 | dependencies = [
 845 |  "dyn-stack",
 846 |  "gemm-common",
 847 |  "num-complex",
 848 |  "num-traits",
 849 |  "paste",
 850 |  "raw-cpuid",
 851 |  "seq-macro",
 852 | ]
 853 | 
 854 | [[package]]
 855 | name = "gemm-c64"
 856 | version = "0.17.0"
 857 | source = "registry+https://github.com/rust-lang/crates.io-index"
 858 | checksum = "a8e34381bc060b47fbd25522a281799ef763cd27f43bbd1783d935774659242a"
 859 | dependencies = [
 860 |  "dyn-stack",
 861 |  "gemm-common",
 862 |  "num-complex",
 863 |  "num-traits",
 864 |  "paste",
 865 |  "raw-cpuid",
 866 |  "seq-macro",
 867 | ]
 868 | 
 869 | [[package]]
 870 | name = "gemm-common"
 871 | version = "0.17.0"
 872 | source = "registry+https://github.com/rust-lang/crates.io-index"
 873 | checksum = "22518a76339b09276f77c3166c44262e55f633712fe8a44fd0573505887feeab"
 874 | dependencies = [
 875 |  "bytemuck",
 876 |  "dyn-stack",
 877 |  "half 2.3.1",
 878 |  "num-complex",
 879 |  "num-traits",
 880 |  "once_cell",
 881 |  "paste",
 882 |  "pulp",
 883 |  "raw-cpuid",
 884 |  "rayon",
 885 |  "seq-macro",
 886 |  "sysctl",
 887 | ]
 888 | 
 889 | [[package]]
 890 | name = "gemm-f16"
 891 | version = "0.17.0"
 892 | source = "registry+https://github.com/rust-lang/crates.io-index"
 893 | checksum = "70409bbf3ef83b38cbe4a58cd4b797c1c27902505bdd926a588ea61b6c550a84"
 894 | dependencies = [
 895 |  "dyn-stack",
 896 |  "gemm-common",
 897 |  "gemm-f32",
 898 |  "half 2.3.1",
 899 |  "num-complex",
 900 |  "num-traits",
 901 |  "paste",
 902 |  "raw-cpuid",
 903 |  "rayon",
 904 |  "seq-macro",
 905 | ]
 906 | 
 907 | [[package]]
 908 | name = "gemm-f32"
 909 | version = "0.17.0"
 910 | source = "registry+https://github.com/rust-lang/crates.io-index"
 911 | checksum = "5ea3068edca27f100964157211782eba19e961aa4d0d2bdac3e1775a51aa7680"
 912 | dependencies = [
 913 |  "dyn-stack",
 914 |  "gemm-common",
 915 |  "num-complex",
 916 |  "num-traits",
 917 |  "paste",
 918 |  "raw-cpuid",
 919 |  "seq-macro",
 920 | ]
 921 | 
 922 | [[package]]
 923 | name = "gemm-f64"
 924 | version = "0.17.0"
 925 | source = "registry+https://github.com/rust-lang/crates.io-index"
 926 | checksum = "5fd41e8f5a60dce8d8acd852a3f4b22f8e18be957e1937731be692c037652510"
 927 | dependencies = [
 928 |  "dyn-stack",
 929 |  "gemm-common",
 930 |  "num-complex",
 931 |  "num-traits",
 932 |  "paste",
 933 |  "raw-cpuid",
 934 |  "seq-macro",
 935 | ]
 936 | 
 937 | [[package]]
 938 | name = "generic-array"
 939 | version = "0.14.7"
 940 | source = "registry+https://github.com/rust-lang/crates.io-index"
 941 | checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a"
 942 | dependencies = [
 943 |  "typenum",
 944 |  "version_check",
 945 | ]
 946 | 
 947 | [[package]]
 948 | name = "getrandom"
 949 | version = "0.2.11"
 950 | source = "registry+https://github.com/rust-lang/crates.io-index"
 951 | checksum = "fe9006bed769170c11f845cf00c7c1e9092aeb3f268e007c3e760ac68008070f"
 952 | dependencies = [
 953 |  "cfg-if",
 954 |  "libc",
 955 |  "wasi",
 956 | ]
 957 | 
 958 | [[package]]
 959 | name = "getset"
 960 | version = "0.1.2"
 961 | source = "registry+https://github.com/rust-lang/crates.io-index"
 962 | checksum = "e45727250e75cc04ff2846a66397da8ef2b3db8e40e0cef4df67950a07621eb9"
 963 | dependencies = [
 964 |  "proc-macro-error",
 965 |  "proc-macro2",
 966 |  "quote",
 967 |  "syn 1.0.109",
 968 | ]
 969 | 
 970 | [[package]]
 971 | name = "glob"
 972 | version = "0.3.1"
 973 | source = "registry+https://github.com/rust-lang/crates.io-index"
 974 | checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b"
 975 | 
 976 | [[package]]
 977 | name = "half"
 978 | version = "1.8.2"
 979 | source = "registry+https://github.com/rust-lang/crates.io-index"
 980 | checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7"
 981 | 
 982 | [[package]]
 983 | name = "half"
 984 | version = "2.3.1"
 985 | source = "registry+https://github.com/rust-lang/crates.io-index"
 986 | checksum = "bc52e53916c08643f1b56ec082790d1e86a32e58dc5268f897f313fbae7b4872"
 987 | dependencies = [
 988 |  "bytemuck",
 989 |  "cfg-if",
 990 |  "crunchy",
 991 |  "num-traits",
 992 |  "rand",
 993 |  "rand_distr",
 994 | ]
 995 | 
 996 | [[package]]
 997 | name = "hashbrown"
 998 | version = "0.12.3"
 999 | source = "registry+https://github.com/rust-lang/crates.io-index"
1000 | checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
1001 | 
1002 | [[package]]
1003 | name = "heck"
1004 | version = "0.4.1"
1005 | source = "registry+https://github.com/rust-lang/crates.io-index"
1006 | checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8"
1007 | 
1008 | [[package]]
1009 | name = "hermit-abi"
1010 | version = "0.1.19"
1011 | source = "registry+https://github.com/rust-lang/crates.io-index"
1012 | checksum = "62b467343b94ba476dcb2500d242dadbb39557df889310ac77c5d99100aaac33"
1013 | dependencies = [
1014 |  "libc",
1015 | ]
1016 | 
1017 | [[package]]
1018 | name = "hermit-abi"
1019 | version = "0.3.3"
1020 | source = "registry+https://github.com/rust-lang/crates.io-index"
1021 | checksum = "d77f7ec81a6d05a3abb01ab6eb7590f6083d08449fe5a1c8b1e620283546ccb7"
1022 | 
1023 | [[package]]
1024 | name = "iana-time-zone"
1025 | version = "0.1.58"
1026 | source = "registry+https://github.com/rust-lang/crates.io-index"
1027 | checksum = "8326b86b6cff230b97d0d312a6c40a60726df3332e721f72a1b035f451663b20"
1028 | dependencies = [
1029 |  "android_system_properties",
1030 |  "core-foundation-sys",
1031 |  "iana-time-zone-haiku",
1032 |  "js-sys",
1033 |  "wasm-bindgen",
1034 |  "windows-core",
1035 | ]
1036 | 
1037 | [[package]]
1038 | name = "iana-time-zone-haiku"
1039 | version = "0.1.2"
1040 | source = "registry+https://github.com/rust-lang/crates.io-index"
1041 | checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
1042 | dependencies = [
1043 |  "cc",
1044 | ]
1045 | 
1046 | [[package]]
1047 | name = "ident_case"
1048 | version = "1.0.1"
1049 | source = "registry+https://github.com/rust-lang/crates.io-index"
1050 | checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39"
1051 | 
1052 | [[package]]
1053 | name = "idna"
1054 | version = "0.5.0"
1055 | source = "registry+https://github.com/rust-lang/crates.io-index"
1056 | checksum = "634d9b1461af396cad843f47fdba5597a4f9e6ddd4bfb6ff5d85028c25cb12f6"
1057 | dependencies = [
1058 |  "unicode-bidi",
1059 |  "unicode-normalization",
1060 | ]
1061 | 
1062 | [[package]]
1063 | name = "indexmap"
1064 | version = "1.9.3"
1065 | source = "registry+https://github.com/rust-lang/crates.io-index"
1066 | checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99"
1067 | dependencies = [
1068 |  "autocfg",
1069 |  "hashbrown",
1070 | ]
1071 | 
1072 | [[package]]
1073 | name = "indicatif"
1074 | version = "0.17.7"
1075 | source = "registry+https://github.com/rust-lang/crates.io-index"
1076 | checksum = "fb28741c9db9a713d93deb3bb9515c20788cef5815265bee4980e87bde7e0f25"
1077 | dependencies = [
1078 |  "console",
1079 |  "instant",
1080 |  "number_prefix",
1081 |  "portable-atomic",
1082 |  "unicode-width",
1083 | ]
1084 | 
1085 | [[package]]
1086 | name = "instant"
1087 | version = "0.1.12"
1088 | source = "registry+https://github.com/rust-lang/crates.io-index"
1089 | checksum = "7a5bbe824c507c5da5956355e86a746d82e0e1464f65d862cc5e71da70e94b2c"
1090 | dependencies = [
1091 |  "cfg-if",
1092 | ]
1093 | 
1094 | [[package]]
1095 | name = "intel-mkl-src"
1096 | version = "0.8.1"
1097 | source = "registry+https://github.com/rust-lang/crates.io-index"
1098 | checksum = "2ee70586cd5b3e772a8739a1bd43eaa90d4f4bf0fb2a4edc202e979937ee7f5e"
1099 | dependencies = [
1100 |  "anyhow",
1101 |  "intel-mkl-tool",
1102 |  "ocipkg",
1103 | ]
1104 | 
1105 | [[package]]
1106 | name = "intel-mkl-tool"
1107 | version = "0.8.1"
1108 | source = "registry+https://github.com/rust-lang/crates.io-index"
1109 | checksum = "887a16b4537d82227af54d3372971cfa5e0cde53322e60f57584056c16ada1b4"
1110 | dependencies = [
1111 |  "anyhow",
1112 |  "log",
1113 |  "walkdir",
1114 | ]
1115 | 
1116 | [[package]]
1117 | name = "itertools"
1118 | version = "0.10.5"
1119 | source = "registry+https://github.com/rust-lang/crates.io-index"
1120 | checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473"
1121 | dependencies = [
1122 |  "either",
1123 | ]
1124 | 
1125 | [[package]]
1126 | name = "itertools"
1127 | version = "0.11.0"
1128 | source = "registry+https://github.com/rust-lang/crates.io-index"
1129 | checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57"
1130 | dependencies = [
1131 |  "either",
1132 | ]
1133 | 
1134 | [[package]]
1135 | name = "itoa"
1136 | version = "1.0.9"
1137 | source = "registry+https://github.com/rust-lang/crates.io-index"
1138 | checksum = "af150ab688ff2122fcef229be89cb50dd66af9e01a4ff320cc137eecc9bacc38"
1139 | 
1140 | [[package]]
1141 | name = "js-sys"
1142 | version = "0.3.66"
1143 | source = "registry+https://github.com/rust-lang/crates.io-index"
1144 | checksum = "cee9c64da59eae3b50095c18d3e74f8b73c0b86d2792824ff01bbce68ba229ca"
1145 | dependencies = [
1146 |  "wasm-bindgen",
1147 | ]
1148 | 
1149 | [[package]]
1150 | name = "lazy_static"
1151 | version = "1.4.0"
1152 | source = "registry+https://github.com/rust-lang/crates.io-index"
1153 | checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646"
1154 | 
1155 | [[package]]
1156 | name = "libc"
1157 | version = "0.2.153"
1158 | source = "registry+https://github.com/rust-lang/crates.io-index"
1159 | checksum = "9c198f91728a82281a64e1f4f9eeb25d82cb32a5de251c6bd1b5154d63a8e7bd"
1160 | 
1161 | [[package]]
1162 | name = "libm"
1163 | version = "0.2.8"
1164 | source = "registry+https://github.com/rust-lang/crates.io-index"
1165 | checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058"
1166 | 
1167 | [[package]]
1168 | name = "libredox"
1169 | version = "0.0.1"
1170 | source = "registry+https://github.com/rust-lang/crates.io-index"
1171 | checksum = "85c833ca1e66078851dba29046874e38f08b2c883700aa29a03ddd3b23814ee8"
1172 | dependencies = [
1173 |  "bitflags 2.4.1",
1174 |  "libc",
1175 |  "redox_syscall",
1176 | ]
1177 | 
1178 | [[package]]
1179 | name = "linux-raw-sys"
1180 | version = "0.4.12"
1181 | source = "registry+https://github.com/rust-lang/crates.io-index"
1182 | checksum = "c4cd1a83af159aa67994778be9070f0ae1bd732942279cabb14f86f986a21456"
1183 | 
1184 | [[package]]
1185 | name = "log"
1186 | version = "0.4.20"
1187 | source = "registry+https://github.com/rust-lang/crates.io-index"
1188 | checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f"
1189 | 
1190 | [[package]]
1191 | name = "macro_rules_attribute"
1192 | version = "0.2.0"
1193 | source = "registry+https://github.com/rust-lang/crates.io-index"
1194 | checksum = "8a82271f7bc033d84bbca59a3ce3e4159938cb08a9c3aebbe54d215131518a13"
1195 | dependencies = [
1196 |  "macro_rules_attribute-proc_macro",
1197 |  "paste",
1198 | ]
1199 | 
1200 | [[package]]
1201 | name = "macro_rules_attribute-proc_macro"
1202 | version = "0.2.0"
1203 | source = "registry+https://github.com/rust-lang/crates.io-index"
1204 | checksum = "b8dd856d451cc0da70e2ef2ce95a18e39a93b7558bedf10201ad28503f918568"
1205 | 
1206 | [[package]]
1207 | name = "malloc_buf"
1208 | version = "0.0.6"
1209 | source = "registry+https://github.com/rust-lang/crates.io-index"
1210 | checksum = "62bb907fe88d54d8d9ce32a3cceab4218ed2f6b7d35617cafe9adf84e43919cb"
1211 | dependencies = [
1212 |  "libc",
1213 | ]
1214 | 
1215 | [[package]]
1216 | name = "mamba-ssm"
1217 | version = "0.1.0"
1218 | dependencies = [
1219 |  "accelerate-src",
1220 |  "anyhow",
1221 |  "candle-core",
1222 |  "candle-metal-kernels",
1223 |  "candle-nn",
1224 |  "candle-transformers",
1225 |  "clap 4.4.11",
1226 |  "criterion",
1227 |  "intel-mkl-src",
1228 |  "rand",
1229 |  "rayon",
1230 |  "serde",
1231 |  "serde_json",
1232 |  "tokenizers",
1233 |  "tracing",
1234 |  "tracing-chrome",
1235 |  "tracing-subscriber",
1236 | ]
1237 | 
1238 | [[package]]
1239 | name = "memchr"
1240 | version = "2.6.4"
1241 | source = "registry+https://github.com/rust-lang/crates.io-index"
1242 | checksum = "f665ee40bc4a3c5590afb1e9677db74a508659dfd71e126420da8274909a0167"
1243 | 
1244 | [[package]]
1245 | name = "memmap2"
1246 | version = "0.9.4"
1247 | source = "registry+https://github.com/rust-lang/crates.io-index"
1248 | checksum = "fe751422e4a8caa417e13c3ea66452215d7d63e19e604f4980461212f3ae1322"
1249 | dependencies = [
1250 |  "libc",
1251 |  "stable_deref_trait",
1252 | ]
1253 | 
1254 | [[package]]
1255 | name = "memoffset"
1256 | version = "0.9.0"
1257 | source = "registry+https://github.com/rust-lang/crates.io-index"
1258 | checksum = "5a634b1c61a95585bd15607c6ab0c4e5b226e695ff2800ba0cdccddf208c406c"
1259 | dependencies = [
1260 |  "autocfg",
1261 | ]
1262 | 
1263 | [[package]]
1264 | name = "metal"
1265 | version = "0.27.0"
1266 | source = "registry+https://github.com/rust-lang/crates.io-index"
1267 | checksum = "c43f73953f8cbe511f021b58f18c3ce1c3d1ae13fe953293e13345bf83217f25"
1268 | dependencies = [
1269 |  "bitflags 2.4.1",
1270 |  "block",
1271 |  "core-graphics-types",
1272 |  "foreign-types",
1273 |  "log",
1274 |  "objc",
1275 |  "paste",
1276 | ]
1277 | 
1278 | [[package]]
1279 | name = "minimal-lexical"
1280 | version = "0.2.1"
1281 | source = "registry+https://github.com/rust-lang/crates.io-index"
1282 | checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
1283 | 
1284 | [[package]]
1285 | name = "miniz_oxide"
1286 | version = "0.7.1"
1287 | source = "registry+https://github.com/rust-lang/crates.io-index"
1288 | checksum = "e7810e0be55b428ada41041c41f32c9f1a42817901b4ccf45fa3d4b6561e74c7"
1289 | dependencies = [
1290 |  "adler",
1291 | ]
1292 | 
1293 | [[package]]
1294 | name = "monostate"
1295 | version = "0.1.10"
1296 | source = "registry+https://github.com/rust-lang/crates.io-index"
1297 | checksum = "e404e13820ea0df0eda93aa294e0c80de76a0daa6bec590d376fbec6d7810394"
1298 | dependencies = [
1299 |  "monostate-impl",
1300 |  "serde",
1301 | ]
1302 | 
1303 | [[package]]
1304 | name = "monostate-impl"
1305 | version = "0.1.10"
1306 | source = "registry+https://github.com/rust-lang/crates.io-index"
1307 | checksum = "531c82a934da419bed3da09bd87d6e98c72f8d4aa755427b3b009c2b8b8c433c"
1308 | dependencies = [
1309 |  "proc-macro2",
1310 |  "quote",
1311 |  "syn 2.0.39",
1312 | ]
1313 | 
1314 | [[package]]
1315 | name = "nom"
1316 | version = "7.1.3"
1317 | source = "registry+https://github.com/rust-lang/crates.io-index"
1318 | checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a"
1319 | dependencies = [
1320 |  "memchr",
1321 |  "minimal-lexical",
1322 | ]
1323 | 
1324 | [[package]]
1325 | name = "nu-ansi-term"
1326 | version = "0.46.0"
1327 | source = "registry+https://github.com/rust-lang/crates.io-index"
1328 | checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84"
1329 | dependencies = [
1330 |  "overload",
1331 |  "winapi",
1332 | ]
1333 | 
1334 | [[package]]
1335 | name = "num-complex"
1336 | version = "0.4.4"
1337 | source = "registry+https://github.com/rust-lang/crates.io-index"
1338 | checksum = "1ba157ca0885411de85d6ca030ba7e2a83a28636056c7c699b07c8b6f7383214"
1339 | dependencies = [
1340 |  "bytemuck",
1341 |  "num-traits",
1342 | ]
1343 | 
1344 | [[package]]
1345 | name = "num-traits"
1346 | version = "0.2.17"
1347 | source = "registry+https://github.com/rust-lang/crates.io-index"
1348 | checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c"
1349 | dependencies = [
1350 |  "autocfg",
1351 |  "libm",
1352 | ]
1353 | 
1354 | [[package]]
1355 | name = "num_cpus"
1356 | version = "1.16.0"
1357 | source = "registry+https://github.com/rust-lang/crates.io-index"
1358 | checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43"
1359 | dependencies = [
1360 |  "hermit-abi 0.3.3",
1361 |  "libc",
1362 | ]
1363 | 
1364 | [[package]]
1365 | name = "number_prefix"
1366 | version = "0.4.0"
1367 | source = "registry+https://github.com/rust-lang/crates.io-index"
1368 | checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
1369 | 
1370 | [[package]]
1371 | name = "objc"
1372 | version = "0.2.7"
1373 | source = "registry+https://github.com/rust-lang/crates.io-index"
1374 | checksum = "915b1b472bc21c53464d6c8461c9d3af805ba1ef837e1cac254428f4a77177b1"
1375 | dependencies = [
1376 |  "malloc_buf",
1377 |  "objc_exception",
1378 | ]
1379 | 
1380 | [[package]]
1381 | name = "objc_exception"
1382 | version = "0.1.2"
1383 | source = "registry+https://github.com/rust-lang/crates.io-index"
1384 | checksum = "ad970fb455818ad6cba4c122ad012fae53ae8b4795f86378bce65e4f6bab2ca4"
1385 | dependencies = [
1386 |  "cc",
1387 | ]
1388 | 
1389 | [[package]]
1390 | name = "oci-spec"
1391 | version = "0.5.8"
1392 | source = "registry+https://github.com/rust-lang/crates.io-index"
1393 | checksum = "98135224dd4faeb24c05a2fac911ed53ea6b09ecb09d7cada1cb79963ab2ee34"
1394 | dependencies = [
1395 |  "derive_builder 0.11.2",
1396 |  "getset",
1397 |  "serde",
1398 |  "serde_json",
1399 |  "thiserror",
1400 | ]
1401 | 
1402 | [[package]]
1403 | name = "ocipkg"
1404 | version = "0.2.8"
1405 | source = "registry+https://github.com/rust-lang/crates.io-index"
1406 | checksum = "60cf01280832705a4e4c4d046d9e67a54b900297c69191457a8fc6d198ddefa2"
1407 | dependencies = [
1408 |  "base16ct",
1409 |  "base64 0.13.1",
1410 |  "chrono",
1411 |  "directories",
1412 |  "flate2",
1413 |  "lazy_static",
1414 |  "log",
1415 |  "oci-spec",
1416 |  "regex",
1417 |  "serde",
1418 |  "serde_json",
1419 |  "sha2",
1420 |  "tar",
1421 |  "thiserror",
1422 |  "toml",
1423 |  "ureq",
1424 |  "url",
1425 |  "uuid",
1426 |  "walkdir",
1427 | ]
1428 | 
1429 | [[package]]
1430 | name = "once_cell"
1431 | version = "1.18.0"
1432 | source = "registry+https://github.com/rust-lang/crates.io-index"
1433 | checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d"
1434 | 
1435 | [[package]]
1436 | name = "onig"
1437 | version = "6.4.0"
1438 | source = "registry+https://github.com/rust-lang/crates.io-index"
1439 | checksum = "8c4b31c8722ad9171c6d77d3557db078cab2bd50afcc9d09c8b315c59df8ca4f"
1440 | dependencies = [
1441 |  "bitflags 1.3.2",
1442 |  "libc",
1443 |  "once_cell",
1444 |  "onig_sys",
1445 | ]
1446 | 
1447 | [[package]]
1448 | name = "onig_sys"
1449 | version = "69.8.1"
1450 | source = "registry+https://github.com/rust-lang/crates.io-index"
1451 | checksum = "7b829e3d7e9cc74c7e315ee8edb185bf4190da5acde74afd7fc59c35b1f086e7"
1452 | dependencies = [
1453 |  "cc",
1454 |  "pkg-config",
1455 | ]
1456 | 
1457 | [[package]]
1458 | name = "oorandom"
1459 | version = "11.1.3"
1460 | source = "registry+https://github.com/rust-lang/crates.io-index"
1461 | checksum = "0ab1bc2a289d34bd04a330323ac98a1b4bc82c9d9fcb1e66b63caa84da26b575"
1462 | 
1463 | [[package]]
1464 | name = "os_str_bytes"
1465 | version = "6.6.1"
1466 | source = "registry+https://github.com/rust-lang/crates.io-index"
1467 | checksum = "e2355d85b9a3786f481747ced0e0ff2ba35213a1f9bd406ed906554d7af805a1"
1468 | 
1469 | [[package]]
1470 | name = "overload"
1471 | version = "0.1.1"
1472 | source = "registry+https://github.com/rust-lang/crates.io-index"
1473 | checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39"
1474 | 
1475 | [[package]]
1476 | name = "paste"
1477 | version = "1.0.14"
1478 | source = "registry+https://github.com/rust-lang/crates.io-index"
1479 | checksum = "de3145af08024dea9fa9914f381a17b8fc6034dfb00f3a84013f7ff43f29ed4c"
1480 | 
1481 | [[package]]
1482 | name = "percent-encoding"
1483 | version = "2.3.1"
1484 | source = "registry+https://github.com/rust-lang/crates.io-index"
1485 | checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e"
1486 | 
1487 | [[package]]
1488 | name = "pin-project-lite"
1489 | version = "0.2.13"
1490 | source = "registry+https://github.com/rust-lang/crates.io-index"
1491 | checksum = "8afb450f006bf6385ca15ef45d71d2288452bc3683ce2e2cacc0d18e4be60b58"
1492 | 
1493 | [[package]]
1494 | name = "pkg-config"
1495 | version = "0.3.27"
1496 | source = "registry+https://github.com/rust-lang/crates.io-index"
1497 | checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964"
1498 | 
1499 | [[package]]
1500 | name = "plotters"
1501 | version = "0.3.5"
1502 | source = "registry+https://github.com/rust-lang/crates.io-index"
1503 | checksum = "d2c224ba00d7cadd4d5c660deaf2098e5e80e07846537c51f9cfa4be50c1fd45"
1504 | dependencies = [
1505 |  "num-traits",
1506 |  "plotters-backend",
1507 |  "plotters-svg",
1508 |  "wasm-bindgen",
1509 |  "web-sys",
1510 | ]
1511 | 
1512 | [[package]]
1513 | name = "plotters-backend"
1514 | version = "0.3.5"
1515 | source = "registry+https://github.com/rust-lang/crates.io-index"
1516 | checksum = "9e76628b4d3a7581389a35d5b6e2139607ad7c75b17aed325f210aa91f4a9609"
1517 | 
1518 | [[package]]
1519 | name = "plotters-svg"
1520 | version = "0.3.5"
1521 | source = "registry+https://github.com/rust-lang/crates.io-index"
1522 | checksum = "38f6d39893cca0701371e3c27294f09797214b86f1fb951b89ade8ec04e2abab"
1523 | dependencies = [
1524 |  "plotters-backend",
1525 | ]
1526 | 
1527 | [[package]]
1528 | name = "portable-atomic"
1529 | version = "1.6.0"
1530 | source = "registry+https://github.com/rust-lang/crates.io-index"
1531 | checksum = "7170ef9988bc169ba16dd36a7fa041e5c4cbeb6a35b76d4c03daded371eae7c0"
1532 | 
1533 | [[package]]
1534 | name = "ppv-lite86"
1535 | version = "0.2.17"
1536 | source = "registry+https://github.com/rust-lang/crates.io-index"
1537 | checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de"
1538 | 
1539 | [[package]]
1540 | name = "proc-macro-error"
1541 | version = "1.0.4"
1542 | source = "registry+https://github.com/rust-lang/crates.io-index"
1543 | checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c"
1544 | dependencies = [
1545 |  "proc-macro-error-attr",
1546 |  "proc-macro2",
1547 |  "quote",
1548 |  "syn 1.0.109",
1549 |  "version_check",
1550 | ]
1551 | 
1552 | [[package]]
1553 | name = "proc-macro-error-attr"
1554 | version = "1.0.4"
1555 | source = "registry+https://github.com/rust-lang/crates.io-index"
1556 | checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869"
1557 | dependencies = [
1558 |  "proc-macro2",
1559 |  "quote",
1560 |  "version_check",
1561 | ]
1562 | 
1563 | [[package]]
1564 | name = "proc-macro2"
1565 | version = "1.0.70"
1566 | source = "registry+https://github.com/rust-lang/crates.io-index"
1567 | checksum = "39278fbbf5fb4f646ce651690877f89d1c5811a3d4acb27700c1cb3cdb78fd3b"
1568 | dependencies = [
1569 |  "unicode-ident",
1570 | ]
1571 | 
1572 | [[package]]
1573 | name = "pulp"
1574 | version = "0.18.6"
1575 | source = "registry+https://github.com/rust-lang/crates.io-index"
1576 | checksum = "16785ee69419641c75affff7c9fdbdb7c0ab26dc9a5fb5218c2a2e9e4ef2087d"
1577 | dependencies = [
1578 |  "bytemuck",
1579 |  "libm",
1580 |  "num-complex",
1581 |  "reborrow",
1582 | ]
1583 | 
1584 | [[package]]
1585 | name = "quote"
1586 | version = "1.0.33"
1587 | source = "registry+https://github.com/rust-lang/crates.io-index"
1588 | checksum = "5267fca4496028628a95160fc423a33e8b2e6af8a5302579e322e4b520293cae"
1589 | dependencies = [
1590 |  "proc-macro2",
1591 | ]
1592 | 
1593 | [[package]]
1594 | name = "rand"
1595 | version = "0.8.5"
1596 | source = "registry+https://github.com/rust-lang/crates.io-index"
1597 | checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
1598 | dependencies = [
1599 |  "libc",
1600 |  "rand_chacha",
1601 |  "rand_core",
1602 | ]
1603 | 
1604 | [[package]]
1605 | name = "rand_chacha"
1606 | version = "0.3.1"
1607 | source = "registry+https://github.com/rust-lang/crates.io-index"
1608 | checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
1609 | dependencies = [
1610 |  "ppv-lite86",
1611 |  "rand_core",
1612 | ]
1613 | 
1614 | [[package]]
1615 | name = "rand_core"
1616 | version = "0.6.4"
1617 | source = "registry+https://github.com/rust-lang/crates.io-index"
1618 | checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
1619 | dependencies = [
1620 |  "getrandom",
1621 | ]
1622 | 
1623 | [[package]]
1624 | name = "rand_distr"
1625 | version = "0.4.3"
1626 | source = "registry+https://github.com/rust-lang/crates.io-index"
1627 | checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31"
1628 | dependencies = [
1629 |  "num-traits",
1630 |  "rand",
1631 | ]
1632 | 
1633 | [[package]]
1634 | name = "raw-cpuid"
1635 | version = "10.7.0"
1636 | source = "registry+https://github.com/rust-lang/crates.io-index"
1637 | checksum = "6c297679cb867470fa8c9f67dbba74a78d78e3e98d7cf2b08d6d71540f797332"
1638 | dependencies = [
1639 |  "bitflags 1.3.2",
1640 | ]
1641 | 
1642 | [[package]]
1643 | name = "rayon"
1644 | version = "1.8.0"
1645 | source = "registry+https://github.com/rust-lang/crates.io-index"
1646 | checksum = "9c27db03db7734835b3f53954b534c91069375ce6ccaa2e065441e07d9b6cdb1"
1647 | dependencies = [
1648 |  "either",
1649 |  "rayon-core",
1650 | ]
1651 | 
1652 | [[package]]
1653 | name = "rayon-cond"
1654 | version = "0.3.0"
1655 | source = "registry+https://github.com/rust-lang/crates.io-index"
1656 | checksum = "059f538b55efd2309c9794130bc149c6a553db90e9d99c2030785c82f0bd7df9"
1657 | dependencies = [
1658 |  "either",
1659 |  "itertools 0.11.0",
1660 |  "rayon",
1661 | ]
1662 | 
1663 | [[package]]
1664 | name = "rayon-core"
1665 | version = "1.12.0"
1666 | source = "registry+https://github.com/rust-lang/crates.io-index"
1667 | checksum = "5ce3fb6ad83f861aac485e76e1985cd109d9a3713802152be56c3b1f0e0658ed"
1668 | dependencies = [
1669 |  "crossbeam-deque",
1670 |  "crossbeam-utils",
1671 | ]
1672 | 
1673 | [[package]]
1674 | name = "reborrow"
1675 | version = "0.5.5"
1676 | source = "registry+https://github.com/rust-lang/crates.io-index"
1677 | checksum = "03251193000f4bd3b042892be858ee50e8b3719f2b08e5833ac4353724632430"
1678 | 
1679 | [[package]]
1680 | name = "redox_syscall"
1681 | version = "0.4.1"
1682 | source = "registry+https://github.com/rust-lang/crates.io-index"
1683 | checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa"
1684 | dependencies = [
1685 |  "bitflags 1.3.2",
1686 | ]
1687 | 
1688 | [[package]]
1689 | name = "redox_users"
1690 | version = "0.4.4"
1691 | source = "registry+https://github.com/rust-lang/crates.io-index"
1692 | checksum = "a18479200779601e498ada4e8c1e1f50e3ee19deb0259c25825a98b5603b2cb4"
1693 | dependencies = [
1694 |  "getrandom",
1695 |  "libredox",
1696 |  "thiserror",
1697 | ]
1698 | 
1699 | [[package]]
1700 | name = "regex"
1701 | version = "1.10.2"
1702 | source = "registry+https://github.com/rust-lang/crates.io-index"
1703 | checksum = "380b951a9c5e80ddfd6136919eef32310721aa4aacd4889a8d39124b026ab343"
1704 | dependencies = [
1705 |  "aho-corasick",
1706 |  "memchr",
1707 |  "regex-automata",
1708 |  "regex-syntax 0.8.2",
1709 | ]
1710 | 
1711 | [[package]]
1712 | name = "regex-automata"
1713 | version = "0.4.3"
1714 | source = "registry+https://github.com/rust-lang/crates.io-index"
1715 | checksum = "5f804c7828047e88b2d32e2d7fe5a105da8ee3264f01902f796c8e067dc2483f"
1716 | dependencies = [
1717 |  "aho-corasick",
1718 |  "memchr",
1719 |  "regex-syntax 0.8.2",
1720 | ]
1721 | 
1722 | [[package]]
1723 | name = "regex-syntax"
1724 | version = "0.7.5"
1725 | source = "registry+https://github.com/rust-lang/crates.io-index"
1726 | checksum = "dbb5fb1acd8a1a18b3dd5be62d25485eb770e05afb408a9627d14d451bae12da"
1727 | 
1728 | [[package]]
1729 | name = "regex-syntax"
1730 | version = "0.8.2"
1731 | source = "registry+https://github.com/rust-lang/crates.io-index"
1732 | checksum = "c08c74e62047bb2de4ff487b251e4a92e24f48745648451635cec7d591162d9f"
1733 | 
1734 | [[package]]
1735 | name = "ring"
1736 | version = "0.17.7"
1737 | source = "registry+https://github.com/rust-lang/crates.io-index"
1738 | checksum = "688c63d65483050968b2a8937f7995f443e27041a0f7700aa59b0822aedebb74"
1739 | dependencies = [
1740 |  "cc",
1741 |  "getrandom",
1742 |  "libc",
1743 |  "spin",
1744 |  "untrusted",
1745 |  "windows-sys 0.48.0",
1746 | ]
1747 | 
1748 | [[package]]
1749 | name = "rustix"
1750 | version = "0.38.28"
1751 | source = "registry+https://github.com/rust-lang/crates.io-index"
1752 | checksum = "72e572a5e8ca657d7366229cdde4bd14c4eb5499a9573d4d366fe1b599daa316"
1753 | dependencies = [
1754 |  "bitflags 2.4.1",
1755 |  "errno",
1756 |  "libc",
1757 |  "linux-raw-sys",
1758 |  "windows-sys 0.52.0",
1759 | ]
1760 | 
1761 | [[package]]
1762 | name = "rustls"
1763 | version = "0.21.10"
1764 | source = "registry+https://github.com/rust-lang/crates.io-index"
1765 | checksum = "f9d5a6813c0759e4609cd494e8e725babae6a2ca7b62a5536a13daaec6fcb7ba"
1766 | dependencies = [
1767 |  "log",
1768 |  "ring",
1769 |  "rustls-webpki",
1770 |  "sct",
1771 | ]
1772 | 
1773 | [[package]]
1774 | name = "rustls-webpki"
1775 | version = "0.101.7"
1776 | source = "registry+https://github.com/rust-lang/crates.io-index"
1777 | checksum = "8b6275d1ee7a1cd780b64aca7726599a1dbc893b1e64144529e55c3c2f745765"
1778 | dependencies = [
1779 |  "ring",
1780 |  "untrusted",
1781 | ]
1782 | 
1783 | [[package]]
1784 | name = "ryu"
1785 | version = "1.0.15"
1786 | source = "registry+https://github.com/rust-lang/crates.io-index"
1787 | checksum = "1ad4cc8da4ef723ed60bced201181d83791ad433213d8c24efffda1eec85d741"
1788 | 
1789 | [[package]]
1790 | name = "safetensors"
1791 | version = "0.4.2"
1792 | source = "registry+https://github.com/rust-lang/crates.io-index"
1793 | checksum = "8d980e6bfb34436fb0a81e42bc41af43f11805bbbca443e7f68e9faaabe669ed"
1794 | dependencies = [
1795 |  "serde",
1796 |  "serde_json",
1797 | ]
1798 | 
1799 | [[package]]
1800 | name = "same-file"
1801 | version = "1.0.6"
1802 | source = "registry+https://github.com/rust-lang/crates.io-index"
1803 | checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
1804 | dependencies = [
1805 |  "winapi-util",
1806 | ]
1807 | 
1808 | [[package]]
1809 | name = "scopeguard"
1810 | version = "1.2.0"
1811 | source = "registry+https://github.com/rust-lang/crates.io-index"
1812 | checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
1813 | 
1814 | [[package]]
1815 | name = "sct"
1816 | version = "0.7.1"
1817 | source = "registry+https://github.com/rust-lang/crates.io-index"
1818 | checksum = "da046153aa2352493d6cb7da4b6e5c0c057d8a1d0a9aa8560baffdd945acd414"
1819 | dependencies = [
1820 |  "ring",
1821 |  "untrusted",
1822 | ]
1823 | 
1824 | [[package]]
1825 | name = "seq-macro"
1826 | version = "0.3.5"
1827 | source = "registry+https://github.com/rust-lang/crates.io-index"
1828 | checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4"
1829 | 
1830 | [[package]]
1831 | name = "serde"
1832 | version = "1.0.193"
1833 | source = "registry+https://github.com/rust-lang/crates.io-index"
1834 | checksum = "25dd9975e68d0cb5aa1120c288333fc98731bd1dd12f561e468ea4728c042b89"
1835 | dependencies = [
1836 |  "serde_derive",
1837 | ]
1838 | 
1839 | [[package]]
1840 | name = "serde_derive"
1841 | version = "1.0.193"
1842 | source = "registry+https://github.com/rust-lang/crates.io-index"
1843 | checksum = "43576ca501357b9b071ac53cdc7da8ef0cbd9493d8df094cd821777ea6e894d3"
1844 | dependencies = [
1845 |  "proc-macro2",
1846 |  "quote",
1847 |  "syn 2.0.39",
1848 | ]
1849 | 
1850 | [[package]]
1851 | name = "serde_json"
1852 | version = "1.0.108"
1853 | source = "registry+https://github.com/rust-lang/crates.io-index"
1854 | checksum = "3d1c7e3eac408d115102c4c24ad393e0821bb3a5df4d506a80f85f7a742a526b"
1855 | dependencies = [
1856 |  "itoa",
1857 |  "ryu",
1858 |  "serde",
1859 | ]
1860 | 
1861 | [[package]]
1862 | name = "serde_plain"
1863 | version = "1.0.2"
1864 | source = "registry+https://github.com/rust-lang/crates.io-index"
1865 | checksum = "9ce1fc6db65a611022b23a0dec6975d63fb80a302cb3388835ff02c097258d50"
1866 | dependencies = [
1867 |  "serde",
1868 | ]
1869 | 
1870 | [[package]]
1871 | name = "sha2"
1872 | version = "0.10.8"
1873 | source = "registry+https://github.com/rust-lang/crates.io-index"
1874 | checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8"
1875 | dependencies = [
1876 |  "cfg-if",
1877 |  "cpufeatures",
1878 |  "digest",
1879 | ]
1880 | 
1881 | [[package]]
1882 | name = "sharded-slab"
1883 | version = "0.1.7"
1884 | source = "registry+https://github.com/rust-lang/crates.io-index"
1885 | checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6"
1886 | dependencies = [
1887 |  "lazy_static",
1888 | ]
1889 | 
1890 | [[package]]
1891 | name = "smallvec"
1892 | version = "1.11.2"
1893 | source = "registry+https://github.com/rust-lang/crates.io-index"
1894 | checksum = "4dccd0940a2dcdf68d092b8cbab7dc0ad8fa938bf95787e1b916b0e3d0e8e970"
1895 | 
1896 | [[package]]
1897 | name = "spin"
1898 | version = "0.9.8"
1899 | source = "registry+https://github.com/rust-lang/crates.io-index"
1900 | checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67"
1901 | 
1902 | [[package]]
1903 | name = "spm_precompiled"
1904 | version = "0.1.4"
1905 | source = "registry+https://github.com/rust-lang/crates.io-index"
1906 | checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326"
1907 | dependencies = [
1908 |  "base64 0.13.1",
1909 |  "nom",
1910 |  "serde",
1911 |  "unicode-segmentation",
1912 | ]
1913 | 
1914 | [[package]]
1915 | name = "stable_deref_trait"
1916 | version = "1.2.0"
1917 | source = "registry+https://github.com/rust-lang/crates.io-index"
1918 | checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3"
1919 | 
1920 | [[package]]
1921 | name = "strsim"
1922 | version = "0.10.0"
1923 | source = "registry+https://github.com/rust-lang/crates.io-index"
1924 | checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623"
1925 | 
1926 | [[package]]
1927 | name = "syn"
1928 | version = "1.0.109"
1929 | source = "registry+https://github.com/rust-lang/crates.io-index"
1930 | checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237"
1931 | dependencies = [
1932 |  "proc-macro2",
1933 |  "quote",
1934 |  "unicode-ident",
1935 | ]
1936 | 
1937 | [[package]]
1938 | name = "syn"
1939 | version = "2.0.39"
1940 | source = "registry+https://github.com/rust-lang/crates.io-index"
1941 | checksum = "23e78b90f2fcf45d3e842032ce32e3f2d1545ba6636271dcbf24fa306d87be7a"
1942 | dependencies = [
1943 |  "proc-macro2",
1944 |  "quote",
1945 |  "unicode-ident",
1946 | ]
1947 | 
1948 | [[package]]
1949 | name = "synstructure"
1950 | version = "0.13.0"
1951 | source = "registry+https://github.com/rust-lang/crates.io-index"
1952 | checksum = "285ba80e733fac80aa4270fbcdf83772a79b80aa35c97075320abfee4a915b06"
1953 | dependencies = [
1954 |  "proc-macro2",
1955 |  "quote",
1956 |  "syn 2.0.39",
1957 |  "unicode-xid",
1958 | ]
1959 | 
1960 | [[package]]
1961 | name = "sysctl"
1962 | version = "0.5.5"
1963 | source = "registry+https://github.com/rust-lang/crates.io-index"
1964 | checksum = "ec7dddc5f0fee506baf8b9fdb989e242f17e4b11c61dfbb0635b705217199eea"
1965 | dependencies = [
1966 |  "bitflags 2.4.1",
1967 |  "byteorder",
1968 |  "enum-as-inner",
1969 |  "libc",
1970 |  "thiserror",
1971 |  "walkdir",
1972 | ]
1973 | 
1974 | [[package]]
1975 | name = "tar"
1976 | version = "0.4.40"
1977 | source = "registry+https://github.com/rust-lang/crates.io-index"
1978 | checksum = "b16afcea1f22891c49a00c751c7b63b2233284064f11a200fc624137c51e2ddb"
1979 | dependencies = [
1980 |  "filetime",
1981 |  "libc",
1982 |  "xattr",
1983 | ]
1984 | 
1985 | [[package]]
1986 | name = "textwrap"
1987 | version = "0.16.0"
1988 | source = "registry+https://github.com/rust-lang/crates.io-index"
1989 | checksum = "222a222a5bfe1bba4a77b45ec488a741b3cb8872e5e499451fd7d0129c9c7c3d"
1990 | 
1991 | [[package]]
1992 | name = "thiserror"
1993 | version = "1.0.50"
1994 | source = "registry+https://github.com/rust-lang/crates.io-index"
1995 | checksum = "f9a7210f5c9a7156bb50aa36aed4c95afb51df0df00713949448cf9e97d382d2"
1996 | dependencies = [
1997 |  "thiserror-impl",
1998 | ]
1999 | 
2000 | [[package]]
2001 | name = "thiserror-impl"
2002 | version = "1.0.50"
2003 | source = "registry+https://github.com/rust-lang/crates.io-index"
2004 | checksum = "266b2e40bc00e5a6c09c3584011e08b06f123c00362c92b975ba9843aaaa14b8"
2005 | dependencies = [
2006 |  "proc-macro2",
2007 |  "quote",
2008 |  "syn 2.0.39",
2009 | ]
2010 | 
2011 | [[package]]
2012 | name = "thread_local"
2013 | version = "1.1.7"
2014 | source = "registry+https://github.com/rust-lang/crates.io-index"
2015 | checksum = "3fdd6f064ccff2d6567adcb3873ca630700f00b5ad3f060c25b5dcfd9a4ce152"
2016 | dependencies = [
2017 |  "cfg-if",
2018 |  "once_cell",
2019 | ]
2020 | 
2021 | [[package]]
2022 | name = "tinytemplate"
2023 | version = "1.2.1"
2024 | source = "registry+https://github.com/rust-lang/crates.io-index"
2025 | checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc"
2026 | dependencies = [
2027 |  "serde",
2028 |  "serde_json",
2029 | ]
2030 | 
2031 | [[package]]
2032 | name = "tinyvec"
2033 | version = "1.6.0"
2034 | source = "registry+https://github.com/rust-lang/crates.io-index"
2035 | checksum = "87cc5ceb3875bb20c2890005a4e226a4651264a5c75edb2421b52861a0a0cb50"
2036 | dependencies = [
2037 |  "tinyvec_macros",
2038 | ]
2039 | 
2040 | [[package]]
2041 | name = "tinyvec_macros"
2042 | version = "0.1.1"
2043 | source = "registry+https://github.com/rust-lang/crates.io-index"
2044 | checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
2045 | 
2046 | [[package]]
2047 | name = "tokenizers"
2048 | version = "0.15.0"
2049 | source = "registry+https://github.com/rust-lang/crates.io-index"
2050 | checksum = "062b8a9613d6017633b80fb55fbb33f1aff006c36225a3025630753398034b3c"
2051 | dependencies = [
2052 |  "aho-corasick",
2053 |  "clap 4.4.11",
2054 |  "derive_builder 0.12.0",
2055 |  "esaxx-rs",
2056 |  "getrandom",
2057 |  "indicatif",
2058 |  "itertools 0.11.0",
2059 |  "lazy_static",
2060 |  "log",
2061 |  "macro_rules_attribute",
2062 |  "monostate",
2063 |  "onig",
2064 |  "paste",
2065 |  "rand",
2066 |  "rayon",
2067 |  "rayon-cond",
2068 |  "regex",
2069 |  "regex-syntax 0.7.5",
2070 |  "serde",
2071 |  "serde_json",
2072 |  "spm_precompiled",
2073 |  "thiserror",
2074 |  "unicode-normalization-alignments",
2075 |  "unicode-segmentation",
2076 |  "unicode_categories",
2077 | ]
2078 | 
2079 | [[package]]
2080 | name = "toml"
2081 | version = "0.5.11"
2082 | source = "registry+https://github.com/rust-lang/crates.io-index"
2083 | checksum = "f4f7f0dd8d50a853a531c426359045b1998f04219d88799810762cd4ad314234"
2084 | dependencies = [
2085 |  "serde",
2086 | ]
2087 | 
2088 | [[package]]
2089 | name = "tracing"
2090 | version = "0.1.40"
2091 | source = "registry+https://github.com/rust-lang/crates.io-index"
2092 | checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef"
2093 | dependencies = [
2094 |  "pin-project-lite",
2095 |  "tracing-attributes",
2096 |  "tracing-core",
2097 | ]
2098 | 
2099 | [[package]]
2100 | name = "tracing-attributes"
2101 | version = "0.1.27"
2102 | source = "registry+https://github.com/rust-lang/crates.io-index"
2103 | checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7"
2104 | dependencies = [
2105 |  "proc-macro2",
2106 |  "quote",
2107 |  "syn 2.0.39",
2108 | ]
2109 | 
2110 | [[package]]
2111 | name = "tracing-chrome"
2112 | version = "0.7.1"
2113 | source = "registry+https://github.com/rust-lang/crates.io-index"
2114 | checksum = "496b3cd5447f7ff527bbbf19b071ad542a000adf297d4127078b4dfdb931f41a"
2115 | dependencies = [
2116 |  "serde_json",
2117 |  "tracing-core",
2118 |  "tracing-subscriber",
2119 | ]
2120 | 
2121 | [[package]]
2122 | name = "tracing-core"
2123 | version = "0.1.32"
2124 | source = "registry+https://github.com/rust-lang/crates.io-index"
2125 | checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54"
2126 | dependencies = [
2127 |  "once_cell",
2128 |  "valuable",
2129 | ]
2130 | 
2131 | [[package]]
2132 | name = "tracing-log"
2133 | version = "0.2.0"
2134 | source = "registry+https://github.com/rust-lang/crates.io-index"
2135 | checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3"
2136 | dependencies = [
2137 |  "log",
2138 |  "once_cell",
2139 |  "tracing-core",
2140 | ]
2141 | 
2142 | [[package]]
2143 | name = "tracing-subscriber"
2144 | version = "0.3.18"
2145 | source = "registry+https://github.com/rust-lang/crates.io-index"
2146 | checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b"
2147 | dependencies = [
2148 |  "nu-ansi-term",
2149 |  "sharded-slab",
2150 |  "smallvec",
2151 |  "thread_local",
2152 |  "tracing-core",
2153 |  "tracing-log",
2154 | ]
2155 | 
2156 | [[package]]
2157 | name = "typenum"
2158 | version = "1.17.0"
2159 | source = "registry+https://github.com/rust-lang/crates.io-index"
2160 | checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825"
2161 | 
2162 | [[package]]
2163 | name = "unicode-bidi"
2164 | version = "0.3.14"
2165 | source = "registry+https://github.com/rust-lang/crates.io-index"
2166 | checksum = "6f2528f27a9eb2b21e69c95319b30bd0efd85d09c379741b0f78ea1d86be2416"
2167 | 
2168 | [[package]]
2169 | name = "unicode-ident"
2170 | version = "1.0.12"
2171 | source = "registry+https://github.com/rust-lang/crates.io-index"
2172 | checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b"
2173 | 
2174 | [[package]]
2175 | name = "unicode-normalization"
2176 | version = "0.1.22"
2177 | source = "registry+https://github.com/rust-lang/crates.io-index"
2178 | checksum = "5c5713f0fc4b5db668a2ac63cdb7bb4469d8c9fed047b1d0292cc7b0ce2ba921"
2179 | dependencies = [
2180 |  "tinyvec",
2181 | ]
2182 | 
2183 | [[package]]
2184 | name = "unicode-normalization-alignments"
2185 | version = "0.1.12"
2186 | source = "registry+https://github.com/rust-lang/crates.io-index"
2187 | checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de"
2188 | dependencies = [
2189 |  "smallvec",
2190 | ]
2191 | 
2192 | [[package]]
2193 | name = "unicode-segmentation"
2194 | version = "1.10.1"
2195 | source = "registry+https://github.com/rust-lang/crates.io-index"
2196 | checksum = "1dd624098567895118886609431a7c3b8f516e41d30e0643f03d94592a147e36"
2197 | 
2198 | [[package]]
2199 | name = "unicode-width"
2200 | version = "0.1.11"
2201 | source = "registry+https://github.com/rust-lang/crates.io-index"
2202 | checksum = "e51733f11c9c4f72aa0c160008246859e340b00807569a0da0e7a1079b27ba85"
2203 | 
2204 | [[package]]
2205 | name = "unicode-xid"
2206 | version = "0.2.4"
2207 | source = "registry+https://github.com/rust-lang/crates.io-index"
2208 | checksum = "f962df74c8c05a667b5ee8bcf162993134c104e96440b663c8daa176dc772d8c"
2209 | 
2210 | [[package]]
2211 | name = "unicode_categories"
2212 | version = "0.1.1"
2213 | source = "registry+https://github.com/rust-lang/crates.io-index"
2214 | checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
2215 | 
2216 | [[package]]
2217 | name = "untrusted"
2218 | version = "0.9.0"
2219 | source = "registry+https://github.com/rust-lang/crates.io-index"
2220 | checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
2221 | 
2222 | [[package]]
2223 | name = "ureq"
2224 | version = "2.9.1"
2225 | source = "registry+https://github.com/rust-lang/crates.io-index"
2226 | checksum = "f8cdd25c339e200129fe4de81451814e5228c9b771d57378817d6117cc2b3f97"
2227 | dependencies = [
2228 |  "base64 0.21.5",
2229 |  "flate2",
2230 |  "log",
2231 |  "once_cell",
2232 |  "rustls",
2233 |  "rustls-webpki",
2234 |  "serde",
2235 |  "serde_json",
2236 |  "url",
2237 |  "webpki-roots",
2238 | ]
2239 | 
2240 | [[package]]
2241 | name = "url"
2242 | version = "2.5.0"
2243 | source = "registry+https://github.com/rust-lang/crates.io-index"
2244 | checksum = "31e6302e3bb753d46e83516cae55ae196fc0c309407cf11ab35cc51a4c2a4633"
2245 | dependencies = [
2246 |  "form_urlencoded",
2247 |  "idna",
2248 |  "percent-encoding",
2249 | ]
2250 | 
2251 | [[package]]
2252 | name = "utf8parse"
2253 | version = "0.2.1"
2254 | source = "registry+https://github.com/rust-lang/crates.io-index"
2255 | checksum = "711b9620af191e0cdc7468a8d14e709c3dcdb115b36f838e601583af800a370a"
2256 | 
2257 | [[package]]
2258 | name = "uuid"
2259 | version = "1.6.1"
2260 | source = "registry+https://github.com/rust-lang/crates.io-index"
2261 | checksum = "5e395fcf16a7a3d8127ec99782007af141946b4795001f876d54fb0d55978560"
2262 | dependencies = [
2263 |  "getrandom",
2264 | ]
2265 | 
2266 | [[package]]
2267 | name = "valuable"
2268 | version = "0.1.0"
2269 | source = "registry+https://github.com/rust-lang/crates.io-index"
2270 | checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d"
2271 | 
2272 | [[package]]
2273 | name = "version_check"
2274 | version = "0.9.4"
2275 | source = "registry+https://github.com/rust-lang/crates.io-index"
2276 | checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
2277 | 
2278 | [[package]]
2279 | name = "walkdir"
2280 | version = "2.4.0"
2281 | source = "registry+https://github.com/rust-lang/crates.io-index"
2282 | checksum = "d71d857dc86794ca4c280d616f7da00d2dbfd8cd788846559a6813e6aa4b54ee"
2283 | dependencies = [
2284 |  "same-file",
2285 |  "winapi-util",
2286 | ]
2287 | 
2288 | [[package]]
2289 | name = "wasi"
2290 | version = "0.11.0+wasi-snapshot-preview1"
2291 | source = "registry+https://github.com/rust-lang/crates.io-index"
2292 | checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
2293 | 
2294 | [[package]]
2295 | name = "wasm-bindgen"
2296 | version = "0.2.89"
2297 | source = "registry+https://github.com/rust-lang/crates.io-index"
2298 | checksum = "0ed0d4f68a3015cc185aff4db9506a015f4b96f95303897bfa23f846db54064e"
2299 | dependencies = [
2300 |  "cfg-if",
2301 |  "wasm-bindgen-macro",
2302 | ]
2303 | 
2304 | [[package]]
2305 | name = "wasm-bindgen-backend"
2306 | version = "0.2.89"
2307 | source = "registry+https://github.com/rust-lang/crates.io-index"
2308 | checksum = "1b56f625e64f3a1084ded111c4d5f477df9f8c92df113852fa5a374dbda78826"
2309 | dependencies = [
2310 |  "bumpalo",
2311 |  "log",
2312 |  "once_cell",
2313 |  "proc-macro2",
2314 |  "quote",
2315 |  "syn 2.0.39",
2316 |  "wasm-bindgen-shared",
2317 | ]
2318 | 
2319 | [[package]]
2320 | name = "wasm-bindgen-macro"
2321 | version = "0.2.89"
2322 | source = "registry+https://github.com/rust-lang/crates.io-index"
2323 | checksum = "0162dbf37223cd2afce98f3d0785506dcb8d266223983e4b5b525859e6e182b2"
2324 | dependencies = [
2325 |  "quote",
2326 |  "wasm-bindgen-macro-support",
2327 | ]
2328 | 
2329 | [[package]]
2330 | name = "wasm-bindgen-macro-support"
2331 | version = "0.2.89"
2332 | source = "registry+https://github.com/rust-lang/crates.io-index"
2333 | checksum = "f0eb82fcb7930ae6219a7ecfd55b217f5f0893484b7a13022ebb2b2bf20b5283"
2334 | dependencies = [
2335 |  "proc-macro2",
2336 |  "quote",
2337 |  "syn 2.0.39",
2338 |  "wasm-bindgen-backend",
2339 |  "wasm-bindgen-shared",
2340 | ]
2341 | 
2342 | [[package]]
2343 | name = "wasm-bindgen-shared"
2344 | version = "0.2.89"
2345 | source = "registry+https://github.com/rust-lang/crates.io-index"
2346 | checksum = "7ab9b36309365056cd639da3134bf87fa8f3d86008abf99e612384a6eecd459f"
2347 | 
2348 | [[package]]
2349 | name = "web-sys"
2350 | version = "0.3.66"
2351 | source = "registry+https://github.com/rust-lang/crates.io-index"
2352 | checksum = "50c24a44ec86bb68fbecd1b3efed7e85ea5621b39b35ef2766b66cd984f8010f"
2353 | dependencies = [
2354 |  "js-sys",
2355 |  "wasm-bindgen",
2356 | ]
2357 | 
2358 | [[package]]
2359 | name = "webpki-roots"
2360 | version = "0.25.3"
2361 | source = "registry+https://github.com/rust-lang/crates.io-index"
2362 | checksum = "1778a42e8b3b90bff8d0f5032bf22250792889a5cdc752aa0020c84abe3aaf10"
2363 | 
2364 | [[package]]
2365 | name = "winapi"
2366 | version = "0.3.9"
2367 | source = "registry+https://github.com/rust-lang/crates.io-index"
2368 | checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
2369 | dependencies = [
2370 |  "winapi-i686-pc-windows-gnu",
2371 |  "winapi-x86_64-pc-windows-gnu",
2372 | ]
2373 | 
2374 | [[package]]
2375 | name = "winapi-i686-pc-windows-gnu"
2376 | version = "0.4.0"
2377 | source = "registry+https://github.com/rust-lang/crates.io-index"
2378 | checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
2379 | 
2380 | [[package]]
2381 | name = "winapi-util"
2382 | version = "0.1.6"
2383 | source = "registry+https://github.com/rust-lang/crates.io-index"
2384 | checksum = "f29e6f9198ba0d26b4c9f07dbe6f9ed633e1f3d5b8b414090084349e46a52596"
2385 | dependencies = [
2386 |  "winapi",
2387 | ]
2388 | 
2389 | [[package]]
2390 | name = "winapi-x86_64-pc-windows-gnu"
2391 | version = "0.4.0"
2392 | source = "registry+https://github.com/rust-lang/crates.io-index"
2393 | checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
2394 | 
2395 | [[package]]
2396 | name = "windows-core"
2397 | version = "0.51.1"
2398 | source = "registry+https://github.com/rust-lang/crates.io-index"
2399 | checksum = "f1f8cf84f35d2db49a46868f947758c7a1138116f7fac3bc844f43ade1292e64"
2400 | dependencies = [
2401 |  "windows-targets 0.48.5",
2402 | ]
2403 | 
2404 | [[package]]
2405 | name = "windows-sys"
2406 | version = "0.45.0"
2407 | source = "registry+https://github.com/rust-lang/crates.io-index"
2408 | checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0"
2409 | dependencies = [
2410 |  "windows-targets 0.42.2",
2411 | ]
2412 | 
2413 | [[package]]
2414 | name = "windows-sys"
2415 | version = "0.48.0"
2416 | source = "registry+https://github.com/rust-lang/crates.io-index"
2417 | checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9"
2418 | dependencies = [
2419 |  "windows-targets 0.48.5",
2420 | ]
2421 | 
2422 | [[package]]
2423 | name = "windows-sys"
2424 | version = "0.52.0"
2425 | source = "registry+https://github.com/rust-lang/crates.io-index"
2426 | checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d"
2427 | dependencies = [
2428 |  "windows-targets 0.52.0",
2429 | ]
2430 | 
2431 | [[package]]
2432 | name = "windows-targets"
2433 | version = "0.42.2"
2434 | source = "registry+https://github.com/rust-lang/crates.io-index"
2435 | checksum = "8e5180c00cd44c9b1c88adb3693291f1cd93605ded80c250a75d472756b4d071"
2436 | dependencies = [
2437 |  "windows_aarch64_gnullvm 0.42.2",
2438 |  "windows_aarch64_msvc 0.42.2",
2439 |  "windows_i686_gnu 0.42.2",
2440 |  "windows_i686_msvc 0.42.2",
2441 |  "windows_x86_64_gnu 0.42.2",
2442 |  "windows_x86_64_gnullvm 0.42.2",
2443 |  "windows_x86_64_msvc 0.42.2",
2444 | ]
2445 | 
2446 | [[package]]
2447 | name = "windows-targets"
2448 | version = "0.48.5"
2449 | source = "registry+https://github.com/rust-lang/crates.io-index"
2450 | checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c"
2451 | dependencies = [
2452 |  "windows_aarch64_gnullvm 0.48.5",
2453 |  "windows_aarch64_msvc 0.48.5",
2454 |  "windows_i686_gnu 0.48.5",
2455 |  "windows_i686_msvc 0.48.5",
2456 |  "windows_x86_64_gnu 0.48.5",
2457 |  "windows_x86_64_gnullvm 0.48.5",
2458 |  "windows_x86_64_msvc 0.48.5",
2459 | ]
2460 | 
2461 | [[package]]
2462 | name = "windows-targets"
2463 | version = "0.52.0"
2464 | source = "registry+https://github.com/rust-lang/crates.io-index"
2465 | checksum = "8a18201040b24831fbb9e4eb208f8892e1f50a37feb53cc7ff887feb8f50e7cd"
2466 | dependencies = [
2467 |  "windows_aarch64_gnullvm 0.52.0",
2468 |  "windows_aarch64_msvc 0.52.0",
2469 |  "windows_i686_gnu 0.52.0",
2470 |  "windows_i686_msvc 0.52.0",
2471 |  "windows_x86_64_gnu 0.52.0",
2472 |  "windows_x86_64_gnullvm 0.52.0",
2473 |  "windows_x86_64_msvc 0.52.0",
2474 | ]
2475 | 
2476 | [[package]]
2477 | name = "windows_aarch64_gnullvm"
2478 | version = "0.42.2"
2479 | source = "registry+https://github.com/rust-lang/crates.io-index"
2480 | checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8"
2481 | 
2482 | [[package]]
2483 | name = "windows_aarch64_gnullvm"
2484 | version = "0.48.5"
2485 | source = "registry+https://github.com/rust-lang/crates.io-index"
2486 | checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8"
2487 | 
2488 | [[package]]
2489 | name = "windows_aarch64_gnullvm"
2490 | version = "0.52.0"
2491 | source = "registry+https://github.com/rust-lang/crates.io-index"
2492 | checksum = "cb7764e35d4db8a7921e09562a0304bf2f93e0a51bfccee0bd0bb0b666b015ea"
2493 | 
2494 | [[package]]
2495 | name = "windows_aarch64_msvc"
2496 | version = "0.42.2"
2497 | source = "registry+https://github.com/rust-lang/crates.io-index"
2498 | checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43"
2499 | 
2500 | [[package]]
2501 | name = "windows_aarch64_msvc"
2502 | version = "0.48.5"
2503 | source = "registry+https://github.com/rust-lang/crates.io-index"
2504 | checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc"
2505 | 
2506 | [[package]]
2507 | name = "windows_aarch64_msvc"
2508 | version = "0.52.0"
2509 | source = "registry+https://github.com/rust-lang/crates.io-index"
2510 | checksum = "bbaa0368d4f1d2aaefc55b6fcfee13f41544ddf36801e793edbbfd7d7df075ef"
2511 | 
2512 | [[package]]
2513 | name = "windows_i686_gnu"
2514 | version = "0.42.2"
2515 | source = "registry+https://github.com/rust-lang/crates.io-index"
2516 | checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f"
2517 | 
2518 | [[package]]
2519 | name = "windows_i686_gnu"
2520 | version = "0.48.5"
2521 | source = "registry+https://github.com/rust-lang/crates.io-index"
2522 | checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e"
2523 | 
2524 | [[package]]
2525 | name = "windows_i686_gnu"
2526 | version = "0.52.0"
2527 | source = "registry+https://github.com/rust-lang/crates.io-index"
2528 | checksum = "a28637cb1fa3560a16915793afb20081aba2c92ee8af57b4d5f28e4b3e7df313"
2529 | 
2530 | [[package]]
2531 | name = "windows_i686_msvc"
2532 | version = "0.42.2"
2533 | source = "registry+https://github.com/rust-lang/crates.io-index"
2534 | checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060"
2535 | 
2536 | [[package]]
2537 | name = "windows_i686_msvc"
2538 | version = "0.48.5"
2539 | source = "registry+https://github.com/rust-lang/crates.io-index"
2540 | checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406"
2541 | 
2542 | [[package]]
2543 | name = "windows_i686_msvc"
2544 | version = "0.52.0"
2545 | source = "registry+https://github.com/rust-lang/crates.io-index"
2546 | checksum = "ffe5e8e31046ce6230cc7215707b816e339ff4d4d67c65dffa206fd0f7aa7b9a"
2547 | 
2548 | [[package]]
2549 | name = "windows_x86_64_gnu"
2550 | version = "0.42.2"
2551 | source = "registry+https://github.com/rust-lang/crates.io-index"
2552 | checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36"
2553 | 
2554 | [[package]]
2555 | name = "windows_x86_64_gnu"
2556 | version = "0.48.5"
2557 | source = "registry+https://github.com/rust-lang/crates.io-index"
2558 | checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e"
2559 | 
2560 | [[package]]
2561 | name = "windows_x86_64_gnu"
2562 | version = "0.52.0"
2563 | source = "registry+https://github.com/rust-lang/crates.io-index"
2564 | checksum = "3d6fa32db2bc4a2f5abeacf2b69f7992cd09dca97498da74a151a3132c26befd"
2565 | 
2566 | [[package]]
2567 | name = "windows_x86_64_gnullvm"
2568 | version = "0.42.2"
2569 | source = "registry+https://github.com/rust-lang/crates.io-index"
2570 | checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3"
2571 | 
2572 | [[package]]
2573 | name = "windows_x86_64_gnullvm"
2574 | version = "0.48.5"
2575 | source = "registry+https://github.com/rust-lang/crates.io-index"
2576 | checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc"
2577 | 
2578 | [[package]]
2579 | name = "windows_x86_64_gnullvm"
2580 | version = "0.52.0"
2581 | source = "registry+https://github.com/rust-lang/crates.io-index"
2582 | checksum = "1a657e1e9d3f514745a572a6846d3c7aa7dbe1658c056ed9c3344c4109a6949e"
2583 | 
2584 | [[package]]
2585 | name = "windows_x86_64_msvc"
2586 | version = "0.42.2"
2587 | source = "registry+https://github.com/rust-lang/crates.io-index"
2588 | checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0"
2589 | 
2590 | [[package]]
2591 | name = "windows_x86_64_msvc"
2592 | version = "0.48.5"
2593 | source = "registry+https://github.com/rust-lang/crates.io-index"
2594 | checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538"
2595 | 
2596 | [[package]]
2597 | name = "windows_x86_64_msvc"
2598 | version = "0.52.0"
2599 | source = "registry+https://github.com/rust-lang/crates.io-index"
2600 | checksum = "dff9641d1cd4be8d1a070daf9e3773c5f67e78b4d9d42263020c057706765c04"
2601 | 
2602 | [[package]]
2603 | name = "xattr"
2604 | version = "1.1.3"
2605 | source = "registry+https://github.com/rust-lang/crates.io-index"
2606 | checksum = "a7dae5072fe1f8db8f8d29059189ac175196e410e40ba42d5d4684ae2f750995"
2607 | dependencies = [
2608 |  "libc",
2609 |  "linux-raw-sys",
2610 |  "rustix",
2611 | ]
2612 | 
2613 | [[package]]
2614 | name = "yoke"
2615 | version = "0.7.3"
2616 | source = "registry+https://github.com/rust-lang/crates.io-index"
2617 | checksum = "65e71b2e4f287f467794c671e2b8f8a5f3716b3c829079a1c44740148eff07e4"
2618 | dependencies = [
2619 |  "serde",
2620 |  "stable_deref_trait",
2621 |  "yoke-derive",
2622 |  "zerofrom",
2623 | ]
2624 | 
2625 | [[package]]
2626 | name = "yoke-derive"
2627 | version = "0.7.3"
2628 | source = "registry+https://github.com/rust-lang/crates.io-index"
2629 | checksum = "9e6936f0cce458098a201c245a11bef556c6a0181129c7034d10d76d1ec3a2b8"
2630 | dependencies = [
2631 |  "proc-macro2",
2632 |  "quote",
2633 |  "syn 2.0.39",
2634 |  "synstructure",
2635 | ]
2636 | 
2637 | [[package]]
2638 | name = "zerofrom"
2639 | version = "0.1.3"
2640 | source = "registry+https://github.com/rust-lang/crates.io-index"
2641 | checksum = "655b0814c5c0b19ade497851070c640773304939a6c0fd5f5fb43da0696d05b7"
2642 | dependencies = [
2643 |  "zerofrom-derive",
2644 | ]
2645 | 
2646 | [[package]]
2647 | name = "zerofrom-derive"
2648 | version = "0.1.3"
2649 | source = "registry+https://github.com/rust-lang/crates.io-index"
2650 | checksum = "e6a647510471d372f2e6c2e6b7219e44d8c574d24fdc11c610a61455782f18c3"
2651 | dependencies = [
2652 |  "proc-macro2",
2653 |  "quote",
2654 |  "syn 2.0.39",
2655 |  "synstructure",
2656 | ]
2657 | 
2658 | [[package]]
2659 | name = "zip"
2660 | version = "0.6.6"
2661 | source = "registry+https://github.com/rust-lang/crates.io-index"
2662 | checksum = "760394e246e4c28189f19d488c058bf16f564016aefac5d32bb1f3b51d5e9261"
2663 | dependencies = [
2664 |  "byteorder",
2665 |  "crc32fast",
2666 |  "crossbeam-utils",
2667 | ]
2668 | 


--------------------------------------------------------------------------------
/Cargo.toml:
--------------------------------------------------------------------------------
1 | [workspace]
2 | resolver = "2"
3 | 
4 | members = ["mamba-ssm"]
5 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2023 Flawedmatrix
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | # mamba-ssm
  2 | 
  3 | Optimized inference-only implementation of [Mamba](#references) [1] written in Rust.
  4 | 
  5 | ## Description
  6 | The primary goal of this project is to provide an inference backend that can run
  7 | Mamba on an Apple Silicon Macbook without having dependencies on CUDA entangled
  8 | with the code. The initial development specifically targets CPU-only as a
  9 | first-class citizen, with linear algebra routines supported by Accelerate or
 10 | Intel MKL.
 11 | 
 12 | The main dependency of this project is [Candle](https://github.com/huggingface/candle),
 13 | so supported platforms are mainly decided by their implementation in that
 14 | framework.
 15 | 
 16 | ### Supported Platforms
 17 | 
 18 | - [x] CPU
 19 | - [x] Accelerate framework (via `--features accelerate`)
 20 | - [ ] Intel MKL (via `features mkl`)
 21 |   - It probably works, but I haven't tested it yet
 22 | - [x] Metal
 23 |   - Still relatively unoptimized
 24 | - [x] CUDA (via `features cuda`)
 25 |   - It works but no optimization was done for CUDA yet.
 26 | 
 27 | ### Supported Features
 28 | 
 29 | - [x] Inference via CLI
 30 | - [x] BF16 (CUDA only)
 31 | - [ ] FP16 support
 32 | - [ ] Quantized models
 33 | - [ ] Web interface for inference
 34 | 
 35 | ## Getting Started
 36 | 
 37 | 1. Prepare a Mamba safetensors model, config.json, and tokenizer.json and move these to the `/.models` directory.
 38 |    - Run `./download.sh` to download [mamba-2.8b-slimpj](https://huggingface.co/state-spaces/mamba-2.8b-slimpj/tree/refs%2Fpr%2F1) and the tokenizer from [gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b/blob/main/tokenizer.json)
 39 | 2. [Install Rust](https://www.rust-lang.org), then run:
 40 | 
 41 | ```bash
 42 | cargo build --release
 43 | target/release/mamba-cli --prompt "Mamba is the"
 44 | ```
 45 | 
 46 | You can also specify the model and config.json used by passing flags:
 47 | 
 48 | ```bash
 49 | target/release/mamba-cli -m models/mamba-2.8b-slimpj/model.safetensors -c models/mamba-2.8b-slimpj/config.json -prompt "Mamba is the"
 50 | ```
 51 | 
 52 | For other usage options such as passing the prompt by file, see the usage:
 53 | 
 54 | ```bash
 55 | target/release/mamba-cli --help
 56 | ```
 57 | 
 58 | ### Building with Apple Accelerate Framework support
 59 | 
 60 | ```bash
 61 | cargo build --release --features accelerate
 62 | ```
 63 | 
 64 | ### Building with Intel MKL framework support
 65 | 
 66 | ```bash
 67 | cargo build --release --features mkl
 68 | ```
 69 | 
 70 | ### Generation speed with CPU
 71 | 
 72 | Currently, with the Mamba 2.8b model, it generates at about 6.5 tokens/s with FP32 on CPU only on a M3 Max MBP.
 73 | 
 74 | ```bash
 75 | $ target/release/mamba-cli --temperature 0 -n 50 -f prompt.txt
 76 | avx: false, neon: true, simd128: false, f16c: false, num_threads: 16, cuda: false, metal: false, accelerate: true, mkl: false
 77 | temp: 0.00 repeat-penalty: 1.10 repeat-last-n: 64
 78 | loaded the model in 1.605674125s
 79 | generating 50 tokens with seed 16889006583945703583
 80 | 
 81 | Prompt processing time (98 tokens at 24.68 token/s)
 82 | I am that merry wanderer of the night.
 83 | I jest to Oberon and make him smile
 84 | When I a fat and bean-fed horse beguile,
 85 | Neighing in likeness of a filly foal:
 86 | And sometime lurk I in a gossip’s bowl,
 87 | In very likeness of a roasted crab,
 88 | And when she drinks, against her lips I bob
 89 | And on her wither’d dewlap pour the ale.
 90 | I am that merry jester of the night;
 91 | When he is sick and sad, I make him smile:
 92 | If his wife be angry with him, then I
 93 | Make him laugh, as if a fool were free.
 94 | But when she
 95 | 50 tokens generated (6.50 token/s)
 96 | ```
 97 | 
 98 | ## References
 99 | 
100 | [1] "Mamba: Linear-Time Sequence Modeling with Selective State Spaces"
101 |     Albert Gu and Tri Dao
102 |     https://arxiv.org/abs/2312.00752
103 | 
104 | [2] "The Annotated S4"
105 |     Sasha Rush and Sidd Karamcheti
106 |     https://srush.github.io/annotated-s4
107 | 
108 | [3] "Error Analysis and Improving the Accuracy of Winograd Convolution for Deep Neural Networks"
109 |     Barbara Barabasz, Andrew Anderson, Kirk M. Soodhalter, David Gregg
110 |     https://arxiv.org/abs/1803.10986
111 | 
112 | [4] "Winograd Convolution for Deep Neural Networks: Efficient Point Selection"
113 |     Syed Asad Alam, Andrew Anderson, Barbara Barabasz, David Gregg
114 |     https://arxiv.org/pdf/2201.10369.pdf
115 | 
116 | ### Code references
117 | - Original implementation: https://github.com/state-spaces/mamba
118 | - This repo was initially adapted from code from the
119 |   mamba-minimal candle-example:
120 |   https://github.com/huggingface/candle/tree/main/candle-examples/examples/mamba-minimal
121 | - Instructive minimal implementation: https://github.com/johnma2006/mamba-minimal
122 | 


--------------------------------------------------------------------------------
/download.sh:
--------------------------------------------------------------------------------
 1 | #!/bin/bash
 2 | 
 3 | # Directory where you want to clone the repo
 4 | DEST_DIR="./models"
 5 | 
 6 | cd $DEST_DIR
 7 | 
 8 | curl -L https://huggingface.co/state-spaces/mamba-2.8b-slimpj/raw/refs%2Fpr%2F1/config.json -o ./config.json
 9 | curl -L https://huggingface.co/EleutherAI/gpt-neox-20b/raw/main/tokenizer.json -o ./tokenizer.json
10 | curl -L https://huggingface.co/state-spaces/mamba-2.8b-slimpj/resolve/refs%2Fpr%2F1/model.safetensors -o ./model.safetensors
11 | 


--------------------------------------------------------------------------------
/mamba-ssm/Cargo.toml:
--------------------------------------------------------------------------------
 1 | [package]
 2 | name = "mamba-ssm"
 3 | version = "0.1.0"
 4 | edition = "2021"
 5 | authors = ["flawedmatrix"]
 6 | license = "MIT"
 7 | repository = "https://github.com/flawedmatrix/mamba-ssm"
 8 | readme = "README.md"
 9 | 
10 | [features]
11 | accelerate = [
12 |     "dep:accelerate-src",
13 |     "candle/accelerate",
14 |     "candle-nn/accelerate",
15 |     "candle-transformers/accelerate",
16 | ]
17 | mkl = [
18 |     "dep:intel-mkl-src",
19 |     "candle/mkl",
20 |     "candle-nn/mkl",
21 |     "candle-transformers/mkl",
22 | ]
23 | cuda = ["candle/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
24 | cudnn = ["candle/cudnn"]
25 | metal = ["candle/metal", "candle-nn/metal"]
26 | 
27 | 
28 | [dependencies]
29 | accelerate-src = { version = "0.3", optional = true }
30 | intel-mkl-src = { version = "0.8", optional = true }
31 | candle = { package = "candle-core", version = "0.4.1" }
32 | candle-nn = "0.4.1"
33 | candle-transformers = "0.4.1"
34 | candle-metal-kernels = "0.4.1"
35 | tokenizers = "0.15"
36 | anyhow = "1"
37 | clap = "4"
38 | serde = "1"
39 | serde_json = "1"
40 | tracing-chrome = "0.7"
41 | tracing-subscriber = "0.3"
42 | tracing = "0.1"
43 | rand = "0.8"
44 | rayon = "1.8"
45 | 
46 | [dev-dependencies]
47 | criterion = "0.4"
48 | 
49 | [[bench]]
50 | name = "conv1d_bench"
51 | path = "benches/conv1d.rs"
52 | harness = false
53 | 
54 | [[bench]]
55 | name = "elemwise_bench"
56 | path = "benches/elemwise.rs"
57 | harness = false
58 | 
59 | [[bench]]
60 | name = "linear_bench"
61 | path = "benches/linear.rs"
62 | harness = false
63 | 


--------------------------------------------------------------------------------
/mamba-ssm/README.md:
--------------------------------------------------------------------------------
1 | # mamba-ssm
2 | 
3 | Core implementation of mamba that can be consumed as a library.
4 | 
5 | See [src/bin/mamba-cli.rs](src/bin/mamba-cli.rs) for usage example.


--------------------------------------------------------------------------------
/mamba-ssm/benches/conv1d.rs:
--------------------------------------------------------------------------------
  1 | use criterion::{criterion_group, criterion_main, Criterion};
  2 | 
  3 | use candle::{Device, Module, Tensor};
  4 | use mamba_ssm::context::Context;
  5 | 
  6 | fn candle_conv1d(c: &mut Criterion) {
  7 |     let device = Device::Cpu;
  8 | 
  9 |     let x = Tensor::randn(0f32, 1., (1, 2000, 5120), &device).unwrap();
 10 | 
 11 |     let conv_cfg = candle_nn::Conv1dConfig {
 12 |         groups: 5120,
 13 |         padding: 3,
 14 |         ..Default::default()
 15 |     };
 16 | 
 17 |     let var_map = candle_nn::VarMap::new();
 18 |     let vb = candle_nn::VarBuilder::from_varmap(&var_map, candle::DType::F32, &device);
 19 | 
 20 |     let candle_conv1d = candle_nn::conv1d(5120, 5120, 4, conv_cfg, vb).unwrap();
 21 | 
 22 |     c.bench_function("candle_nn::conv1d single batch", |b| {
 23 |         b.iter(|| candle_conv1d.forward(&x.t().unwrap()).unwrap().t().unwrap())
 24 |     });
 25 | }
 26 | 
 27 | fn custom_conv1d(c: &mut Criterion) {
 28 |     let device = Device::Cpu;
 29 | 
 30 |     let x = Tensor::randn(0f32, 1., (1, 2000, 5120), &device).unwrap();
 31 | 
 32 |     let conv_cfg = candle_nn::Conv1dConfig {
 33 |         groups: 5120,
 34 |         padding: 3,
 35 |         ..Default::default()
 36 |     };
 37 | 
 38 |     let var_map = candle_nn::VarMap::new();
 39 |     let vb = candle_nn::VarBuilder::from_varmap(&var_map, candle::DType::F32, &device);
 40 | 
 41 |     let ctx = Context::new(candle::DType::F32, &device);
 42 |     ctx.freeze();
 43 | 
 44 |     let custom_conv1d = mamba_ssm::nn::conv1d(5120, 5120, 4, conv_cfg, vb, ctx).unwrap();
 45 |     c.bench_function("mamba_ssm::nn::conv1d single batch", |b| {
 46 |         b.iter(|| custom_conv1d.forward(&x).unwrap())
 47 |     });
 48 | }
 49 | 
 50 | fn candle_conv1d_inference(c: &mut Criterion) {
 51 |     let device = Device::Cpu;
 52 | 
 53 |     let x = Tensor::randn(0f32, 1., (1, 1, 5120), &device).unwrap();
 54 | 
 55 |     let conv_cfg = candle_nn::Conv1dConfig {
 56 |         groups: 5120,
 57 |         padding: 3,
 58 |         ..Default::default()
 59 |     };
 60 | 
 61 |     let var_map = candle_nn::VarMap::new();
 62 |     let vb = candle_nn::VarBuilder::from_varmap(&var_map, candle::DType::F32, &device);
 63 | 
 64 |     let candle_conv1d = candle_nn::conv1d(5120, 5120, 4, conv_cfg, vb).unwrap();
 65 | 
 66 |     c.bench_function("candle_nn::conv1d inference", |b| {
 67 |         b.iter(|| candle_conv1d.forward(&x.t().unwrap()).unwrap().t().unwrap())
 68 |     });
 69 | }
 70 | 
 71 | fn custom_conv1d_inference(c: &mut Criterion) {
 72 |     let device = Device::Cpu;
 73 | 
 74 |     let x = Tensor::randn(0f32, 1., (1, 1, 5120), &device).unwrap();
 75 |     let cached = Tensor::randn(0f32, 1., (3, 5120), &device).unwrap();
 76 | 
 77 |     let conv_cfg = candle_nn::Conv1dConfig {
 78 |         groups: 5120,
 79 |         padding: 3,
 80 |         ..Default::default()
 81 |     };
 82 | 
 83 |     let var_map = candle_nn::VarMap::new();
 84 |     let vb = candle_nn::VarBuilder::from_varmap(&var_map, candle::DType::F32, &device);
 85 | 
 86 |     let ctx = Context::new(candle::DType::F32, &device);
 87 |     let mut ctx = ctx.pp(0);
 88 |     _ = ctx.get((3, 5120), "ph").unwrap();
 89 |     ctx.set("ph", cached).unwrap();
 90 |     ctx.freeze();
 91 | 
 92 |     let custom_conv1d = mamba_ssm::nn::conv1d(5120, 5120, 4, conv_cfg, vb, ctx).unwrap();
 93 |     c.bench_function("mamba_ssm::nn::conv1d inference", |b| {
 94 |         b.iter(|| custom_conv1d.forward(&x).unwrap())
 95 |     });
 96 | }
 97 | 
 98 | criterion_group!(
 99 |     conv1d_benches,
100 |     candle_conv1d,
101 |     custom_conv1d,
102 |     candle_conv1d_inference,
103 |     custom_conv1d_inference
104 | );
105 | criterion_main!(conv1d_benches);
106 | 


--------------------------------------------------------------------------------
/mamba-ssm/benches/elemwise.rs:
--------------------------------------------------------------------------------
 1 | use criterion::{criterion_group, criterion_main, Criterion};
 2 | 
 3 | use candle::{Device, Tensor};
 4 | 
 5 | use mamba_ssm::ops::elemwise;
 6 | 
 7 | fn candle_elemwise_mul(c: &mut Criterion) {
 8 |     let device = Device::Cpu;
 9 | 
10 |     let x = Tensor::randn(0f32, 1., (5000, 5000), &device).unwrap();
11 |     let y = Tensor::randn(0f32, 1., (5000, 5000), &device).unwrap();
12 | 
13 |     c.bench_function("elemwise mul", |b| b.iter(|| x.mul(&y).unwrap()));
14 | }
15 | 
16 | fn candle_elemwise_add(c: &mut Criterion) {
17 |     let device = Device::Cpu;
18 | 
19 |     let x = Tensor::randn(0f32, 1., (5000, 5000), &device).unwrap();
20 |     let y = Tensor::randn(0f32, 1., (5000, 5000), &device).unwrap();
21 | 
22 |     c.bench_function("elemwise add", |b| b.iter(|| x.add(&y).unwrap()));
23 | }
24 | 
25 | fn custom_elemwise_add(c: &mut Criterion) {
26 |     let device = Device::Cpu;
27 | 
28 |     let x = Tensor::randn(0f32, 1., (5000, 5000), &device).unwrap();
29 |     let y = Tensor::randn(0f32, 1., (5000, 5000), &device).unwrap();
30 | 
31 |     c.bench_function("custom elemwise add", |b| {
32 |         b.iter(|| x.apply_op2(&y, elemwise::BinaryOp::Add).unwrap())
33 |     });
34 | }
35 | 
36 | fn custom_elemwise_mul(c: &mut Criterion) {
37 |     let device = Device::Cpu;
38 | 
39 |     let x = Tensor::randn(0f32, 1., (5000, 5000), &device).unwrap();
40 |     let y = Tensor::randn(0f32, 1., (5000, 5000), &device).unwrap();
41 | 
42 |     c.bench_function("custom elemwise mul", |b| {
43 |         b.iter(|| x.apply_op2(&y, elemwise::BinaryOp::Mul).unwrap())
44 |     });
45 | }
46 | 
47 | criterion_group!(
48 |     elemwise_benches,
49 |     candle_elemwise_mul,
50 |     candle_elemwise_add,
51 |     custom_elemwise_mul,
52 |     custom_elemwise_add,
53 | );
54 | criterion_main!(elemwise_benches);
55 | 


--------------------------------------------------------------------------------
/mamba-ssm/benches/linear.rs:
--------------------------------------------------------------------------------
 1 | use criterion::{criterion_group, criterion_main, Criterion};
 2 | 
 3 | use candle::{Device, Module, Tensor};
 4 | 
 5 | fn candle_gemm_mul(c: &mut Criterion) {
 6 |     let device = Device::Cpu;
 7 | 
 8 |     let x = Tensor::randn(0f32, 1., (1, 2560), &device).unwrap();
 9 |     let y = Tensor::randn(0f32, 1., (2560, 10240), &device).unwrap();
10 | 
11 |     c.bench_function("gemm mul", |b| b.iter(|| x.matmul(&y).unwrap()));
12 | }
13 | 
14 | fn candle_gemm_mul_t(c: &mut Criterion) {
15 |     let device = Device::Cpu;
16 | 
17 |     let x = Tensor::randn(0f32, 1., (1, 2560), &device).unwrap();
18 |     let y = Tensor::randn(0f32, 1., (10240, 2560), &device).unwrap();
19 | 
20 |     c.bench_function("gemm mul transposed", |b| {
21 |         b.iter(|| x.matmul(&y.t().unwrap()).unwrap())
22 |     });
23 | }
24 | 
25 | fn candle_linear(c: &mut Criterion) {
26 |     let device = Device::Cpu;
27 | 
28 |     let x = Tensor::randn(0f32, 1., (1, 1, 2560), &device).unwrap();
29 | 
30 |     let var_map = candle_nn::VarMap::new();
31 |     let vb = candle_nn::VarBuilder::from_varmap(&var_map, candle::DType::F32, &device);
32 | 
33 |     let layer = candle_nn::linear(2560, 10240, vb).unwrap();
34 | 
35 |     c.bench_function("candle_nn::linear single batch", |b| {
36 |         b.iter(|| layer.forward(&x).unwrap())
37 |     });
38 | }
39 | 
40 | fn candle_linear_no_bias(c: &mut Criterion) {
41 |     let device = Device::Cpu;
42 | 
43 |     let x = Tensor::randn(0f32, 1., (1, 1, 2560), &device).unwrap();
44 | 
45 |     let var_map = candle_nn::VarMap::new();
46 |     let vb = candle_nn::VarBuilder::from_varmap(&var_map, candle::DType::F32, &device);
47 | 
48 |     let layer = candle_nn::linear_no_bias(2560, 10240, vb).unwrap();
49 | 
50 |     c.bench_function("candle_nn::linear_no_bias single batch", |b| {
51 |         b.iter(|| layer.forward(&x).unwrap())
52 |     });
53 | }
54 | 
55 | fn custom_linear(c: &mut Criterion) {
56 |     let device = Device::Cpu;
57 | 
58 |     let x = Tensor::randn(0f32, 1., (1, 1, 2560), &device).unwrap();
59 | 
60 |     let var_map = candle_nn::VarMap::new();
61 |     let vb = candle_nn::VarBuilder::from_varmap(&var_map, candle::DType::F32, &device);
62 | 
63 |     let layer = mamba_ssm::nn::linear(2560, 10240, vb).unwrap();
64 | 
65 |     c.bench_function("mamba_ssm::nn::linear single batch", |b| {
66 |         b.iter(|| layer.forward(&x).unwrap())
67 |     });
68 | }
69 | 
70 | fn custom_linear_no_bias(c: &mut Criterion) {
71 |     let device = Device::Cpu;
72 | 
73 |     let x = Tensor::randn(0f32, 1., (1, 1, 2560), &device).unwrap();
74 | 
75 |     let var_map = candle_nn::VarMap::new();
76 |     let vb = candle_nn::VarBuilder::from_varmap(&var_map, candle::DType::F32, &device);
77 | 
78 |     let layer = mamba_ssm::nn::linear_no_bias(2560, 10240, vb).unwrap();
79 | 
80 |     c.bench_function("mamba_ssm::nn::linear_no_bias single batch", |b| {
81 |         b.iter(|| layer.forward(&x).unwrap())
82 |     });
83 | }
84 | 
85 | criterion_group!(
86 |     linear_benches,
87 |     candle_gemm_mul,
88 |     candle_gemm_mul_t,
89 |     candle_linear,
90 |     candle_linear_no_bias,
91 |     custom_linear,
92 |     custom_linear_no_bias,
93 | );
94 | criterion_main!(linear_benches);
95 | 


--------------------------------------------------------------------------------
/mamba-ssm/src/bin/mamba-cli.rs:
--------------------------------------------------------------------------------
  1 | use std::path::PathBuf;
  2 | 
  3 | use anyhow::{Error as E, Result};
  4 | use candle::{DType, Device};
  5 | use clap::{ArgGroup, Parser};
  6 | 
  7 | use candle_nn::VarBuilder;
  8 | use tokenizers::Tokenizer;
  9 | 
 10 | use rand::Rng;
 11 | 
 12 | use mamba_ssm::{
 13 |     context::Context,
 14 |     model::{Config, Model},
 15 |     TextGeneration,
 16 | };
 17 | 
 18 | #[derive(Parser, Debug)]
 19 | #[command(author, version, about, long_about = None)]
 20 | #[clap(group(ArgGroup::new("prompt_options")
 21 |                 .required(true)
 22 |                 .args(&["prompt", "prompt_file"])
 23 | ))]
 24 | struct Args {
 25 |     /// Enable tracing (generates a trace-timestamp.json file).
 26 |     #[arg(long)]
 27 |     tracing: bool,
 28 | 
 29 |     #[arg(long, short = 'p')]
 30 |     prompt: Option<String>,
 31 | 
 32 |     #[arg(long, short = 'f')]
 33 |     prompt_file: Option<PathBuf>,
 34 | 
 35 |     /// The temperature used to generate samples.
 36 |     #[arg(long, default_value_t = 0.7)]
 37 |     temperature: f64,
 38 | 
 39 |     /// Nucleus sampling probability cutoff.
 40 |     #[arg(long, default_value_t = 0.92)]
 41 |     top_p: f64,
 42 | 
 43 |     /// The seed to use when generating random samples.
 44 |     #[arg(long)]
 45 |     seed: Option<u64>,
 46 | 
 47 |     /// The length of the sample to generate (in tokens).
 48 |     #[arg(long, short = 'n', default_value_t = 500)]
 49 |     sample_len: usize,
 50 | 
 51 |     #[arg(long, default_value = "./models/tokenizer.json")]
 52 |     tokenizer_file: PathBuf,
 53 | 
 54 |     #[arg(long, short = 'm', default_value = "./models/model.safetensors")]
 55 |     weights_file: PathBuf,
 56 | 
 57 |     #[arg(long, short = 'c', default_value = "./models/config.json")]
 58 |     config_file: PathBuf,
 59 | 
 60 |     /// Penalty to be applied for repeating tokens, 1. means no penalty.
 61 |     #[arg(long, default_value_t = 1.1)]
 62 |     repeat_penalty: f32,
 63 | 
 64 |     /// The context size to consider for the repeat penalty.
 65 |     #[arg(long, default_value_t = 320)]
 66 |     repeat_last_n: usize,
 67 | 
 68 |     /// Use BF16 floating point format for calculations instead of F32. CUDA only.
 69 |     #[arg(long)]
 70 |     bf16: bool,
 71 | }
 72 | 
 73 | fn main() -> Result<()> {
 74 |     use tracing_chrome::ChromeLayerBuilder;
 75 |     use tracing_subscriber::prelude::*;
 76 | 
 77 |     let args = Args::parse();
 78 | 
 79 |     let _guard = if args.tracing {
 80 |         let (chrome_layer, guard) = ChromeLayerBuilder::new().include_args(true).build();
 81 |         tracing_subscriber::registry().with(chrome_layer).init();
 82 |         Some(guard)
 83 |     } else {
 84 |         None
 85 |     };
 86 | 
 87 |     println!(
 88 |         "avx: {}, neon: {}, simd128: {}, f16c: {}, num_threads: {}, cuda: {}, metal: {}, accelerate: {}, mkl: {}",
 89 |         candle::utils::with_avx(),
 90 |         candle::utils::with_neon(),
 91 |         candle::utils::with_simd128(),
 92 |         candle::utils::with_f16c(),
 93 |         candle::utils::get_num_threads(),
 94 |         candle::utils::cuda_is_available(),
 95 |         candle::utils::metal_is_available(),
 96 |         candle::utils::has_accelerate(),
 97 |         candle::utils::has_mkl()
 98 |     );
 99 |     println!(
100 |         "temp: {:.2} top-p: {:.2} repeat-penalty: {:.2} repeat-last-n: {}",
101 |         args.temperature, args.top_p, args.repeat_penalty, args.repeat_last_n
102 |     );
103 | 
104 |     let tokenizer = Tokenizer::from_file(args.tokenizer_file).map_err(E::msg)?;
105 | 
106 |     let start = std::time::Instant::now();
107 |     let config: Config = serde_json::from_slice(&std::fs::read(args.config_file)?)?;
108 | 
109 |     // TODO: Implement GPU-based inference
110 |     let device = if candle::utils::cuda_is_available() {
111 |         Device::cuda_if_available(0)?
112 |     } else if candle::utils::metal_is_available() {
113 |         Device::new_metal(0)?
114 |     } else {
115 |         Device::Cpu
116 |     };
117 | 
118 |     let dtype = if args.bf16 { DType::BF16 } else { DType::F32 };
119 | 
120 |     let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[args.weights_file], dtype, &device)? };
121 | 
122 |     let ctx = Context::new(dtype, &device);
123 | 
124 |     let model = Model::new(&config, vb.pp("backbone"), ctx)?;
125 |     println!("loaded the model in {:?}", start.elapsed());
126 | 
127 |     let seed = args.seed.unwrap_or(rand::thread_rng().gen());
128 |     println!("generating {} tokens with seed {seed}", args.sample_len);
129 | 
130 |     let prompt = if args.prompt.is_some() {
131 |         args.prompt.unwrap()
132 |     } else {
133 |         std::fs::read_to_string(args.prompt_file.unwrap())?
134 |     };
135 | 
136 |     let mut pipeline = TextGeneration::new(
137 |         model,
138 |         tokenizer,
139 |         seed,
140 |         Some(args.temperature),
141 |         Some(args.top_p),
142 |         args.repeat_penalty,
143 |         args.repeat_last_n,
144 |         &device,
145 |     );
146 |     pipeline.run(&prompt, args.sample_len)?;
147 |     Ok(())
148 | }
149 | 


--------------------------------------------------------------------------------
/mamba-ssm/src/context.rs:
--------------------------------------------------------------------------------
  1 | use std::sync::{
  2 |     atomic::{AtomicBool, Ordering::Relaxed},
  3 |     Arc,
  4 | };
  5 | 
  6 | use candle::{DType, Device, Result, Shape, Tensor};
  7 | use candle_nn::VarMap;
  8 | 
  9 | /// Wrapper type for candle_nn::VarMap that supports prefixes
 10 | #[derive(Clone)]
 11 | pub struct Context {
 12 |     // Internally uses Arc<Mutex>
 13 |     vm: VarMap,
 14 |     path: Vec<String>,
 15 | 
 16 |     dtype: DType,
 17 |     device: Device,
 18 | 
 19 |     frozen: Arc<AtomicBool>,
 20 | }
 21 | 
 22 | impl Context {
 23 |     pub fn new(dtype: DType, device: &Device) -> Self {
 24 |         Self {
 25 |             vm: VarMap::new(),
 26 |             path: Vec::<String>::new(),
 27 |             dtype,
 28 |             device: device.clone(),
 29 |             frozen: Arc::new(false.into()),
 30 |         }
 31 |     }
 32 |     /// Returns a new `Context` adding `s` to the current prefix, like `cd`ing
 33 |     /// into a directory.
 34 |     pub fn pp<S: ToString>(&self, s: S) -> Self {
 35 |         let mut path = self.path.clone();
 36 |         path.push(s.to_string());
 37 |         Self {
 38 |             vm: self.vm.clone(),
 39 |             path,
 40 |             dtype: self.dtype,
 41 |             device: self.device.clone(),
 42 |             frozen: self.frozen.clone(),
 43 |         }
 44 |     }
 45 |     pub fn set<K: AsRef<str>, V: AsRef<Tensor>>(&mut self, name: K, value: V) -> Result<()> {
 46 |         if self.frozen.load(Relaxed) {
 47 |             return Ok(());
 48 |         }
 49 |         let name = self.path(name.as_ref());
 50 |         let v_ref = value.as_ref();
 51 |         if v_ref.dtype() != self.dtype {
 52 |             candle::bail!(
 53 |                 "cannot save context {name} because the dtype doesn't match: {} != {}",
 54 |                 v_ref.dtype().as_str(),
 55 |                 self.dtype.as_str()
 56 |             )
 57 |         }
 58 |         if !v_ref.device().same_device(self.device()) {
 59 |             candle::bail!(
 60 |                 "cannot save context {name} because the device doesn't match: {:?} != {:?}",
 61 |                 v_ref.device(),
 62 |                 self.device()
 63 |             )
 64 |         }
 65 | 
 66 |         self.vm.set_one(name, value)
 67 |     }
 68 | 
 69 |     fn path(&self, tensor_name: &str) -> String {
 70 |         if self.path.is_empty() {
 71 |             tensor_name.to_string()
 72 |         } else {
 73 |             [&self.path.join("."), tensor_name].join(".")
 74 |         }
 75 |     }
 76 | 
 77 |     // Gets the variable, or initializes it with 0s if it isn't found
 78 |     pub fn get<S: Into<Shape>>(&self, shape: S, path: &str) -> Result<Tensor> {
 79 |         let name = self.path(path);
 80 |         let zero_init = candle_nn::Init::Const(0.);
 81 |         self.vm
 82 |             .get(shape, name.as_str(), zero_init, self.dtype, self.device())
 83 |     }
 84 |     /// This returns true only if a tensor with the passed in name is available. E.g. when passed
 85 |     /// `a`, true is returned if `prefix.a` exists but false is returned if only `prefix.a.b`
 86 |     /// exists.
 87 |     pub fn contains_tensor(&self, tensor_name: &str) -> bool {
 88 |         let path = self.path(tensor_name);
 89 | 
 90 |         let tensor_data = self.vm.data().lock().unwrap();
 91 |         tensor_data.contains_key(path.as_str())
 92 |     }
 93 | 
 94 |     /// The device used for variables.
 95 |     pub fn device(&self) -> &Device {
 96 |         &self.device
 97 |     }
 98 | 
 99 |     /// The dtype used for variables.
100 |     pub fn dtype(&self) -> DType {
101 |         self.dtype
102 |     }
103 | 
104 |     /// Prevents setting variables on this context and all parent/child contexts
105 |     pub fn freeze(&self) {
106 |         self.frozen.store(true, Relaxed)
107 |     }
108 | 
109 |     /// Allows setting variables on this context and all parent/child contexts
110 |     pub fn unfreeze(&self) {
111 |         self.frozen.store(false, Relaxed)
112 |     }
113 | }
114 | 
115 | impl std::fmt::Debug for Context {
116 |     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
117 |         f.debug_struct("Context").field("path", &self.path).finish()
118 |     }
119 | }
120 | 


--------------------------------------------------------------------------------
/mamba-ssm/src/lib.rs:
--------------------------------------------------------------------------------
  1 | #[cfg(feature = "mkl")]
  2 | extern crate intel_mkl_src;
  3 | 
  4 | #[cfg(feature = "accelerate")]
  5 | extern crate accelerate_src;
  6 | 
  7 | use anyhow::{anyhow, Error as E, Result};
  8 | 
  9 | use candle::{DType, Device, Module, Tensor};
 10 | use candle_transformers::generation::LogitsProcessor;
 11 | use tokenizers::Tokenizer;
 12 | 
 13 | pub mod model;
 14 | use model::Model;
 15 | 
 16 | mod utils;
 17 | use utils::TokenOutputStream;
 18 | 
 19 | pub mod context;
 20 | pub mod nn;
 21 | pub mod ops;
 22 | 
 23 | pub struct TextGeneration {
 24 |     model: Model,
 25 |     device: Device,
 26 |     tokenizer: TokenOutputStream,
 27 |     logits_processor: LogitsProcessor,
 28 |     repeat_penalty: f32,
 29 |     repeat_last_n: usize,
 30 | }
 31 | 
 32 | impl TextGeneration {
 33 |     #[allow(clippy::too_many_arguments)]
 34 |     pub fn new(
 35 |         model: Model,
 36 |         tokenizer: Tokenizer,
 37 |         seed: u64,
 38 |         temp: Option<f64>,
 39 |         top_p: Option<f64>,
 40 |         repeat_penalty: f32,
 41 |         repeat_last_n: usize,
 42 |         device: &Device,
 43 |     ) -> Self {
 44 |         let logits_processor = LogitsProcessor::new(seed, temp, top_p);
 45 |         Self {
 46 |             model,
 47 |             tokenizer: TokenOutputStream::new(tokenizer),
 48 |             logits_processor,
 49 |             repeat_penalty,
 50 |             repeat_last_n,
 51 |             device: device.clone(),
 52 |         }
 53 |     }
 54 | 
 55 |     pub fn run(&mut self, prompt: &str, sample_len: usize) -> Result<()> {
 56 |         if prompt.is_empty() {
 57 |             anyhow::bail!("Prompt cannot be empty");
 58 |         }
 59 |         use std::io::Write;
 60 |         self.tokenizer.clear();
 61 |         let mut tokens = self
 62 |             .tokenizer
 63 |             .tokenizer()
 64 |             .encode(prompt, true)
 65 |             .map_err(E::msg)?
 66 |             .get_ids()
 67 |             .to_vec();
 68 | 
 69 |         // Prompt processing
 70 |         let start_process = std::time::Instant::now();
 71 |         // Just process all but the last token of the prompt to give the model
 72 |         // the last token to work with when generation starts
 73 |         let (last, tokens_m1) = tokens
 74 |             .split_last()
 75 |             .ok_or(anyhow!("Tokens cannot be empty"))?;
 76 |         let mut next_token = *last;
 77 |         let input = Tensor::new(tokens_m1, &self.device)?.unsqueeze(0)?;
 78 |         let _ = self.model.forward(&input)?;
 79 |         let dt = start_process.elapsed();
 80 |         let prompt_len = tokens.len() - 1;
 81 |         println!(
 82 |             "\nPrompt processing time ({} tokens at {:.2} token/s)",
 83 |             prompt_len,
 84 |             prompt_len as f64 / dt.as_secs_f64(),
 85 |         );
 86 |         for &t in tokens.iter() {
 87 |             if let Some(t) = self.tokenizer.next_token(t)? {
 88 |                 print!("{t}")
 89 |             }
 90 |         }
 91 |         std::io::stdout().flush()?;
 92 | 
 93 |         let mut generated_tokens = 0usize;
 94 |         let eos_token = match self.tokenizer.get_token("<|endoftext|>") {
 95 |             Some(token) => token,
 96 |             None => anyhow::bail!("cannot find the </s> token"),
 97 |         };
 98 |         let start_gen = std::time::Instant::now();
 99 |         for _ in 0..sample_len {
100 |             let input = Tensor::new(&[next_token], &self.device)?.unsqueeze(0)?;
101 |             let logits = self.model.forward(&input)?;
102 |             let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
103 |             let logits = if self.repeat_penalty == 1. {
104 |                 logits
105 |             } else {
106 |                 let start_at = tokens.len().saturating_sub(self.repeat_last_n);
107 |                 candle_transformers::utils::apply_repeat_penalty(
108 |                     &logits,
109 |                     self.repeat_penalty,
110 |                     &tokens[start_at..],
111 |                 )?
112 |             };
113 | 
114 |             next_token = self.logits_processor.sample(&logits)?;
115 |             tokens.push(next_token);
116 | 
117 |             generated_tokens += 1;
118 |             if next_token == eos_token {
119 |                 break;
120 |             }
121 |             if let Some(t) = self.tokenizer.next_token(next_token)? {
122 |                 print!("{t}");
123 |                 std::io::stdout().flush()?;
124 |             }
125 |         }
126 |         let dt = start_gen.elapsed();
127 |         if let Some(rest) = self.tokenizer.decode_rest().map_err(E::msg)? {
128 |             print!("{rest}");
129 |         }
130 |         std::io::stdout().flush()?;
131 |         println!(
132 |             "\n{generated_tokens} tokens generated ({:.2} token/s)",
133 |             generated_tokens as f64 / dt.as_secs_f64(),
134 |         );
135 |         Ok(())
136 |     }
137 | }
138 | 


--------------------------------------------------------------------------------
/mamba-ssm/src/model.rs:
--------------------------------------------------------------------------------
  1 | /// This is mostly an adaptation of the example mamba implementation found at
  2 | /// https://github.com/huggingface/candle/tree/37c539f2b7dfc8aa67a10b611dc12e5e0428be00/candle-examples/examples/mamba-minimal
  3 | /// which is itself based on
  4 | /// https://github.com/johnma2006/mamba-minimal/blob/master/model.py
  5 | /// Simple, minimal implementation of Mamba in one file of PyTorch.
  6 | ///
  7 | /// Refer to the paper for details:
  8 | /// Mamba: Linear-Time Sequence Modeling with Selective State Spaces (Albert Gu and Tri Dao)
  9 | /// https://arxiv.org/abs/2312.00752
 10 | ///
 11 | use candle::{Module, Result, Tensor, D};
 12 | use candle_nn::{RmsNorm, VarBuilder};
 13 | 
 14 | use crate::context::Context;
 15 | use crate::nn::selective_ssm::SSM;
 16 | use crate::nn::{conv1d, linear_no_bias, Conv1d, Linear};
 17 | 
 18 | #[derive(Debug, Clone, serde::Deserialize)]
 19 | pub struct Config {
 20 |     /// The model dimension (number of channels in the input sequence)
 21 |     d_model: usize,
 22 |     n_layer: usize,
 23 |     vocab_size: usize,
 24 |     pad_vocab_size_multiple: usize,
 25 | }
 26 | 
 27 | impl Config {
 28 |     fn vocab_size(&self) -> usize {
 29 |         let pad = self.pad_vocab_size_multiple;
 30 |         (self.vocab_size + pad - 1) / pad * pad
 31 |     }
 32 | 
 33 |     /// The rank of Δ
 34 |     fn dt_rank(&self) -> usize {
 35 |         (self.d_model + 15) / 16
 36 |     }
 37 | 
 38 |     /// Size of the conv1d kernel that's built into the model, but isn't in the
 39 |     /// config.json
 40 |     fn d_conv(&self) -> usize {
 41 |         4
 42 |     }
 43 | 
 44 |     /// N, the SSM state dimension
 45 |     fn d_state(&self) -> usize {
 46 |         16
 47 |     }
 48 | 
 49 |     /// This is the model dimension multiplied by the expansion factor 2.
 50 |     fn d_inner(&self) -> usize {
 51 |         self.d_model * 2
 52 |     }
 53 | }
 54 | 
 55 | #[derive(Clone, Debug)]
 56 | pub struct MambaBlock {
 57 |     in_proj: Linear,
 58 |     conv1d: Conv1d,
 59 |     ssm: SSM,
 60 |     out_proj: Linear,
 61 | }
 62 | 
 63 | impl MambaBlock {
 64 |     pub fn new(cfg: &Config, vb: VarBuilder, ctx: Context) -> Result<Self> {
 65 |         let d_inner = cfg.d_inner();
 66 |         let d_conv = cfg.d_conv();
 67 | 
 68 |         let d_state = cfg.d_state();
 69 |         let dt_rank = cfg.dt_rank();
 70 | 
 71 |         let in_proj = linear_no_bias(cfg.d_model, d_inner * 2, vb.pp("in_proj"))?;
 72 |         let conv_cfg = candle_nn::Conv1dConfig {
 73 |             groups: d_inner,
 74 |             padding: d_conv - 1,
 75 |             ..Default::default()
 76 |         };
 77 |         let conv1d = conv1d(
 78 |             d_inner,
 79 |             d_inner,
 80 |             d_conv,
 81 |             conv_cfg,
 82 |             vb.pp("conv1d"),
 83 |             ctx.pp("conv1d"),
 84 |         )?;
 85 |         let ssm = SSM::new(d_inner, d_state, dt_rank, vb.clone(), ctx.pp("ssm"))?;
 86 |         let out_proj = linear_no_bias(d_inner, cfg.d_model, vb.pp("out_proj"))?;
 87 | 
 88 |         Ok(Self {
 89 |             in_proj,
 90 |             conv1d,
 91 |             ssm,
 92 |             out_proj,
 93 |         })
 94 |     }
 95 | }
 96 | 
 97 | impl Module for MambaBlock {
 98 |     fn forward(&self, xs: &Tensor) -> Result<Tensor> {
 99 |         // Project the input to two larger inputs as per the Mamba
100 |         // architecture in Figure 3 of the paper.
101 |         let xs_and_res = xs.apply(&self.in_proj)?.chunk(2, D::Minus1)?;
102 |         let (xs, res) = (&xs_and_res[0], &xs_and_res[1]);
103 | 
104 |         // Conv1d -> SiLU -> SSM
105 |         let xs = xs.apply(&self.conv1d)?;
106 |         let xs = candle_nn::ops::silu(&xs)?;
107 |         let xs = xs.apply(&self.ssm)?;
108 | 
109 |         // Multiplicative gate
110 |         let ys = (xs * candle_nn::ops::silu(res))?;
111 |         ys.apply(&self.out_proj)
112 |     }
113 | }
114 | 
115 | #[derive(Clone, Debug)]
116 | pub struct ResMambaBlock {
117 |     mixer: MambaBlock,
118 |     norm: RmsNorm,
119 | }
120 | 
121 | impl ResMambaBlock {
122 |     pub fn new(cfg: &Config, vb: VarBuilder, ctx: Context) -> Result<Self> {
123 |         let norm = candle_nn::rms_norm(cfg.d_model, 1e-5, vb.pp("norm"))?;
124 |         let mixer = MambaBlock::new(cfg, vb.pp("mixer"), ctx.pp("mamba"))?;
125 |         Ok(Self { mixer, norm })
126 |     }
127 | }
128 | 
129 | impl Module for ResMambaBlock {
130 |     fn forward(&self, xs: &Tensor) -> Result<Tensor> {
131 |         xs.apply(&self.norm)?.apply(&self.mixer)? + xs
132 |     }
133 | }
134 | 
135 | #[derive(Clone, Debug)]
136 | pub struct Model {
137 |     embedding: candle_nn::Embedding,
138 |     layers: Vec<ResMambaBlock>,
139 |     norm_f: RmsNorm,
140 |     lm_head: Linear,
141 | }
142 | 
143 | impl Model {
144 |     pub fn new(cfg: &Config, vb: VarBuilder, ctx: Context) -> Result<Self> {
145 |         let embedding = candle_nn::embedding(cfg.vocab_size(), cfg.d_model, vb.pp("embedding"))?;
146 |         let mut layers = Vec::with_capacity(cfg.n_layer);
147 |         let vb_l = vb.pp("layers");
148 |         for layer_idx in 0..cfg.n_layer {
149 |             let layer = ResMambaBlock::new(cfg, vb_l.pp(layer_idx), ctx.pp(layer_idx))?;
150 |             layers.push(layer)
151 |         }
152 |         let norm_f = candle_nn::rms_norm(cfg.d_model, 1e-5, vb.pp("norm_f"))?;
153 |         let lm_head = Linear::from_weights(embedding.embeddings().clone(), None)?;
154 |         Ok(Self {
155 |             embedding,
156 |             layers,
157 |             norm_f,
158 |             lm_head,
159 |         })
160 |     }
161 | }
162 | 
163 | impl Module for Model {
164 |     fn forward(&self, input_ids: &Tensor) -> Result<Tensor> {
165 |         let (_batch_size, seq_len) = input_ids.dims2()?;
166 |         let mut xs = self.embedding.forward(input_ids)?;
167 |         for layer in self.layers.iter() {
168 |             xs = layer.forward(&xs)?
169 |         }
170 |         xs.narrow(1, seq_len - 1, 1)?
171 |             .apply(&self.norm_f)?
172 |             .apply(&self.lm_head)
173 |     }
174 | }
175 | 


--------------------------------------------------------------------------------
/mamba-ssm/src/nn/conv1d.rs:
--------------------------------------------------------------------------------
  1 | use candle::{IndexOp, Module, Result, Tensor};
  2 | use candle_nn::{Conv1dConfig, VarBuilder};
  3 | 
  4 | use crate::context::Context;
  5 | 
  6 | /// Faster Conv1d with some more traceability
  7 | #[derive(Debug, Clone)]
  8 | pub struct Conv1d {
  9 |     inner: Conv1dImpl,
 10 |     span: tracing::Span,
 11 | }
 12 | 
 13 | impl Module for Conv1d {
 14 |     fn forward(&self, x: &Tensor) -> Result<Tensor> {
 15 |         let _enter = self.span.enter();
 16 |         match self.inner {
 17 |             Conv1dImpl::TC(ref c) => c.forward(x),
 18 |             Conv1dImpl::Direct(ref c) => {
 19 |                 let (_, seq_len, _) = x.dims3()?;
 20 |                 c.forward(&x.t()?)?
 21 |                     .narrow(candle::D::Minus1, 0, seq_len)?
 22 |                     .t()
 23 |             }
 24 |         }
 25 |     }
 26 | }
 27 | 
 28 | #[derive(Debug, Clone)]
 29 | enum Conv1dImpl {
 30 |     TC(TCConv1d),
 31 |     Direct(candle_nn::Conv1d),
 32 | }
 33 | 
 34 | pub fn conv1d(
 35 |     in_channels: usize,
 36 |     out_channels: usize,
 37 |     kernel_size: usize,
 38 |     cfg: Conv1dConfig,
 39 |     vs: VarBuilder,
 40 |     ctx: Context,
 41 | ) -> Result<Conv1d> {
 42 |     let span = tracing::span!(
 43 |         tracing::Level::TRACE,
 44 |         "conv1d",
 45 |         in_channels,
 46 |         out_channels,
 47 |         kernel_size
 48 |     );
 49 |     if in_channels == out_channels && cfg.groups == out_channels && kernel_size == 4 {
 50 |         let inner = TCConv1d::new(in_channels, out_channels, kernel_size, cfg, vs, ctx)?;
 51 |         Ok(Conv1d {
 52 |             inner: Conv1dImpl::TC(inner),
 53 |             span,
 54 |         })
 55 |     } else {
 56 |         let inner = candle_nn::conv1d(in_channels, out_channels, kernel_size, cfg, vs)?;
 57 |         Ok(Conv1d {
 58 |             inner: Conv1dImpl::Direct(inner),
 59 |             span,
 60 |         })
 61 |     }
 62 | }
 63 | 
 64 | /// Winograd transform matrices for convolution F(2,4) evaluated at points
 65 | /// [0, 0.5, -1.5, 1.5] generated with the help of https://github.com/andravin/wincnn
 66 | /// See:
 67 | /// [1] Error Analysis and Improving the Accuracy of Winograd Convolution for Deep Neural Networks
 68 | ///     Barbara Barabasz, Andrew Anderson, Kirk M. Soodhalter, David Gregg
 69 | ///     https://arxiv.org/abs/1803.10986
 70 | /// [2] Winograd Convolution for Deep Neural Networks: Efficient Point Selection
 71 | ///     Syed Asad Alam, Andrew Anderson, Barbara Barabasz, David Gregg
 72 | ///     https://arxiv.org/pdf/2201.10369.pdf
 73 | ///
 74 | /// Originally used [0, -1, 1, 0.5] as the set of points from [1] but now using
 75 | /// a set of points from [2] shown to have better FP error (E=4.69E-08 for 1D
 76 | /// conv of kernel size 3 for n=5)
 77 | const A_T: &[[f32; 5]; 2] = &[[1., 1., 1., 1., 0.], [0., 0.5, -1.5, 1.5, 1.]];
 78 | const G: &[[f32; 4]; 5] = &[
 79 |     [8. / 9., 0., 0., 0.],
 80 |     [-1., -0.5, -0.25, -0.125],
 81 |     [-1. / 9., 1. / 6., -0.25, 0.375],
 82 |     [2. / 9., 1. / 3., 0.5, 0.75],
 83 |     [0., 0., 0., 1.],
 84 | ];
 85 | const B_T: &[[f32; 5]; 5] = &[
 86 |     [1.125, -2.25, -0.5, 1.0, 0.],
 87 |     [0., -2.25, 0., 1.0, 0.],
 88 |     [0., 0.75, -2.0, 1.0, 0.],
 89 |     [0., -0.75, 1.0, 1.0, 0.],
 90 |     [0., 1.125, -2.25, -0.5, 1.],
 91 | ];
 92 | 
 93 | /// Basic implementation of Toom-Cook (Winograd) conv1d (F(2,4)). Only
 94 | /// implemented for the use case of Mamba.
 95 | #[derive(Debug, Clone)]
 96 | pub struct TCConv1d {
 97 |     // Winograd transform matrices
 98 |     at: Tensor,
 99 |     bt: Tensor,
100 | 
101 |     /// The cached result of multiplying G with the stack of kernels
102 |     gk: Tensor,
103 |     bias: Option<Tensor>,
104 |     cfg: Conv1dConfig,
105 | 
106 |     ctx: Context,
107 | }
108 | 
109 | impl TCConv1d {
110 |     fn new(
111 |         in_channels: usize,
112 |         out_channels: usize,
113 |         kernel_size: usize,
114 |         cfg: Conv1dConfig,
115 |         vb: VarBuilder,
116 |         ctx: Context,
117 |     ) -> Result<Self> {
118 |         if in_channels != cfg.groups {
119 |             candle::bail!(
120 |                 "Mismatch between in_channels ({in_channels}) and groups ({})",
121 |                 cfg.groups
122 |             );
123 |         }
124 |         let init_ws = candle_nn::init::DEFAULT_KAIMING_NORMAL;
125 |         let ws = vb.get_with_hints((out_channels, 1, kernel_size), "weight", init_ws)?;
126 |         let bound = 1. / (in_channels as f64).sqrt();
127 | 
128 |         // Change the kernel's layout to be (kernel_size, channels)
129 |         let kernel = ws.permute((1, 2, 0))?.squeeze(0)?;
130 | 
131 |         let dtype = ws.dtype();
132 | 
133 |         let at = Tensor::new(A_T, vb.device())?.to_dtype(dtype)?;
134 |         let bt = Tensor::new(B_T, vb.device())?.to_dtype(dtype)?;
135 |         // Save the result of multiplying G and the stack of kernels since we
136 |         // don't need to use the kernel weights directly anymore
137 |         let g = Tensor::new(G, vb.device())?.to_dtype(dtype)?;
138 |         let gk = g.matmul(&kernel)?;
139 | 
140 |         let init_bs = candle_nn::Init::Uniform {
141 |             lo: -bound,
142 |             up: bound,
143 |         };
144 |         let bs = vb.get_with_hints(out_channels, "bias", init_bs)?;
145 | 
146 |         Ok(Self {
147 |             at,
148 |             bt,
149 |             gk,
150 |             bias: Some(bs),
151 |             cfg,
152 |             ctx,
153 |         })
154 |     }
155 | }
156 | 
157 | impl Module for TCConv1d {
158 |     fn forward(&self, x: &Tensor) -> Result<Tensor> {
159 |         // Assume out_channels == in_channels
160 |         let (b_size, seq_len, c_in) = x.dims3()?;
161 | 
162 |         let batches = (0..b_size)
163 |             .map(|i| {
164 |                 let view = x.i((i, .., ..))?;
165 |                 let mut ctx = self.ctx.pp(i);
166 |                 let padding = ctx.get((self.cfg.padding, c_in), "ph")?;
167 |                 let view = Tensor::cat(&[padding, view], 0)?;
168 | 
169 |                 let conv = tc_conv1d(&view, &self.at, &self.bt, &self.gk);
170 | 
171 |                 // Since view is [seq_len + padding, c_in], this saves off
172 |                 // padding inputs from the end of the view
173 |                 ctx.set("ph", view.i((seq_len.., ..))?)?;
174 | 
175 |                 conv
176 |             })
177 |             .collect::<candle::Result<Vec<_>>>()?;
178 |         let output = Tensor::stack(&batches, 0)?;
179 | 
180 |         match &self.bias {
181 |             None => Ok(output),
182 |             Some(bias) => {
183 |                 let b = bias.dims1()?;
184 |                 let bias = bias.reshape((1, 1, b))?;
185 |                 Ok(output.broadcast_add(&bias)?)
186 |             }
187 |         }
188 |     }
189 | }
190 | 
191 | /// Given an (L+K-1)xD tensor, applies Toom-Cook 1D convolution on it via the
192 | /// provided transformation matrices, returns an LxD tensor.
193 | /// Assumes input is already left-padded and that
194 | /// groups == in_channels == out_channels
195 | fn tc_conv1d(input: &Tensor, at: &Tensor, bt: &Tensor, gk: &Tensor) -> Result<Tensor> {
196 |     // TODO: Support convolutions other than F(2,4)
197 |     let tile_size = 5;
198 |     let tile_stride = 2;
199 | 
200 |     let tiles = tile_conv_input(&input, tile_size, tile_stride)?;
201 |     let outputs = tiles
202 |         .iter()
203 |         .map(|tile| {
204 |             let (chunk_size, _) = tile.dims2()?;
205 |             let right_padding = tile_size - chunk_size;
206 |             assert!(
207 |                 tile_stride > right_padding,
208 |                 "conv tiling returned too small of a chunk size {chunk_size}"
209 |             );
210 |             let tile = if right_padding > 0 {
211 |                 tile.pad_with_zeros(0, 0, right_padding)?
212 |             } else {
213 |                 tile.clone()
214 |             };
215 |             let btd = bt.matmul(&tile)?;
216 |             let elemwise_prod = gk.mul(&btd)?;
217 |             let res = at.matmul(&elemwise_prod)?;
218 |             if right_padding > 0 {
219 |                 res.narrow(0, 0, tile_stride - right_padding)
220 |             } else {
221 |                 Ok(res)
222 |             }
223 |         })
224 |         .collect::<Result<Vec<_>>>()?;
225 |     Tensor::cat(&outputs, 0)
226 | }
227 | 
228 | /// Given a LxD tensor, returns a vec of tensor views each with length
229 | /// tile_size, advancing with tile_stride. Assumes input is already left-padded
230 | /// as desired. If the input is not cleanly tilable by the tile size and tile
231 | /// stride, the last tensor view will contain the remainders.
232 | fn tile_conv_input(input: &Tensor, tile_size: usize, tile_stride: usize) -> Result<Vec<Tensor>> {
233 |     let length = input.dims()[0];
234 |     if length < tile_size {
235 |         return Ok(vec![input.clone()]);
236 |     }
237 | 
238 |     let tilable_size = (length - tile_size) + tile_stride;
239 |     let remainder = tilable_size % tile_stride;
240 |     let num_tiles = tilable_size / tile_stride;
241 |     let mut res = vec![];
242 |     for i in 0..num_tiles {
243 |         res.push(input.i((tile_stride * i..(tile_stride * i) + tile_size, ..))?);
244 |     }
245 |     if remainder > 0 {
246 |         res.push(input.i((tile_stride * num_tiles.., ..))?);
247 |     }
248 |     Ok(res)
249 | }
250 | 
251 | #[cfg(test)]
252 | mod tests {
253 |     use candle::{IndexOp, Tensor};
254 | 
255 |     use super::{tc_conv1d, tile_conv_input, A_T, B_T, G};
256 | 
257 |     /// For some reason candle's conv1d implementation returns the wrong output
258 |     /// for groups > 1 in test so here's a reimplementation of it for
259 |     /// c_in == groups
260 |     /// There doesn't seem to be a problem with it at runtime though
261 |     fn candle_conv1d(
262 |         input: &Tensor,
263 |         kernel: &Tensor,
264 |         padding: usize,
265 |         stride: usize,
266 |         dilation: usize,
267 |         groups: usize,
268 |     ) -> candle::Result<Tensor> {
269 |         let (_c_out, c_in_k, k_size) = kernel.dims3()?;
270 |         let (b_size, c_in, l_in) = input.dims3()?;
271 |         if c_in != groups {
272 |             candle::bail!("c_in must match number of groups")
273 |         }
274 | 
275 |         let blocks = (0..groups)
276 |             .map(|i| {
277 |                 let v = input.i((.., i, ..))?.reshape((b_size, 1, l_in))?;
278 |                 let k = kernel.i((i, .., ..))?.reshape((1, c_in_k, k_size))?;
279 |                 v.conv1d(&k, padding, stride, dilation, 1)
280 |             })
281 |             .collect::<candle::Result<Vec<_>>>()?;
282 |         Tensor::cat(&blocks, 1)
283 |     }
284 | 
285 |     #[test]
286 |     fn test_tc_conv1d() -> candle::Result<()> {
287 |         let device = &candle::Device::Cpu;
288 |         let x = Tensor::new(
289 |             &[
290 |                 [1f32, 2., 3., 4., 5.],
291 |                 [1.1, 2.1, 3.1, 4.1, 5.1],
292 |                 [1.2, 2.2, 3.2, 4.2, 5.2],
293 |                 [1.3, 2.3, 3.3, 4.3, 5.3],
294 |                 [1.4, 2.4, 3.4, 4.4, 5.4],
295 |                 [1.5, 2.5, 3.5, 4.5, 5.5],
296 |                 [1.6, 2.6, 3.6, 4.6, 5.6],
297 |                 [1.7, 2.7, 3.7, 4.7, 5.7],
298 |                 [1.8, 2.8, 3.8, 4.8, 5.8],
299 |                 [1.9, 2.9, 3.9, 4.9, 5.9],
300 |                 [2., 3., 4., 5., 6.],
301 |             ],
302 |             device,
303 |         )?;
304 |         let (seq_len, _channels) = x.dims2()?;
305 | 
306 |         let kernels = Tensor::new(
307 |             &[
308 |                 [0.5f32, 0.25, 0.125, 0.125, 0.125],
309 |                 [1.5, 1.25, 1.125, 1.125, 1.125],
310 |                 [2.5, 2.25, 2.125, 2.125, 2.125],
311 |                 [3.5, 3.25, 3.125, 3.125, 3.125],
312 |             ],
313 |             device,
314 |         )?;
315 |         let (kernel_size, channels) = kernels.dims2()?;
316 | 
317 |         // Compare with original candle conv1d implementation
318 |         let x_copy = x.unsqueeze(0)?.t()?;
319 |         let kernels_copy = kernels.unsqueeze(0)?;
320 |         let kernel = kernels_copy.permute((2, 0, 1))?;
321 | 
322 |         let original_conv = candle_conv1d(&x_copy, &kernel, kernel_size - 1, 1, 1, channels)?
323 |             .narrow(candle::D::Minus1, 0, seq_len)?
324 |             .t()?
325 |             .squeeze(0)?;
326 | 
327 |         // Perform optimized conv1d
328 |         let at = Tensor::new(A_T, device)?;
329 |         let bt = Tensor::new(B_T, device)?;
330 |         let g = Tensor::new(G, device)?;
331 |         let gk = g.matmul(&kernels)?;
332 | 
333 |         let x = x.pad_with_zeros(0, 3, 0)?;
334 |         let conv_output = tc_conv1d(&x, &at, &bt, &gk)?;
335 |         assert_eq!(conv_output.dims(), &[seq_len, channels]);
336 | 
337 |         println!("Original\n{}", original_conv);
338 |         println!("Conv Output\n{}", conv_output);
339 | 
340 |         let conv_diff = (original_conv - conv_output)?
341 |             .sum_all()?
342 |             .to_scalar::<f32>()?;
343 |         assert!(conv_diff.abs() < 1e-4);
344 |         Ok(())
345 |     }
346 | 
347 |     #[test]
348 |     fn test_tile_conv_input() -> candle::Result<()> {
349 |         // Simulate data with padding already added
350 |         let x = Tensor::new(
351 |             &[
352 |                 [1f32, 2., 3., 4.],
353 |                 [1.1, 2.1, 3.1, 4.1],
354 |                 [1.2, 2.2, 3.2, 4.2],
355 |                 [1.3, 2.3, 3.3, 4.3],
356 |                 [1.4, 2.4, 3.4, 4.4],
357 |                 [1.5, 2.5, 3.5, 4.5],
358 |                 [1.6, 2.6, 3.6, 4.6],
359 |                 [1.7, 2.7, 3.7, 4.7],
360 |                 [1.8, 2.8, 3.8, 4.8],
361 |                 [1.9, 2.9, 3.9, 4.9],
362 |             ],
363 |             &candle::Device::Cpu,
364 |         )?;
365 |         let x = x.pad_with_zeros(0, 3, 0)?;
366 | 
367 |         let views = tile_conv_input(&x, 5, 2)?;
368 | 
369 |         assert_eq!(views.len(), 5);
370 |         assert_eq!(
371 |             views[0].to_vec2::<f32>()?,
372 |             &[
373 |                 [0f32, 0., 0., 0.],
374 |                 [0., 0., 0., 0.],
375 |                 [0., 0., 0., 0.],
376 |                 [1., 2., 3., 4.],
377 |                 [1.1, 2.1, 3.1, 4.1],
378 |             ],
379 |         );
380 |         assert_eq!(
381 |             views[1].to_vec2::<f32>()?,
382 |             &[
383 |                 [0f32, 0., 0., 0.],
384 |                 [1., 2., 3., 4.],
385 |                 [1.1, 2.1, 3.1, 4.1],
386 |                 [1.2, 2.2, 3.2, 4.2],
387 |                 [1.3, 2.3, 3.3, 4.3],
388 |             ],
389 |         );
390 |         Ok(())
391 |     }
392 | 
393 |     #[test]
394 |     fn test_tile_conv_input_odd_size() -> candle::Result<()> {
395 |         // Simulate data with padding already added
396 |         let x = Tensor::new(
397 |             &[
398 |                 [1f32, 2., 3., 4.],
399 |                 [1.1, 2.1, 3.1, 4.1],
400 |                 [1.2, 2.2, 3.2, 4.2],
401 |                 [1.3, 2.3, 3.3, 4.3],
402 |                 [1.4, 2.4, 3.4, 4.4],
403 |                 [1.5, 2.5, 3.5, 4.5],
404 |                 [1.6, 2.6, 3.6, 4.6],
405 |                 [1.7, 2.7, 3.7, 4.7],
406 |                 [1.8, 2.8, 3.8, 4.8],
407 |                 [1.9, 2.9, 3.9, 4.9],
408 |                 [2., 3., 4., 5.],
409 |             ],
410 |             &candle::Device::Cpu,
411 |         )?;
412 |         let x = x.pad_with_zeros(0, 3, 0)?;
413 | 
414 |         let views = tile_conv_input(&x, 5, 2)?;
415 | 
416 |         assert_eq!(views.len(), 6);
417 |         assert_eq!(
418 |             views[0].to_vec2::<f32>()?,
419 |             &[
420 |                 [0f32, 0., 0., 0.],
421 |                 [0., 0., 0., 0.],
422 |                 [0., 0., 0., 0.],
423 |                 [1., 2., 3., 4.],
424 |                 [1.1, 2.1, 3.1, 4.1],
425 |             ],
426 |         );
427 |         assert_eq!(
428 |             views[5].to_vec2::<f32>()?,
429 |             &[
430 |                 [1.7f32, 2.7, 3.7, 4.7],
431 |                 [1.8, 2.8, 3.8, 4.8],
432 |                 [1.9, 2.9, 3.9, 4.9],
433 |                 [2., 3., 4., 5.],
434 |             ],
435 |         );
436 |         Ok(())
437 |     }
438 | 
439 |     #[test]
440 |     fn test_tile_conv_input_large_tile() -> candle::Result<()> {
441 |         // Simulate data with padding already added
442 |         let x = Tensor::new(
443 |             &[
444 |                 [1f32, 2., 3., 4.],
445 |                 [1.1, 2.1, 3.1, 4.1],
446 |                 [1.2, 2.2, 3.2, 4.2],
447 |                 [1.3, 2.3, 3.3, 4.3],
448 |                 [1.4, 2.4, 3.4, 4.4],
449 |                 [1.5, 2.5, 3.5, 4.5],
450 |                 [1.6, 2.6, 3.6, 4.6],
451 |                 [1.7, 2.7, 3.7, 4.7],
452 |                 [1.8, 2.8, 3.8, 4.8],
453 |                 [1.9, 2.9, 3.9, 4.9],
454 |             ],
455 |             &candle::Device::Cpu,
456 |         )?;
457 |         let x = x.pad_with_zeros(0, 3, 0)?;
458 | 
459 |         let views = tile_conv_input(&x, 10, 7)?;
460 | 
461 |         assert_eq!(views.len(), 2);
462 |         assert_eq!(
463 |             views[0].to_vec2::<f32>()?,
464 |             &[
465 |                 [0f32, 0., 0., 0.],
466 |                 [0., 0., 0., 0.],
467 |                 [0., 0., 0., 0.],
468 |                 [1., 2., 3., 4.],
469 |                 [1.1, 2.1, 3.1, 4.1],
470 |                 [1.2, 2.2, 3.2, 4.2],
471 |                 [1.3, 2.3, 3.3, 4.3],
472 |                 [1.4, 2.4, 3.4, 4.4],
473 |                 [1.5, 2.5, 3.5, 4.5],
474 |                 [1.6, 2.6, 3.6, 4.6],
475 |             ],
476 |         );
477 |         assert_eq!(
478 |             views[1].to_vec2::<f32>()?,
479 |             &[
480 |                 [1.4f32, 2.4, 3.4, 4.4],
481 |                 [1.5, 2.5, 3.5, 4.5],
482 |                 [1.6, 2.6, 3.6, 4.6],
483 |                 [1.7, 2.7, 3.7, 4.7],
484 |                 [1.8, 2.8, 3.8, 4.8],
485 |                 [1.9, 2.9, 3.9, 4.9],
486 |             ],
487 |         );
488 |         Ok(())
489 |     }
490 | }
491 | 


--------------------------------------------------------------------------------
/mamba-ssm/src/nn/linear.rs:
--------------------------------------------------------------------------------
  1 | use candle::{Module, Result, Tensor};
  2 | use candle_nn::VarBuilder;
  3 | 
  4 | /// Linear with some more traceability.
  5 | /// The layer applies a linear transformation to the incoming data,
  6 | /// `y = x@w.t() + b` with an optional bias.
  7 | #[derive(Debug, Clone)]
  8 | pub struct Linear {
  9 |     inner: LinearInner,
 10 |     span: tracing::Span,
 11 | }
 12 | 
 13 | impl Linear {
 14 |     pub fn from_weights(weights: Tensor, bias: Option<Tensor>) -> Result<Self> {
 15 |         let dims = weights.shape().clone();
 16 |         assert!(dims.elem_count() >= 2);
 17 |         let in_dim = dims.dims()[0];
 18 |         let out_dim = dims.dims()[1];
 19 |         let inner = LinearInner::new(weights, bias)?;
 20 |         let span = tracing::span!(tracing::Level::TRACE, "linear", in_dim, out_dim);
 21 |         Ok(Self { inner, span })
 22 |     }
 23 | }
 24 | 
 25 | pub fn linear(in_dim: usize, out_dim: usize, vb: VarBuilder) -> Result<Linear> {
 26 |     let vb_prefix = vb.prefix();
 27 | 
 28 |     let init_ws = candle_nn::init::DEFAULT_KAIMING_NORMAL;
 29 |     let ws = vb.get_with_hints((out_dim, in_dim), "weight", init_ws)?;
 30 |     let bound = 1. / (in_dim as f64).sqrt();
 31 |     let init_bs = candle_nn::Init::Uniform {
 32 |         lo: -bound,
 33 |         up: bound,
 34 |     };
 35 | 
 36 |     let bs = vb.get_with_hints(out_dim, "bias", init_bs)?;
 37 |     let span = tracing::span!(tracing::Level::TRACE, "linear", vb_prefix, in_dim, out_dim);
 38 |     Ok(Linear {
 39 |         inner: LinearInner::new(ws, Some(bs))?,
 40 |         span,
 41 |     })
 42 | }
 43 | 
 44 | pub fn linear_no_bias(in_dim: usize, out_dim: usize, vb: VarBuilder) -> Result<Linear> {
 45 |     let vb_prefix = vb.prefix();
 46 |     let init_ws = candle_nn::init::DEFAULT_KAIMING_NORMAL;
 47 |     let ws = vb.get_with_hints((out_dim, in_dim), "weight", init_ws)?;
 48 |     let span = tracing::span!(
 49 |         tracing::Level::TRACE,
 50 |         "linear_no_bias",
 51 |         vb_prefix,
 52 |         in_dim,
 53 |         out_dim
 54 |     );
 55 |     Ok(Linear {
 56 |         inner: LinearInner::new(ws, None)?,
 57 |         span,
 58 |     })
 59 | }
 60 | 
 61 | impl Module for Linear {
 62 |     fn forward(&self, xs: &Tensor) -> Result<Tensor> {
 63 |         let _enter = self.span.enter();
 64 |         self.inner.forward(xs)
 65 |     }
 66 | }
 67 | 
 68 | /// LinearInner is basically the same as candle_nn::Linear
 69 | #[derive(Clone, Debug)]
 70 | struct LinearInner {
 71 |     weight: Tensor,
 72 |     bias: Option<Tensor>,
 73 | }
 74 | 
 75 | impl LinearInner {
 76 |     fn new(weight: Tensor, bias: Option<Tensor>) -> Result<Self> {
 77 |         Ok(Self { weight, bias })
 78 |     }
 79 | }
 80 | 
 81 | impl Module for LinearInner {
 82 |     fn forward(&self, x: &Tensor) -> Result<Tensor> {
 83 |         let x = match *x.dims() {
 84 |             [b1, b2, _, _] => x.matmul(&self.weight.broadcast_left((b1, b2))?.t()?)?,
 85 |             [bsize, _, _] => x.matmul(&self.weight.broadcast_left(bsize)?.t()?)?,
 86 |             _ => x.matmul(&self.weight.t()?)?,
 87 |         };
 88 |         match &self.bias {
 89 |             None => Ok(x),
 90 |             Some(bias) => x.broadcast_add(bias),
 91 |         }
 92 |     }
 93 | }
 94 | 
 95 | #[cfg(test)]
 96 | mod tests {
 97 |     use candle::{Module, Tensor};
 98 | 
 99 |     #[test]
100 |     fn test_linear() -> candle::Result<()> {
101 |         let device = &candle::Device::Cpu;
102 | 
103 |         let x = Tensor::randn(0f32, 1., (1, 10, 2000), &device)?;
104 | 
105 |         let var_map = candle_nn::VarMap::new();
106 |         let vb = candle_nn::VarBuilder::from_varmap(&var_map, candle::DType::F32, &device);
107 | 
108 |         let candle_linear = candle_nn::linear(2000, 2000, vb.clone())?;
109 |         let custom_linear = super::linear(2000, 2000, vb)?;
110 | 
111 |         let candle_res = candle_linear.forward(&x)?;
112 |         let custom_res = custom_linear.forward(&x)?;
113 | 
114 |         let res_diff = (candle_res - custom_res)?.sum_all()?.to_scalar::<f32>()?;
115 | 
116 |         // The difference should be 0 but this adds a bit of leniency
117 |         assert!(res_diff.abs() < 1e-8);
118 |         Ok(())
119 |     }
120 | }
121 | 


--------------------------------------------------------------------------------
/mamba-ssm/src/nn/mod.rs:
--------------------------------------------------------------------------------
1 | pub mod conv1d;
2 | pub mod linear;
3 | pub mod selective_ssm;
4 | 
5 | pub use conv1d::*;
6 | pub use linear::*;
7 | 


--------------------------------------------------------------------------------
/mamba-ssm/src/nn/selective_ssm.rs:
--------------------------------------------------------------------------------
  1 | use candle::{IndexOp, Module, Result, Tensor, D};
  2 | 
  3 | use candle_nn::VarBuilder;
  4 | 
  5 | use crate::{
  6 |     context::Context,
  7 |     nn::{linear, linear_no_bias, Linear},
  8 | };
  9 | 
 10 | #[derive(Clone, Debug)]
 11 | pub struct SSM {
 12 |     dt_rank: usize,
 13 | 
 14 |     at: Tensor,
 15 |     d: Tensor,
 16 | 
 17 |     x_proj: Linear,
 18 |     dt_proj: Linear,
 19 | 
 20 |     ctx: Context,
 21 | 
 22 |     span: tracing::Span,
 23 |     iter_span: tracing::Span,
 24 | }
 25 | 
 26 | impl SSM {
 27 |     pub fn new(
 28 |         d_inner: usize,
 29 |         d_state: usize,
 30 |         dt_rank: usize,
 31 |         vb: VarBuilder,
 32 |         ctx: Context,
 33 |     ) -> Result<Self> {
 34 |         let x_proj = linear_no_bias(d_inner, dt_rank + d_state * 2, vb.pp("x_proj"))?;
 35 |         let dt_proj = linear(dt_rank, d_inner, vb.pp("dt_proj"))?;
 36 | 
 37 |         let a_log = vb.get((d_inner, d_state), "A_log")?;
 38 |         let d = vb.get(d_inner, "D")?;
 39 | 
 40 |         let dtype = ctx.dtype();
 41 | 
 42 |         // Store the transpose of A as an optimization
 43 |         let at = a_log.to_dtype(dtype)?.exp()?.neg()?.t()?;
 44 |         let d = d.to_dtype(dtype)?;
 45 | 
 46 |         let span = tracing::span!(tracing::Level::TRACE, "SSM", d_inner, d_state, dt_rank);
 47 |         let iter_span = tracing::span!(tracing::Level::TRACE, "iter");
 48 | 
 49 |         Ok(Self {
 50 |             dt_rank,
 51 |             at,
 52 |             d,
 53 |             x_proj,
 54 |             dt_proj,
 55 | 
 56 |             ctx,
 57 | 
 58 |             span,
 59 |             iter_span,
 60 |         })
 61 |     }
 62 | 
 63 |     fn selective_scan(
 64 |         &self,
 65 |         u: &Tensor,     // The input (batch_size, seq_len, d_inner) aka (B L D)
 66 |         delta: &Tensor, // (batch_size, seq_len, d_inner) aka (B L D)
 67 |         b: &Tensor,     // (batch_size, seq_len, d_state) aka (B L N)
 68 |         c: &Tensor,     // (batch_size, seq_len, d_state) aka (B L N)
 69 |     ) -> Result<Tensor> {
 70 |         let (batch_size, seq_len, d_inner) = u.dims3()?;
 71 | 
 72 |         let at = &self.at; // (d_state, d_inner) aka (N, D)
 73 |         let d_state = at.dim(0)?;
 74 | 
 75 |         // Precompute delta * u since they are both the same size
 76 |         let delta_u = (delta * u)?;
 77 | 
 78 |         let mut bys = Vec::with_capacity(batch_size);
 79 |         for batch in 0..batch_size {
 80 |             // Load the SSM state from the context, or initialize with 0s
 81 |             let mut ctx = self.ctx.pp(batch);
 82 |             let mut x = ctx.get((d_state, d_inner), "ht")?; // h_(t-1), size (N D)
 83 | 
 84 |             let mut ys = Vec::with_capacity(seq_len);
 85 |             for i in 0..seq_len {
 86 |                 let _enter_iter = self.iter_span.enter();
 87 |                 // We want to avoid carrying around BLDN elements in memory, so
 88 |                 // let's compute everything on the fly with at most (N x D) elements
 89 | 
 90 |                 let delta_i = delta.i((batch, i, ..))?.reshape((1, d_inner))?; // (size 1 D)
 91 |                 let delta_a_i = delta_i.broadcast_mul(&at)?.exp()?; // exp(Δ_i * A) (size N D)
 92 | 
 93 |                 let delta_u_i = delta_u.i((batch, i, ..))?.reshape((1, d_inner))?; // (size 1 D)
 94 |                 let b_i = b.i((batch, i, ..))?.reshape((d_state, 1))?; // size (N 1)
 95 | 
 96 |                 // Performs the following operation:
 97 |                 // (B_i * u_i) x Δ_i simplified discretization
 98 |                 // The (N 1) x (1 D) broadcast mul can also be done as a matmul)
 99 |                 let delta_b_u_i = b_i.matmul(&delta_u_i)?; // (size N D)
100 | 
101 |                 x = ((delta_a_i * x)? + delta_b_u_i)?; // h_t = (ΔA_i * h_(t-1)) + ΔBu_i
102 | 
103 |                 let c_i = c.i((batch, i, ..))?.reshape((1, d_state))?;
104 |                 let y = c_i.matmul(&x)?; // y = C x h_t  { matrix mul (1 N) x (N D) -> (1 D) }
105 | 
106 |                 let u_i = u.i((batch, i, ..))?;
107 |                 let du_i = (&self.d * &u_i)?.reshape((1, d_inner)); // Calculate D * u (size 1 D)
108 |                 ys.push((y + du_i)?)
109 |             }
110 |             ctx.set("ht", x)?; // Save off the resulting SSM to the context
111 |             let ys = Tensor::cat(&ys, 0)?; // [(1,D) x L] -> (L, D)
112 |             bys.push(ys)
113 |         }
114 |         Tensor::stack(&bys, 0)
115 |     }
116 | }
117 | 
118 | impl Module for SSM {
119 |     fn forward(&self, xs: &Tensor) -> Result<Tensor> {
120 |         let _enter = self.span.enter();
121 | 
122 |         let (n, _d_in) = self.at.dims2()?;
123 | 
124 |         let x_dbl = xs.apply(&self.x_proj)?;
125 | 
126 |         let delta = x_dbl.narrow(D::Minus1, 0, self.dt_rank)?;
127 |         let b = x_dbl.narrow(D::Minus1, self.dt_rank, n)?;
128 |         let c = x_dbl.narrow(D::Minus1, self.dt_rank + n, n)?;
129 | 
130 |         let delta = delta.contiguous()?.apply(&self.dt_proj)?;
131 |         // softplus without threshold
132 |         let delta = (delta.exp()? + 1.)?.log()?;
133 |         let ss = self.selective_scan(xs, &delta, &b, &c)?;
134 |         Ok(ss)
135 |     }
136 | }
137 | 


--------------------------------------------------------------------------------
/mamba-ssm/src/ops/elemwise.rs:
--------------------------------------------------------------------------------
  1 | use candle::{bail, CpuStorage, CustomOp2, Layout, Shape};
  2 | use rayon::prelude::*;
  3 | 
  4 | pub enum BinaryOp {
  5 |     Add,
  6 |     Mul,
  7 | }
  8 | 
  9 | type C = CpuStorage;
 10 | 
 11 | impl CustomOp2 for BinaryOp {
 12 |     fn name(&self) -> &'static str {
 13 |         match self {
 14 |             Self::Add => "custom-add",
 15 |             Self::Mul => "custom-mul",
 16 |         }
 17 |     }
 18 | 
 19 |     /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
 20 |     /// offsets etc so the associated layout should be used to access it.
 21 |     fn cpu_fwd(
 22 |         &self,
 23 |         s1: &CpuStorage,
 24 |         l1: &Layout,
 25 |         s2: &CpuStorage,
 26 |         l2: &Layout,
 27 |     ) -> candle::Result<(CpuStorage, Shape)> {
 28 |         if l1.shape() != l2.shape() {
 29 |             bail!("lhs and rhs mismatched shapes");
 30 |         }
 31 |         match (s1, s2) {
 32 |             (C::BF16(_v1), C::BF16(_v2)) => bail!("TODO: Implement BF16 binary op"),
 33 |             (C::F16(_v1), C::F16(_v2)) => bail!("TODO: Implement F16 binary op"),
 34 |             (C::F32(v1), C::F32(v2)) => Ok((C::F32(self.f(v1, l1, v2, l2)?), l1.shape().clone())),
 35 |             _ => bail!("Op not implemented"),
 36 |         }
 37 |     }
 38 | }
 39 | 
 40 | impl BinaryOp {
 41 |     fn f(
 42 |         &self,
 43 |         lhs: &[f32],
 44 |         lhs_l: &Layout,
 45 |         rhs: &[f32],
 46 |         rhs_l: &Layout,
 47 |     ) -> candle::Result<Vec<f32>> {
 48 |         match (lhs_l.contiguous_offsets(), rhs_l.contiguous_offsets()) {
 49 |             (Some((o_l1, o_l2)), Some((o_r1, o_r2))) => {
 50 |                 let n = o_l2 - o_l1;
 51 |                 let mut v: Vec<f32> = Vec::with_capacity(o_l2 - o_l1);
 52 |                 unsafe {
 53 |                     v.set_len(n);
 54 |                 }
 55 |                 match self {
 56 |                     BinaryOp::Add => {
 57 |                         binary_add_f32(&lhs[o_l1..o_l2], &rhs[o_r1..o_r2], v.as_mut_slice())
 58 |                     }
 59 |                     BinaryOp::Mul => {
 60 |                         binary_mul_f32(&lhs[o_l1..o_l2], &rhs[o_r1..o_r2], v.as_mut_slice())
 61 |                     }
 62 |                 };
 63 |                 Ok(v)
 64 |             }
 65 |             _ => bail!("lhs and rhs layouts must be C contiguous"),
 66 |         }
 67 |     }
 68 | }
 69 | 
 70 | fn binary_add_f32(lhs: &[f32], rhs: &[f32], dst: &mut [f32]) {
 71 |     let chunk_size = dst.len() / 8;
 72 |     let operand_chunks = lhs.par_chunks(chunk_size).zip(rhs.par_chunks(chunk_size));
 73 |     dst.par_chunks_mut(chunk_size).zip(operand_chunks).for_each(
 74 |         |(dst_chunk, (lhs_chunk, rhs_chunk))| {
 75 |             let operands = lhs_chunk.iter().zip(rhs_chunk.iter());
 76 |             dst_chunk
 77 |                 .iter_mut()
 78 |                 .zip(operands)
 79 |                 .for_each(|(dst_val, (lhs, rhs))| {
 80 |                     *dst_val = lhs + rhs;
 81 |                 });
 82 |         },
 83 |     );
 84 | }
 85 | 
 86 | fn binary_mul_f32(lhs: &[f32], rhs: &[f32], dst: &mut [f32]) {
 87 |     let chunk_size = dst.len() / 8;
 88 |     let operand_chunks = lhs.par_chunks(chunk_size).zip(rhs.par_chunks(chunk_size));
 89 |     dst.par_chunks_mut(chunk_size).zip(operand_chunks).for_each(
 90 |         |(dst_chunk, (lhs_chunk, rhs_chunk))| {
 91 |             let operands = lhs_chunk.iter().zip(rhs_chunk.iter());
 92 |             dst_chunk
 93 |                 .iter_mut()
 94 |                 .zip(operands)
 95 |                 .for_each(|(dst_val, (lhs, rhs))| {
 96 |                     *dst_val = lhs * rhs;
 97 |                 });
 98 |         },
 99 |     );
100 | }
101 | 
102 | #[cfg(test)]
103 | mod tests {
104 |     use crate::ops::elemwise::{binary_add_f32, binary_mul_f32};
105 | 
106 |     #[test]
107 |     fn test_binary_add_f32() {
108 |         let lhs: Vec<f32> = (0i16..100).map(|i| i.into()).collect();
109 |         let rhs: Vec<f32> = (0i16..100).rev().map(|i| i.into()).collect();
110 |         let expected: Vec<f32> = vec![99.0; 100];
111 |         let mut dst: Vec<f32> = vec![0.0; 100];
112 | 
113 |         binary_add_f32(lhs.as_slice(), rhs.as_slice(), dst.as_mut_slice());
114 | 
115 |         println!("Expected {expected:?}");
116 |         for i in 0..100 {
117 |             assert_eq!(dst[i], expected[i]);
118 |         }
119 |     }
120 | 
121 |     #[test]
122 |     fn test_binary_mul_f32() {
123 |         let lhs: Vec<f32> = (0i16..100).map(|i| i.into()).collect();
124 |         let rhs: Vec<f32> = (0i16..100).rev().map(|i| i.into()).collect();
125 |         let expected: Vec<f32> = (0i16..100)
126 |             .zip((0i16..100).rev())
127 |             .map(|(a, b)| (a * b).into())
128 |             .collect();
129 |         let mut dst: Vec<f32> = vec![0.0; 100];
130 | 
131 |         binary_mul_f32(lhs.as_slice(), rhs.as_slice(), dst.as_mut_slice());
132 | 
133 |         println!("Expected {expected:?}");
134 |         for i in 0..100 {
135 |             assert_eq!(dst[i], expected[i]);
136 |         }
137 |     }
138 | }
139 | 


--------------------------------------------------------------------------------
/mamba-ssm/src/ops/mod.rs:
--------------------------------------------------------------------------------
1 | pub mod elemwise;
2 | 


--------------------------------------------------------------------------------
/mamba-ssm/src/utils.rs:
--------------------------------------------------------------------------------
 1 | /// Utilities copied from candle-examples: https://github.com/huggingface/candle/tree/37c539f2b7dfc8aa67a10b611dc12e5e0428be00/candle-examples/examples
 2 | use candle::Result;
 3 | 
 4 | /// This is a wrapper around a tokenizer to ensure that tokens can be returned to the user in a
 5 | /// streaming way rather than having to wait for the full decoding.
 6 | pub struct TokenOutputStream {
 7 |     tokenizer: tokenizers::Tokenizer,
 8 |     tokens: Vec<u32>,
 9 |     prev_index: usize,
10 |     current_index: usize,
11 | }
12 | 
13 | #[allow(dead_code)]
14 | impl TokenOutputStream {
15 |     pub fn new(tokenizer: tokenizers::Tokenizer) -> Self {
16 |         Self {
17 |             tokenizer,
18 |             tokens: Vec::new(),
19 |             prev_index: 0,
20 |             current_index: 0,
21 |         }
22 |     }
23 | 
24 |     pub fn into_inner(self) -> tokenizers::Tokenizer {
25 |         self.tokenizer
26 |     }
27 | 
28 |     fn decode(&self, tokens: &[u32]) -> Result<String> {
29 |         match self.tokenizer.decode(tokens, true) {
30 |             Ok(str) => Ok(str),
31 |             Err(err) => candle::bail!("cannot decode: {err}"),
32 |         }
33 |     }
34 | 
35 |     // https://github.com/huggingface/text-generation-inference/blob/5ba53d44a18983a4de32d122f4cb46f4a17d9ef6/server/text_generation_server/models/model.py#L68
36 |     pub fn next_token(&mut self, token: u32) -> Result<Option<String>> {
37 |         let prev_text = if self.tokens.is_empty() {
38 |             String::new()
39 |         } else {
40 |             let tokens = &self.tokens[self.prev_index..self.current_index];
41 |             self.decode(tokens)?
42 |         };
43 |         self.tokens.push(token);
44 |         let text = self.decode(&self.tokens[self.prev_index..])?;
45 |         if text.len() > prev_text.len() && text.chars().last().unwrap().is_ascii() {
46 |             let text = text.split_at(prev_text.len());
47 |             self.prev_index = self.current_index;
48 |             self.current_index = self.tokens.len();
49 |             Ok(Some(text.1.to_string()))
50 |         } else {
51 |             Ok(None)
52 |         }
53 |     }
54 | 
55 |     pub fn decode_rest(&self) -> Result<Option<String>> {
56 |         let prev_text = if self.tokens.is_empty() {
57 |             String::new()
58 |         } else {
59 |             let tokens = &self.tokens[self.prev_index..self.current_index];
60 |             self.decode(tokens)?
61 |         };
62 |         let text = self.decode(&self.tokens[self.prev_index..])?;
63 |         if text.len() > prev_text.len() {
64 |             let text = text.split_at(prev_text.len());
65 |             Ok(Some(text.1.to_string()))
66 |         } else {
67 |             Ok(None)
68 |         }
69 |     }
70 | 
71 |     pub fn decode_all(&self) -> Result<String> {
72 |         self.decode(&self.tokens)
73 |     }
74 | 
75 |     pub fn get_token(&self, token_s: &str) -> Option<u32> {
76 |         self.tokenizer.get_vocab(true).get(token_s).copied()
77 |     }
78 | 
79 |     pub fn tokenizer(&self) -> &tokenizers::Tokenizer {
80 |         &self.tokenizer
81 |     }
82 | 
83 |     pub fn clear(&mut self) {
84 |         self.tokens.clear();
85 |         self.prev_index = 0;
86 |         self.current_index = 0;
87 |     }
88 | }
89 | 


--------------------------------------------------------------------------------
/models/README.txt:
--------------------------------------------------------------------------------
1 | Move safetensors models, config.json, and tokenizer.json
2 | here.


--------------------------------------------------------------------------------
/prompts/chat.txt:
--------------------------------------------------------------------------------
 1 | <|system|>
 2 | You are a friendly chatbot that answers as helpfully as possible.
 3 | <|user|>
 4 | Hi, how are you?
 5 | <|assistant|>
 6 | Good thanks!
 7 | <|user|>
 8 | Please tell me the tallest mountain on Earth.
 9 | <|assistant|>
10 | Sure. The tallest mountain in Earth is Mount Everest, located in the Himalayas.
11 | <|user|>
12 | Thanks!
13 | <|assistant|>
14 | You're welcome!
15 | <|user|>
16 | What is the meaning of life, the universe, and everything?
17 | 


--------------------------------------------------------------------------------