Upload DASSForAudioClassification

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,1089 @@
+{
+  "architectures": [
+    "DASSForAudioClassification"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_dass.DASSConfig",
+    "AutoModelForAudioClassification": "modeling_dass.DASSForAudioClassification"
+  },
+  "depths": [
+    2,
+    2,
+    20,
+    2
+  ],
+  "dims": [
+    96,
+    192,
+    384,
+    768
+  ],
+  "drop_path_rate": 0.2,
+  "embed_dim": 96,
+  "id2label": {
+    "0": "Speech",
+    "1": "Male speech, man speaking",
+    "2": "Female speech, woman speaking",
+    "3": "Child speech, kid speaking",
+    "4": "Conversation",
+    "5": "Narration, monologue",
+    "6": "Babbling",
+    "7": "Speech synthesizer",
+    "8": "Shout",
+    "9": "Bellow",
+    "10": "Whoop",
+    "11": "Yell",
+    "12": "Battle cry",
+    "13": "Children shouting",
+    "14": "Screaming",
+    "15": "Whispering",
+    "16": "Laughter",
+    "17": "Baby laughter",
+    "18": "Giggle",
+    "19": "Snicker",
+    "20": "Belly laugh",
+    "21": "Chuckle, chortle",
+    "22": "Crying, sobbing",
+    "23": "Baby cry, infant cry",
+    "24": "Whimper",
+    "25": "Wail, moan",
+    "26": "Sigh",
+    "27": "Singing",
+    "28": "Choir",
+    "29": "Yodeling",
+    "30": "Chant",
+    "31": "Mantra",
+    "32": "Male singing",
+    "33": "Female singing",
+    "34": "Child singing",
+    "35": "Synthetic singing",
+    "36": "Rapping",
+    "37": "Humming",
+    "38": "Groan",
+    "39": "Grunt",
+    "40": "Whistling",
+    "41": "Breathing",
+    "42": "Wheeze",
+    "43": "Snoring",
+    "44": "Gasp",
+    "45": "Pant",
+    "46": "Snort",
+    "47": "Cough",
+    "48": "Throat clearing",
+    "49": "Sneeze",
+    "50": "Sniff",
+    "51": "Run",
+    "52": "Shuffle",
+    "53": "Walk, footsteps",
+    "54": "Chewing, mastication",
+    "55": "Biting",
+    "56": "Gargling",
+    "57": "Stomach rumble",
+    "58": "Burping, eructation",
+    "59": "Hiccup",
+    "60": "Fart",
+    "61": "Hands",
+    "62": "Finger snapping",
+    "63": "Clapping",
+    "64": "Heart sounds, heartbeat",
+    "65": "Heart murmur",
+    "66": "Cheering",
+    "67": "Applause",
+    "68": "Chatter",
+    "69": "Crowd",
+    "70": "Hubbub, speech noise, speech babble",
+    "71": "Children playing",
+    "72": "Animal",
+    "73": "Domestic animals, pets",
+    "74": "Dog",
+    "75": "Bark",
+    "76": "Yip",
+    "77": "Howl",
+    "78": "Bow-wow",
+    "79": "Growling",
+    "80": "Whimper (dog)",
+    "81": "Cat",
+    "82": "Purr",
+    "83": "Meow",
+    "84": "Hiss",
+    "85": "Caterwaul",
+    "86": "Livestock, farm animals, working animals",
+    "87": "Horse",
+    "88": "Clip-clop",
+    "89": "Neigh, whinny",
+    "90": "Cattle, bovinae",
+    "91": "Moo",
+    "92": "Cowbell",
+    "93": "Pig",
+    "94": "Oink",
+    "95": "Goat",
+    "96": "Bleat",
+    "97": "Sheep",
+    "98": "Fowl",
+    "99": "Chicken, rooster",
+    "100": "Cluck",
+    "101": "Crowing, cock-a-doodle-doo",
+    "102": "Turkey",
+    "103": "Gobble",
+    "104": "Duck",
+    "105": "Quack",
+    "106": "Goose",
+    "107": "Honk",
+    "108": "Wild animals",
+    "109": "Roaring cats (lions, tigers)",
+    "110": "Roar",
+    "111": "Bird",
+    "112": "Bird vocalization, bird call, bird song",
+    "113": "Chirp, tweet",
+    "114": "Squawk",
+    "115": "Pigeon, dove",
+    "116": "Coo",
+    "117": "Crow",
+    "118": "Caw",
+    "119": "Owl",
+    "120": "Hoot",
+    "121": "Bird flight, flapping wings",
+    "122": "Canidae, dogs, wolves",
+    "123": "Rodents, rats, mice",
+    "124": "Mouse",
+    "125": "Patter",
+    "126": "Insect",
+    "127": "Cricket",
+    "128": "Mosquito",
+    "129": "Fly, housefly",
+    "130": "Buzz",
+    "131": "Bee, wasp, etc.",
+    "132": "Frog",
+    "133": "Croak",
+    "134": "Snake",
+    "135": "Rattle",
+    "136": "Whale vocalization",
+    "137": "Music",
+    "138": "Musical instrument",
+    "139": "Plucked string instrument",
+    "140": "Guitar",
+    "141": "Electric guitar",
+    "142": "Bass guitar",
+    "143": "Acoustic guitar",
+    "144": "Steel guitar, slide guitar",
+    "145": "Tapping (guitar technique)",
+    "146": "Strum",
+    "147": "Banjo",
+    "148": "Sitar",
+    "149": "Mandolin",
+    "150": "Zither",
+    "151": "Ukulele",
+    "152": "Keyboard (musical)",
+    "153": "Piano",
+    "154": "Electric piano",
+    "155": "Organ",
+    "156": "Electronic organ",
+    "157": "Hammond organ",
+    "158": "Synthesizer",
+    "159": "Sampler",
+    "160": "Harpsichord",
+    "161": "Percussion",
+    "162": "Drum kit",
+    "163": "Drum machine",
+    "164": "Drum",
+    "165": "Snare drum",
+    "166": "Rimshot",
+    "167": "Drum roll",
+    "168": "Bass drum",
+    "169": "Timpani",
+    "170": "Tabla",
+    "171": "Cymbal",
+    "172": "Hi-hat",
+    "173": "Wood block",
+    "174": "Tambourine",
+    "175": "Rattle (instrument)",
+    "176": "Maraca",
+    "177": "Gong",
+    "178": "Tubular bells",
+    "179": "Mallet percussion",
+    "180": "Marimba, xylophone",
+    "181": "Glockenspiel",
+    "182": "Vibraphone",
+    "183": "Steelpan",
+    "184": "Orchestra",
+    "185": "Brass instrument",
+    "186": "French horn",
+    "187": "Trumpet",
+    "188": "Trombone",
+    "189": "Bowed string instrument",
+    "190": "String section",
+    "191": "Violin, fiddle",
+    "192": "Pizzicato",
+    "193": "Cello",
+    "194": "Double bass",
+    "195": "Wind instrument, woodwind instrument",
+    "196": "Flute",
+    "197": "Saxophone",
+    "198": "Clarinet",
+    "199": "Harp",
+    "200": "Bell",
+    "201": "Church bell",
+    "202": "Jingle bell",
+    "203": "Bicycle bell",
+    "204": "Tuning fork",
+    "205": "Chime",
+    "206": "Wind chime",
+    "207": "Change ringing (campanology)",
+    "208": "Harmonica",
+    "209": "Accordion",
+    "210": "Bagpipes",
+    "211": "Didgeridoo",
+    "212": "Shofar",
+    "213": "Theremin",
+    "214": "Singing bowl",
+    "215": "Scratching (performance technique)",
+    "216": "Pop music",
+    "217": "Hip hop music",
+    "218": "Beatboxing",
+    "219": "Rock music",
+    "220": "Heavy metal",
+    "221": "Punk rock",
+    "222": "Grunge",
+    "223": "Progressive rock",
+    "224": "Rock and roll",
+    "225": "Psychedelic rock",
+    "226": "Rhythm and blues",
+    "227": "Soul music",
+    "228": "Reggae",
+    "229": "Country",
+    "230": "Swing music",
+    "231": "Bluegrass",
+    "232": "Funk",
+    "233": "Folk music",
+    "234": "Middle Eastern music",
+    "235": "Jazz",
+    "236": "Disco",
+    "237": "Classical music",
+    "238": "Opera",
+    "239": "Electronic music",
+    "240": "House music",
+    "241": "Techno",
+    "242": "Dubstep",
+    "243": "Drum and bass",
+    "244": "Electronica",
+    "245": "Electronic dance music",
+    "246": "Ambient music",
+    "247": "Trance music",
+    "248": "Music of Latin America",
+    "249": "Salsa music",
+    "250": "Flamenco",
+    "251": "Blues",
+    "252": "Music for children",
+    "253": "New-age music",
+    "254": "Vocal music",
+    "255": "A capella",
+    "256": "Music of Africa",
+    "257": "Afrobeat",
+    "258": "Christian music",
+    "259": "Gospel music",
+    "260": "Music of Asia",
+    "261": "Carnatic music",
+    "262": "Music of Bollywood",
+    "263": "Ska",
+    "264": "Traditional music",
+    "265": "Independent music",
+    "266": "Song",
+    "267": "Background music",
+    "268": "Theme music",
+    "269": "Jingle (music)",
+    "270": "Soundtrack music",
+    "271": "Lullaby",
+    "272": "Video game music",
+    "273": "Christmas music",
+    "274": "Dance music",
+    "275": "Wedding music",
+    "276": "Happy music",
+    "277": "Funny music",
+    "278": "Sad music",
+    "279": "Tender music",
+    "280": "Exciting music",
+    "281": "Angry music",
+    "282": "Scary music",
+    "283": "Wind",
+    "284": "Rustling leaves",
+    "285": "Wind noise (microphone)",
+    "286": "Thunderstorm",
+    "287": "Thunder",
+    "288": "Water",
+    "289": "Rain",
+    "290": "Raindrop",
+    "291": "Rain on surface",
+    "292": "Stream",
+    "293": "Waterfall",
+    "294": "Ocean",
+    "295": "Waves, surf",
+    "296": "Steam",
+    "297": "Gurgling",
+    "298": "Fire",
+    "299": "Crackle",
+    "300": "Vehicle",
+    "301": "Boat, Water vehicle",
+    "302": "Sailboat, sailing ship",
+    "303": "Rowboat, canoe, kayak",
+    "304": "Motorboat, speedboat",
+    "305": "Ship",
+    "306": "Motor vehicle (road)",
+    "307": "Car",
+    "308": "Vehicle horn, car horn, honking",
+    "309": "Toot",
+    "310": "Car alarm",
+    "311": "Power windows, electric windows",
+    "312": "Skidding",
+    "313": "Tire squeal",
+    "314": "Car passing by",
+    "315": "Race car, auto racing",
+    "316": "Truck",
+    "317": "Air brake",
+    "318": "Air horn, truck horn",
+    "319": "Reversing beeps",
+    "320": "Ice cream truck, ice cream van",
+    "321": "Bus",
+    "322": "Emergency vehicle",
+    "323": "Police car (siren)",
+    "324": "Ambulance (siren)",
+    "325": "Fire engine, fire truck (siren)",
+    "326": "Motorcycle",
+    "327": "Traffic noise, roadway noise",
+    "328": "Rail transport",
+    "329": "Train",
+    "330": "Train whistle",
+    "331": "Train horn",
+    "332": "Railroad car, train wagon",
+    "333": "Train wheels squealing",
+    "334": "Subway, metro, underground",
+    "335": "Aircraft",
+    "336": "Aircraft engine",
+    "337": "Jet engine",
+    "338": "Propeller, airscrew",
+    "339": "Helicopter",
+    "340": "Fixed-wing aircraft, airplane",
+    "341": "Bicycle",
+    "342": "Skateboard",
+    "343": "Engine",
+    "344": "Light engine (high frequency)",
+    "345": "Dental drill, dentist's drill",
+    "346": "Lawn mower",
+    "347": "Chainsaw",
+    "348": "Medium engine (mid frequency)",
+    "349": "Heavy engine (low frequency)",
+    "350": "Engine knocking",
+    "351": "Engine starting",
+    "352": "Idling",
+    "353": "Accelerating, revving, vroom",
+    "354": "Door",
+    "355": "Doorbell",
+    "356": "Ding-dong",
+    "357": "Sliding door",
+    "358": "Slam",
+    "359": "Knock",
+    "360": "Tap",
+    "361": "Squeak",
+    "362": "Cupboard open or close",
+    "363": "Drawer open or close",
+    "364": "Dishes, pots, and pans",
+    "365": "Cutlery, silverware",
+    "366": "Chopping (food)",
+    "367": "Frying (food)",
+    "368": "Microwave oven",
+    "369": "Blender",
+    "370": "Water tap, faucet",
+    "371": "Sink (filling or washing)",
+    "372": "Bathtub (filling or washing)",
+    "373": "Hair dryer",
+    "374": "Toilet flush",
+    "375": "Toothbrush",
+    "376": "Electric toothbrush",
+    "377": "Vacuum cleaner",
+    "378": "Zipper (clothing)",
+    "379": "Keys jangling",
+    "380": "Coin (dropping)",
+    "381": "Scissors",
+    "382": "Electric shaver, electric razor",
+    "383": "Shuffling cards",
+    "384": "Typing",
+    "385": "Typewriter",
+    "386": "Computer keyboard",
+    "387": "Writing",
+    "388": "Alarm",
+    "389": "Telephone",
+    "390": "Telephone bell ringing",
+    "391": "Ringtone",
+    "392": "Telephone dialing, DTMF",
+    "393": "Dial tone",
+    "394": "Busy signal",
+    "395": "Alarm clock",
+    "396": "Siren",
+    "397": "Civil defense siren",
+    "398": "Buzzer",
+    "399": "Smoke detector, smoke alarm",
+    "400": "Fire alarm",
+    "401": "Foghorn",
+    "402": "Whistle",
+    "403": "Steam whistle",
+    "404": "Mechanisms",
+    "405": "Ratchet, pawl",
+    "406": "Clock",
+    "407": "Tick",
+    "408": "Tick-tock",
+    "409": "Gears",
+    "410": "Pulleys",
+    "411": "Sewing machine",
+    "412": "Mechanical fan",
+    "413": "Air conditioning",
+    "414": "Cash register",
+    "415": "Printer",
+    "416": "Camera",
+    "417": "Single-lens reflex camera",
+    "418": "Tools",
+    "419": "Hammer",
+    "420": "Jackhammer",
+    "421": "Sawing",
+    "422": "Filing (rasp)",
+    "423": "Sanding",
+    "424": "Power tool",
+    "425": "Drill",
+    "426": "Explosion",
+    "427": "Gunshot, gunfire",
+    "428": "Machine gun",
+    "429": "Fusillade",
+    "430": "Artillery fire",
+    "431": "Cap gun",
+    "432": "Fireworks",
+    "433": "Firecracker",
+    "434": "Burst, pop",
+    "435": "Eruption",
+    "436": "Boom",
+    "437": "Wood",
+    "438": "Chop",
+    "439": "Splinter",
+    "440": "Crack",
+    "441": "Glass",
+    "442": "Chink, clink",
+    "443": "Shatter",
+    "444": "Liquid",
+    "445": "Splash, splatter",
+    "446": "Slosh",
+    "447": "Squish",
+    "448": "Drip",
+    "449": "Pour",
+    "450": "Trickle, dribble",
+    "451": "Gush",
+    "452": "Fill (with liquid)",
+    "453": "Spray",
+    "454": "Pump (liquid)",
+    "455": "Stir",
+    "456": "Boiling",
+    "457": "Sonar",
+    "458": "Arrow",
+    "459": "Whoosh, swoosh, swish",
+    "460": "Thump, thud",
+    "461": "Thunk",
+    "462": "Electronic tuner",
+    "463": "Effects unit",
+    "464": "Chorus effect",
+    "465": "Basketball bounce",
+    "466": "Bang",
+    "467": "Slap, smack",
+    "468": "Whack, thwack",
+    "469": "Smash, crash",
+    "470": "Breaking",
+    "471": "Bouncing",
+    "472": "Whip",
+    "473": "Flap",
+    "474": "Scratch",
+    "475": "Scrape",
+    "476": "Rub",
+    "477": "Roll",
+    "478": "Crushing",
+    "479": "Crumpling, crinkling",
+    "480": "Tearing",
+    "481": "Beep, bleep",
+    "482": "Ping",
+    "483": "Ding",
+    "484": "Clang",
+    "485": "Squeal",
+    "486": "Creak",
+    "487": "Rustle",
+    "488": "Whir",
+    "489": "Clatter",
+    "490": "Sizzle",
+    "491": "Clicking",
+    "492": "Clickety-clack",
+    "493": "Rumble",
+    "494": "Plop",
+    "495": "Jingle, tinkle",
+    "496": "Hum",
+    "497": "Zing",
+    "498": "Boing",
+    "499": "Crunch",
+    "500": "Silence",
+    "501": "Sine wave",
+    "502": "Harmonic",
+    "503": "Chirp tone",
+    "504": "Sound effect",
+    "505": "Pulse",
+    "506": "Inside, small room",
+    "507": "Inside, large room or hall",
+    "508": "Inside, public space",
+    "509": "Outside, urban or manmade",
+    "510": "Outside, rural or natural",
+    "511": "Reverberation",
+    "512": "Echo",
+    "513": "Noise",
+    "514": "Environmental noise",
+    "515": "Static",
+    "516": "Mains hum",
+    "517": "Distortion",
+    "518": "Sidetone",
+    "519": "Cacophony",
+    "520": "White noise",
+    "521": "Pink noise",
+    "522": "Throbbing",
+    "523": "Vibration",
+    "524": "Television",
+    "525": "Radio",
+    "526": "Field recording"
+  },
+  "label2id": {
+    "A capella": 255,
+    "Accelerating, revving, vroom": 353,
+    "Accordion": 209,
+    "Acoustic guitar": 143,
+    "Afrobeat": 257,
+    "Air brake": 317,
+    "Air conditioning": 413,
+    "Air horn, truck horn": 318,
+    "Aircraft": 335,
+    "Aircraft engine": 336,
+    "Alarm": 388,
+    "Alarm clock": 395,
+    "Ambient music": 246,
+    "Ambulance (siren)": 324,
+    "Angry music": 281,
+    "Animal": 72,
+    "Applause": 67,
+    "Arrow": 458,
+    "Artillery fire": 430,
+    "Babbling": 6,
+    "Baby cry, infant cry": 23,
+    "Baby laughter": 17,
+    "Background music": 267,
+    "Bagpipes": 210,
+    "Bang": 466,
+    "Banjo": 147,
+    "Bark": 75,
+    "Basketball bounce": 465,
+    "Bass drum": 168,
+    "Bass guitar": 142,
+    "Bathtub (filling or washing)": 372,
+    "Battle cry": 12,
+    "Beatboxing": 218,
+    "Bee, wasp, etc.": 131,
+    "Beep, bleep": 481,
+    "Bell": 200,
+    "Bellow": 9,
+    "Belly laugh": 20,
+    "Bicycle": 341,
+    "Bicycle bell": 203,
+    "Bird": 111,
+    "Bird flight, flapping wings": 121,
+    "Bird vocalization, bird call, bird song": 112,
+    "Biting": 55,
+    "Bleat": 96,
+    "Blender": 369,
+    "Bluegrass": 231,
+    "Blues": 251,
+    "Boat, Water vehicle": 301,
+    "Boiling": 456,
+    "Boing": 498,
+    "Boom": 436,
+    "Bouncing": 471,
+    "Bow-wow": 78,
+    "Bowed string instrument": 189,
+    "Brass instrument": 185,
+    "Breaking": 470,
+    "Breathing": 41,
+    "Burping, eructation": 58,
+    "Burst, pop": 434,
+    "Bus": 321,
+    "Busy signal": 394,
+    "Buzz": 130,
+    "Buzzer": 398,
+    "Cacophony": 519,
+    "Camera": 416,
+    "Canidae, dogs, wolves": 122,
+    "Cap gun": 431,
+    "Car": 307,
+    "Car alarm": 310,
+    "Car passing by": 314,
+    "Carnatic music": 261,
+    "Cash register": 414,
+    "Cat": 81,
+    "Caterwaul": 85,
+    "Cattle, bovinae": 90,
+    "Caw": 118,
+    "Cello": 193,
+    "Chainsaw": 347,
+    "Change ringing (campanology)": 207,
+    "Chant": 30,
+    "Chatter": 68,
+    "Cheering": 66,
+    "Chewing, mastication": 54,
+    "Chicken, rooster": 99,
+    "Child singing": 34,
+    "Child speech, kid speaking": 3,
+    "Children playing": 71,
+    "Children shouting": 13,
+    "Chime": 205,
+    "Chink, clink": 442,
+    "Chirp tone": 503,
+    "Chirp, tweet": 113,
+    "Choir": 28,
+    "Chop": 438,
+    "Chopping (food)": 366,
+    "Chorus effect": 464,
+    "Christian music": 258,
+    "Christmas music": 273,
+    "Chuckle, chortle": 21,
+    "Church bell": 201,
+    "Civil defense siren": 397,
+    "Clang": 484,
+    "Clapping": 63,
+    "Clarinet": 198,
+    "Classical music": 237,
+    "Clatter": 489,
+    "Clickety-clack": 492,
+    "Clicking": 491,
+    "Clip-clop": 88,
+    "Clock": 406,
+    "Cluck": 100,
+    "Coin (dropping)": 380,
+    "Computer keyboard": 386,
+    "Conversation": 4,
+    "Coo": 116,
+    "Cough": 47,
+    "Country": 229,
+    "Cowbell": 92,
+    "Crack": 440,
+    "Crackle": 299,
+    "Creak": 486,
+    "Cricket": 127,
+    "Croak": 133,
+    "Crow": 117,
+    "Crowd": 69,
+    "Crowing, cock-a-doodle-doo": 101,
+    "Crumpling, crinkling": 479,
+    "Crunch": 499,
+    "Crushing": 478,
+    "Crying, sobbing": 22,
+    "Cupboard open or close": 362,
+    "Cutlery, silverware": 365,
+    "Cymbal": 171,
+    "Dance music": 274,
+    "Dental drill, dentist's drill": 345,
+    "Dial tone": 393,
+    "Didgeridoo": 211,
+    "Ding": 483,
+    "Ding-dong": 356,
+    "Disco": 236,
+    "Dishes, pots, and pans": 364,
+    "Distortion": 517,
+    "Dog": 74,
+    "Domestic animals, pets": 73,
+    "Door": 354,
+    "Doorbell": 355,
+    "Double bass": 194,
+    "Drawer open or close": 363,
+    "Drill": 425,
+    "Drip": 448,
+    "Drum": 164,
+    "Drum and bass": 243,
+    "Drum kit": 162,
+    "Drum machine": 163,
+    "Drum roll": 167,
+    "Dubstep": 242,
+    "Duck": 104,
+    "Echo": 512,
+    "Effects unit": 463,
+    "Electric guitar": 141,
+    "Electric piano": 154,
+    "Electric shaver, electric razor": 382,
+    "Electric toothbrush": 376,
+    "Electronic dance music": 245,
+    "Electronic music": 239,
+    "Electronic organ": 156,
+    "Electronic tuner": 462,
+    "Electronica": 244,
+    "Emergency vehicle": 322,
+    "Engine": 343,
+    "Engine knocking": 350,
+    "Engine starting": 351,
+    "Environmental noise": 514,
+    "Eruption": 435,
+    "Exciting music": 280,
+    "Explosion": 426,
+    "Fart": 60,
+    "Female singing": 33,
+    "Female speech, woman speaking": 2,
+    "Field recording": 526,
+    "Filing (rasp)": 422,
+    "Fill (with liquid)": 452,
+    "Finger snapping": 62,
+    "Fire": 298,
+    "Fire alarm": 400,
+    "Fire engine, fire truck (siren)": 325,
+    "Firecracker": 433,
+    "Fireworks": 432,
+    "Fixed-wing aircraft, airplane": 340,
+    "Flamenco": 250,
+    "Flap": 473,
+    "Flute": 196,
+    "Fly, housefly": 129,
+    "Foghorn": 401,
+    "Folk music": 233,
+    "Fowl": 98,
+    "French horn": 186,
+    "Frog": 132,
+    "Frying (food)": 367,
+    "Funk": 232,
+    "Funny music": 277,
+    "Fusillade": 429,
+    "Gargling": 56,
+    "Gasp": 44,
+    "Gears": 409,
+    "Giggle": 18,
+    "Glass": 441,
+    "Glockenspiel": 181,
+    "Goat": 95,
+    "Gobble": 103,
+    "Gong": 177,
+    "Goose": 106,
+    "Gospel music": 259,
+    "Groan": 38,
+    "Growling": 79,
+    "Grunge": 222,
+    "Grunt": 39,
+    "Guitar": 140,
+    "Gunshot, gunfire": 427,
+    "Gurgling": 297,
+    "Gush": 451,
+    "Hair dryer": 373,
+    "Hammer": 419,
+    "Hammond organ": 157,
+    "Hands": 61,
+    "Happy music": 276,
+    "Harmonic": 502,
+    "Harmonica": 208,
+    "Harp": 199,
+    "Harpsichord": 160,
+    "Heart murmur": 65,
+    "Heart sounds, heartbeat": 64,
+    "Heavy engine (low frequency)": 349,
+    "Heavy metal": 220,
+    "Helicopter": 339,
+    "Hi-hat": 172,
+    "Hiccup": 59,
+    "Hip hop music": 217,
+    "Hiss": 84,
+    "Honk": 107,
+    "Hoot": 120,
+    "Horse": 87,
+    "House music": 240,
+    "Howl": 77,
+    "Hubbub, speech noise, speech babble": 70,
+    "Hum": 496,
+    "Humming": 37,
+    "Ice cream truck, ice cream van": 320,
+    "Idling": 352,
+    "Independent music": 265,
+    "Insect": 126,
+    "Inside, large room or hall": 507,
+    "Inside, public space": 508,
+    "Inside, small room": 506,
+    "Jackhammer": 420,
+    "Jazz": 235,
+    "Jet engine": 337,
+    "Jingle (music)": 269,
+    "Jingle bell": 202,
+    "Jingle, tinkle": 495,
+    "Keyboard (musical)": 152,
+    "Keys jangling": 379,
+    "Knock": 359,
+    "Laughter": 16,
+    "Lawn mower": 346,
+    "Light engine (high frequency)": 344,
+    "Liquid": 444,
+    "Livestock, farm animals, working animals": 86,
+    "Lullaby": 271,
+    "Machine gun": 428,
+    "Mains hum": 516,
+    "Male singing": 32,
+    "Male speech, man speaking": 1,
+    "Mallet percussion": 179,
+    "Mandolin": 149,
+    "Mantra": 31,
+    "Maraca": 176,
+    "Marimba, xylophone": 180,
+    "Mechanical fan": 412,
+    "Mechanisms": 404,
+    "Medium engine (mid frequency)": 348,
+    "Meow": 83,
+    "Microwave oven": 368,
+    "Middle Eastern music": 234,
+    "Moo": 91,
+    "Mosquito": 128,
+    "Motor vehicle (road)": 306,
+    "Motorboat, speedboat": 304,
+    "Motorcycle": 326,
+    "Mouse": 124,
+    "Music": 137,
+    "Music for children": 252,
+    "Music of Africa": 256,
+    "Music of Asia": 260,
+    "Music of Bollywood": 262,
+    "Music of Latin America": 248,
+    "Musical instrument": 138,
+    "Narration, monologue": 5,
+    "Neigh, whinny": 89,
+    "New-age music": 253,
+    "Noise": 513,
+    "Ocean": 294,
+    "Oink": 94,
+    "Opera": 238,
+    "Orchestra": 184,
+    "Organ": 155,
+    "Outside, rural or natural": 510,
+    "Outside, urban or manmade": 509,
+    "Owl": 119,
+    "Pant": 45,
+    "Patter": 125,
+    "Percussion": 161,
+    "Piano": 153,
+    "Pig": 93,
+    "Pigeon, dove": 115,
+    "Ping": 482,
+    "Pink noise": 521,
+    "Pizzicato": 192,
+    "Plop": 494,
+    "Plucked string instrument": 139,
+    "Police car (siren)": 323,
+    "Pop music": 216,
+    "Pour": 449,
+    "Power tool": 424,
+    "Power windows, electric windows": 311,
+    "Printer": 415,
+    "Progressive rock": 223,
+    "Propeller, airscrew": 338,
+    "Psychedelic rock": 225,
+    "Pulleys": 410,
+    "Pulse": 505,
+    "Pump (liquid)": 454,
+    "Punk rock": 221,
+    "Purr": 82,
+    "Quack": 105,
+    "Race car, auto racing": 315,
+    "Radio": 525,
+    "Rail transport": 328,
+    "Railroad car, train wagon": 332,
+    "Rain": 289,
+    "Rain on surface": 291,
+    "Raindrop": 290,
+    "Rapping": 36,
+    "Ratchet, pawl": 405,
+    "Rattle": 135,
+    "Rattle (instrument)": 175,
+    "Reggae": 228,
+    "Reverberation": 511,
+    "Reversing beeps": 319,
+    "Rhythm and blues": 226,
+    "Rimshot": 166,
+    "Ringtone": 391,
+    "Roar": 110,
+    "Roaring cats (lions, tigers)": 109,
+    "Rock and roll": 224,
+    "Rock music": 219,
+    "Rodents, rats, mice": 123,
+    "Roll": 477,
+    "Rowboat, canoe, kayak": 303,
+    "Rub": 476,
+    "Rumble": 493,
+    "Run": 51,
+    "Rustle": 487,
+    "Rustling leaves": 284,
+    "Sad music": 278,
+    "Sailboat, sailing ship": 302,
+    "Salsa music": 249,
+    "Sampler": 159,
+    "Sanding": 423,
+    "Sawing": 421,
+    "Saxophone": 197,
+    "Scary music": 282,
+    "Scissors": 381,
+    "Scrape": 475,
+    "Scratch": 474,
+    "Scratching (performance technique)": 215,
+    "Screaming": 14,
+    "Sewing machine": 411,
+    "Shatter": 443,
+    "Sheep": 97,
+    "Ship": 305,
+    "Shofar": 212,
+    "Shout": 8,
+    "Shuffle": 52,
+    "Shuffling cards": 383,
+    "Sidetone": 518,
+    "Sigh": 26,
+    "Silence": 500,
+    "Sine wave": 501,
+    "Singing": 27,
+    "Singing bowl": 214,
+    "Single-lens reflex camera": 417,
+    "Sink (filling or washing)": 371,
+    "Siren": 396,
+    "Sitar": 148,
+    "Sizzle": 490,
+    "Ska": 263,
+    "Skateboard": 342,
+    "Skidding": 312,
+    "Slam": 358,
+    "Slap, smack": 467,
+    "Sliding door": 357,
+    "Slosh": 446,
+    "Smash, crash": 469,
+    "Smoke detector, smoke alarm": 399,
+    "Snake": 134,
+    "Snare drum": 165,
+    "Sneeze": 49,
+    "Snicker": 19,
+    "Sniff": 50,
+    "Snoring": 43,
+    "Snort": 46,
+    "Sonar": 457,
+    "Song": 266,
+    "Soul music": 227,
+    "Sound effect": 504,
+    "Soundtrack music": 270,
+    "Speech": 0,
+    "Speech synthesizer": 7,
+    "Splash, splatter": 445,
+    "Splinter": 439,
+    "Spray": 453,
+    "Squawk": 114,
+    "Squeak": 361,
+    "Squeal": 485,
+    "Squish": 447,
+    "Static": 515,
+    "Steam": 296,
+    "Steam whistle": 403,
+    "Steel guitar, slide guitar": 144,
+    "Steelpan": 183,
+    "Stir": 455,
+    "Stomach rumble": 57,
+    "Stream": 292,
+    "String section": 190,
+    "Strum": 146,
+    "Subway, metro, underground": 334,
+    "Swing music": 230,
+    "Synthesizer": 158,
+    "Synthetic singing": 35,
+    "Tabla": 170,
+    "Tambourine": 174,
+    "Tap": 360,
+    "Tapping (guitar technique)": 145,
+    "Tearing": 480,
+    "Techno": 241,
+    "Telephone": 389,
+    "Telephone bell ringing": 390,
+    "Telephone dialing, DTMF": 392,
+    "Television": 524,
+    "Tender music": 279,
+    "Theme music": 268,
+    "Theremin": 213,
+    "Throat clearing": 48,
+    "Throbbing": 522,
+    "Thump, thud": 460,
+    "Thunder": 287,
+    "Thunderstorm": 286,
+    "Thunk": 461,
+    "Tick": 407,
+    "Tick-tock": 408,
+    "Timpani": 169,
+    "Tire squeal": 313,
+    "Toilet flush": 374,
+    "Tools": 418,
+    "Toot": 309,
+    "Toothbrush": 375,
+    "Traditional music": 264,
+    "Traffic noise, roadway noise": 327,
+    "Train": 329,
+    "Train horn": 331,
+    "Train wheels squealing": 333,
+    "Train whistle": 330,
+    "Trance music": 247,
+    "Trickle, dribble": 450,
+    "Trombone": 188,
+    "Truck": 316,
+    "Trumpet": 187,
+    "Tubular bells": 178,
+    "Tuning fork": 204,
+    "Turkey": 102,
+    "Typewriter": 385,
+    "Typing": 384,
+    "Ukulele": 151,
+    "Vacuum cleaner": 377,
+    "Vehicle": 300,
+    "Vehicle horn, car horn, honking": 308,
+    "Vibraphone": 182,
+    "Vibration": 523,
+    "Video game music": 272,
+    "Violin, fiddle": 191,
+    "Vocal music": 254,
+    "Wail, moan": 25,
+    "Walk, footsteps": 53,
+    "Water": 288,
+    "Water tap, faucet": 370,
+    "Waterfall": 293,
+    "Waves, surf": 295,
+    "Wedding music": 275,
+    "Whack, thwack": 468,
+    "Whale vocalization": 136,
+    "Wheeze": 42,
+    "Whimper": 24,
+    "Whimper (dog)": 80,
+    "Whip": 472,
+    "Whir": 488,
+    "Whispering": 15,
+    "Whistle": 402,
+    "Whistling": 40,
+    "White noise": 520,
+    "Whoop": 10,
+    "Whoosh, swoosh, swish": 459,
+    "Wild animals": 108,
+    "Wind": 283,
+    "Wind chime": 206,
+    "Wind instrument, woodwind instrument": 195,
+    "Wind noise (microphone)": 285,
+    "Wood": 437,
+    "Wood block": 173,
+    "Writing": 387,
+    "Yell": 11,
+    "Yip": 76,
+    "Yodeling": 29,
+    "Zing": 497,
+    "Zipper (clothing)": 378,
+    "Zither": 150
+  },
+  "max_length": 1024,
+  "model_type": "dass",
+  "num_classes": 527,
+  "num_mel_bins": 128,
+  "patch_size": 4,
+  "torch_dtype": "float32",
+  "transformers_version": "4.50.0.dev0",
+  "use_checkpoint": false
+}

configuration_dass.py

@@ -0,0 +1,91 @@
+# coding=utf-8
+"""Distilled Audio State-Space Model (DASS) configuration"""
+
+from typing import Any, Dict
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+class DASSConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`DASSModel`]. It is used to instantiate a DASS
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to that of the
+    [DASS-small](https://github.com/Saurabhbhati/DASS/) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        patch_size (`int`, *optional*, defaults to 4):
+            The size (resolution) of each patch.
+        embed_dim (`int`, *optional*, defaults to 96):
+            Dimensionality of patch embedding.
+        depths (`list(int)`, *optional*, defaults to `[2, 2, 8, 2]`):
+            Depth of each layer in the DASS encoder.
+        dims (`list(int)`, *optional*, defaults to `[96, 192, 384, 768]`):
+            Dimensionality of each layer in the DASS encoder.
+        drop_path_rate (`float`, *optional*, defaults to 0.2):
+            Stochastic depth rate.
+        num_classes (`int`, *optional*, defaults to 527):
+            Number of classes for classification.
+        max_length (`int`, *optional*, defaults to 1024):
+            Temporal dimension of the spectrograms.
+        num_mel_bins (`int`, *optional*, defaults to 128):
+            Frequency dimension of the spectrograms (number of Mel-frequency bins).
+        use_checkpoint (`bool`, *optional*, defaults to `False`):
+            Whether to use checkpointing to save memory.
+
+    Example:
+
+    ```python
+    >>> from transformers import DASSConfig, DASSModel
+
+    >>> # Initializing a DASS small style configuration
+    >>> configuration = DASSConfig()
+
+    >>> # Initializing a model (with random weights) from the DASS small style configuration
+    >>> model = DASSModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "dass"
+
+    def __init__(
+        self,
+        patch_size: int = 4,
+        embed_dim: int = 96,
+        depths: list = [2, 2, 8, 2],
+        dims: list =[96, 192, 384, 768],
+        drop_path_rate: float = 0.2,
+        num_classes: int = 527,
+        max_length: int = 1024,
+        num_mel_bins: int = 128,
+        use_checkpoint: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        self.depths = depths
+        self.dims = dims
+        self.drop_path_rate = drop_path_rate
+        self.num_classes = num_classes
+        self.max_length = max_length
+        self.num_mel_bins = num_mel_bins
+        self.use_checkpoint = use_checkpoint
+
+    # Overwritten from the parent class: DASS is not compatible with `generate`, but has a config parameter sharing the
+    # same name (`max_length`). Sharing the same name triggers checks regarding the config -> generation_config
+    # generative parameters deprecation cycle, overwriting this function prevents this from happening.
+    def _get_non_default_generation_parameters(self) -> Dict[str, Any]:
+        return {}
+
+
+__all__ = ["DASSConfig"]

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd691a24b6075e62531fe8ee242d5eb196f045ab678bc468715d648a217bd8e2
+size 194645532

modeling_dass.py

@@ -0,0 +1,1228 @@
+# coding=utf-8
+# VMamba backbone is from https://github.com/MzeroMiko/VMamba/blob/main/vmamba.py
+# DASSLayer, DASSModel, DASSForAudioClassification are implemnted based on VMamba and AST
+# 
+"""Distilled Audio State-Space Model (DASS) model"""
+
+import math
+import torch 
+import warnings
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from functools import partial
+from typing import Optional, Callable, Any, Union
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+from transformers.utils import logging
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_dass import DASSConfig
+
+logger = logging.get_logger(__name__)
+
+# General docstring
+_CONFIG_FOR_DOC = "DASSConfig"
+
+WITH_TRITON = True
+# WITH_TRITON = False
+try:
+    import triton
+    import triton.language as tl
+except:
+    WITH_TRITON = False
+    warnings.warn("Triton not installed, fall back to pytorch implements.")
+
+# to make sure cached_property can be loaded for triton
+if WITH_TRITON:
+    try:
+        from functools import cached_property
+    except:
+        warnings.warn("if you are using py37, add this line to functools.py: "
+            "cached_property = lambda func: property(lru_cache()(func))")
+
+# torch implementation ========================================
+def cross_scan_fwd(x: torch.Tensor, in_channel_first=True, out_channel_first=True, scans=0):
+    if in_channel_first:
+        B, C, H, W = x.shape
+        if scans == 0:
+            y = x.new_empty((B, 4, C, H * W))
+            y[:, 0, :, :] = x.flatten(2, 3)
+            y[:, 1, :, :] = x.transpose(dim0=2, dim1=3).flatten(2, 3)
+            y[:, 2:4, :, :] = torch.flip(y[:, 0:2, :, :], dims=[-1])
+        elif scans == 1:
+            y = x.view(B, 1, C, H * W).repeat(1, 4, 1, 1)
+        elif scans == 2:
+            y = x.view(B, 1, C, H * W).repeat(1, 2, 1, 1)
+            y = torch.cat([y, y.flip(dims=[-1])], dim=1)
+        elif scans == 3:
+            y = x.new_empty((B, 4, C, H * W))
+            y[:, 0, :, :] = x.flatten(2, 3)
+            y[:, 1, :, :] = torch.rot90(x, 1, dims=(2, 3)).flatten(2, 3)
+            y[:, 2, :, :] = torch.rot90(x, 2, dims=(2, 3)).flatten(2, 3)
+            y[:, 3, :, :] = torch.rot90(x, 3, dims=(2, 3)).flatten(2, 3)
+    else:
+        B, H, W, C = x.shape
+        if scans == 0:
+            y = x.new_empty((B, H * W, 4, C))
+            y[:, :, 0, :] = x.flatten(1, 2)
+            y[:, :, 1, :] = x.transpose(dim0=1, dim1=2).flatten(1, 2)
+            y[:, :, 2:4, :] = torch.flip(y[:, :, 0:2, :], dims=[1])
+        elif scans == 1:
+            y = x.view(B, H * W, 1, C).repeat(1, 1, 4, 1)
+        elif scans == 2:
+            y = x.view(B, H * W, 1, C).repeat(1, 1, 2, 1)
+            y = torch.cat([y, y.flip(dims=[1])], dim=2)
+        elif scans == 3:
+            y = x.new_empty((B, H * W, 4, C))
+            y[:, :, 0, :] = x.flatten(1, 2)
+            y[:, :, 1, :] = torch.rot90(x, 1, dims=(1, 2)).flatten(1, 2)
+            y[:, :, 2, :] = torch.rot90(x, 2, dims=(1, 2)).flatten(1, 2)
+            y[:, :, 3, :] = torch.rot90(x, 3, dims=(1, 2)).flatten(1, 2)
+
+    if in_channel_first and (not out_channel_first):
+        y = y.permute(0, 3, 1, 2).contiguous()
+    elif (not in_channel_first) and out_channel_first:
+        y = y.permute(0, 2, 3, 1).contiguous()
+
+    return y
+
+
+def cross_merge_fwd(y: torch.Tensor, in_channel_first=True, out_channel_first=True, scans=0):
+    if out_channel_first:
+        B, K, D, H, W = y.shape
+        y = y.view(B, K, D, -1)
+        if scans == 0:
+            y = y[:, 0:2] + y[:, 2:4].flip(dims=[-1]).view(B, 2, D, -1)
+            y = y[:, 0] + y[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).contiguous().view(B, D, -1)
+        elif scans == 1:
+            y = y.sum(1)
+        elif scans == 2:
+            y = y[:, 0:2] + y[:, 2:4].flip(dims=[-1]).view(B, 2, D, -1)
+            y = y.sum(1)
+        elif scans == 3:
+            oy = y[:, 0, :, :].contiguous().view(B, D, -1)
+            oy = oy + torch.rot90(y.view(B, K, D, W, H)[:, 1, :, :, :], -1, dims=(2, 3)).flatten(2, 3)
+            oy = oy + torch.rot90(y.view(B, K, D, H, W)[:, 2, :, :, :], -2, dims=(2, 3)).flatten(2, 3)
+            oy = oy + torch.rot90(y.view(B, K, D, W, H)[:, 3, :, :, :], -3, dims=(2, 3)).flatten(2, 3)
+            y = oy
+    else:
+        B, H, W, K, D = y.shape
+        y = y.view(B, -1, K, D)
+        if scans == 0:
+            y = y[:, :, 0:2] + y[:, :, 2:4].flip(dims=[1]).view(B, -1, 2, D)
+            y = y[:, :, 0] + y[:, :, 1].view(B, W, H, -1).transpose(dim0=1, dim1=2).contiguous().view(B, -1, D)        
+        elif scans == 1:
+            y = y.sum(2)
+        elif scans == 2:
+            y = y[:, :, 0:2] + y[:, :, 2:4].flip(dims=[1]).view(B, -1, 2, D)
+            y = y.sum(2)
+        elif scans == 3:
+            oy = y[:, :, 0, :].contiguous().view(B, -1, D)
+            oy = oy + torch.rot90(y.view(B, W, H, K, D)[:, :, :, 1, :], -1, dims=(1, 2)).flatten(1, 2)
+            oy = oy + torch.rot90(y.view(B, H, W, K, D)[:, :, :, 2, :], -2, dims=(1, 2)).flatten(1, 2)
+            oy = oy + torch.rot90(y.view(B, W, H, K, D)[:, :, :, 3, :], -3, dims=(1, 2)).flatten(1, 2)
+            y = oy
+            
+    if in_channel_first and (not out_channel_first):
+        y = y.permute(0, 2, 1).contiguous()
+    elif (not in_channel_first) and out_channel_first:
+        y = y.permute(0, 2, 1).contiguous()
+    
+    return y
+
+
+def cross_scan1b1_fwd(x: torch.Tensor, in_channel_first=True, out_channel_first=True, scans=0):
+    if in_channel_first:
+        B, _, C, H, W = x.shape
+        if scans == 0:
+            y = torch.stack([
+                x[:, 0].flatten(2, 3),
+                x[:, 1].transpose(dim0=2, dim1=3).flatten(2, 3),
+                torch.flip(x[:, 2].flatten(2, 3), dims=[-1]),
+                torch.flip(x[:, 3].transpose(dim0=2, dim1=3).flatten(2, 3), dims=[-1]),
+            ], dim=1)
+        elif scans == 1:
+            y = x.flatten(2, 3)
+        elif scans == 2:
+            y = torch.stack([
+                x[:, 0].flatten(2, 3),
+                x[:, 1].flatten(2, 3),
+                torch.flip(x[:, 2].flatten(2, 3), dims=[-1]),
+                torch.flip(x[:, 3].flatten(2, 3), dims=[-1]),
+            ], dim=1)
+        elif scans == 3:
+            y = torch.stack([
+                x[:, 0, :, :, :].flatten(2, 3),
+                torch.rot90(x[:, 1, :, :, :], 1, dims=(2, 3)).flatten(2, 3),
+                torch.rot90(x[:, 2, :, :, :], 2, dims=(2, 3)).flatten(2, 3),
+                torch.rot90(x[:, 3, :, :, :], 3, dims=(2, 3)).flatten(2, 3),
+            ], dim=1)
+
+    else:
+        B, H, W, _, C = x.shape
+        if scans == 0:
+            y = torch.stack([
+                x[:, :, :, 0].flatten(1, 2),
+                x[:, :, :, 1].transpose(dim0=1, dim1=2).flatten(1, 2),
+                torch.flip(x[:, :, :, 2].flatten(1, 2), dims=[1]),
+                torch.flip(x[:, :, :, 3].transpose(dim0=1, dim1=2).flatten(1, 2), dims=[1]),
+            ], dim=2)
+        elif scans == 1:
+            y = x.flatten(1, 2)
+        elif scans == 2:
+            y = torch.stack([
+                x[:, 0].flatten(1, 2),
+                x[:, 1].flatten(1, 2),
+                torch.flip(x[:, 2].flatten(1, 2), dims=[-1]),
+                torch.flip(x[:, 3].flatten(1, 2), dims=[-1]),
+            ], dim=2)
+        elif scans == 3:
+            y = torch.stack([
+                x[:, :, :, 0, :].flatten(1, 2),
+                torch.rot90(x[:, :, :, 1, :], 1, dims=(1, 2)).flatten(1, 2),
+                torch.rot90(x[:, :, :, 2, :], 2, dims=(1, 2)).flatten(1, 2),
+                torch.rot90(x[:, :, :, 3, :], 3, dims=(1, 2)).flatten(1, 2),
+            ], dim=1)
+
+    if in_channel_first and (not out_channel_first):
+        y = y.permute(0, 3, 1, 2).contiguous()
+    elif (not in_channel_first) and out_channel_first:
+        y = y.permute(0, 2, 3, 1).contiguous()
+
+    return y
+
+
+def cross_merge1b1_fwd(y: torch.Tensor, in_channel_first=True, out_channel_first=True, scans=0):
+    if out_channel_first:
+        B, K, D, H, W = y.shape
+        y = y.view(B, K, D, -1)
+        if scans == 0:
+            y = torch.stack([
+                y[:, 0],
+                y[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).flatten(2, 3),
+                torch.flip(y[:, 2], dims=[-1]),
+                torch.flip(y[:, 3].view(B, -1, W, H).transpose(dim0=2, dim1=3).flatten(2, 3), dims=[-1]),
+            ], dim=1)
+        elif scans == 1:
+            y = y
+        elif scans == 2:
+            y = torch.stack([
+                y[:, 0],
+                y[:, 1],
+                torch.flip(y[:, 2], dims=[-1]),
+                torch.flip(y[:, 3], dims=[-1]),
+            ], dim=1)
+        elif scans == 3:
+            y = torch.stack([
+                y[:, 0, :, :].contiguous().view(B, D, -1),
+                torch.rot90(y.view(B, K, D, W, H)[:, 1, :, :, :], -1, dims=(2, 3)).flatten(2, 3),
+                torch.rot90(y.view(B, K, D, H, W)[:, 2, :, :, :], -2, dims=(2, 3)).flatten(2, 3),
+                torch.rot90(y.view(B, K, D, W, H)[:, 3, :, :, :], -3, dims=(2, 3)).flatten(2, 3),
+            ], dim=1)
+    else:
+        B, H, W, K, D = y.shape
+        y = y.view(B, -1, K, D)
+        if scans == 0:
+            y = torch.stack([
+                y[:, :, 0],
+                y[:, :, 1].view(B, W, H, -1).transpose(dim0=1, dim1=2).flatten(1, 2),
+                torch.flip(y[:, :, 2], dims=[1]),
+                torch.flip(y[:, :, 3].view(B, W, H, -1).transpose(dim0=1, dim1=2).flatten(1, 2), dims=[1]),
+            ], dim=2)
+        elif scans == 1:
+            y = y
+        elif scans == 2:
+            y = torch.stack([
+                y[:, :, 0],
+                y[:, :, 1],
+                torch.flip(y[:, :, 2], dims=[1]),
+                torch.flip(y[:, :, 3], dims=[1]),
+            ], dim=2)
+        elif scans == 3:
+            y = torch.stack([
+                y[:, :, 0, :].contiguous().view(B, -1, D),
+                torch.rot90(y.view(B, W, H, K, D)[:, :, :, 1, :], -1, dims=(1, 2)).flatten(1, 2),
+                torch.rot90(y.view(B, H, W, K, D)[:, :, :, 2, :], -2, dims=(1, 2)).flatten(1, 2),
+                torch.rot90(y.view(B, W, H, K, D)[:, :, :, 3, :], -3, dims=(1, 2)).flatten(1, 2),
+            ], dim=2)
+
+    if out_channel_first and (not in_channel_first):
+        y = y.permute(0, 3, 1, 2).contiguous()
+    elif (not out_channel_first) and in_channel_first:
+        y = y.permute(0, 2, 3, 1).contiguous()
+
+    return y
+
+
+class CrossScanF(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=0):
+        # x: (B, C, H, W) | (B, H, W, C) | (B, 4, C, H, W) | (B, H, W, 4, C)
+        # y: (B, 4, C, H * W) | (B, H * W, 4, C)
+        ctx.in_channel_first = in_channel_first
+        ctx.out_channel_first = out_channel_first
+        ctx.one_by_one = one_by_one
+        ctx.scans = scans
+
+        if one_by_one:
+            B, K, C, H, W = x.shape
+            if not in_channel_first:
+                B, H, W, K, C = x.shape
+        else:
+            B, C, H, W = x.shape
+            if not in_channel_first:
+                B, H, W, C = x.shape
+        ctx.shape = (B, C, H, W)
+
+        _fn = cross_scan1b1_fwd if one_by_one else cross_scan_fwd
+        y = _fn(x, in_channel_first, out_channel_first, scans)
+
+        return y
+    
+    @staticmethod
+    def backward(ctx, ys: torch.Tensor):
+        # out: (b, k, d, l)
+        in_channel_first = ctx.in_channel_first
+        out_channel_first = ctx.out_channel_first
+        one_by_one = ctx.one_by_one
+        scans = ctx.scans
+        B, C, H, W = ctx.shape
+
+        ys = ys.view(B, -1, C, H, W) if out_channel_first else ys.view(B, H, W, -1, C)
+        _fn = cross_merge1b1_fwd if one_by_one else cross_merge_fwd
+        y = _fn(ys, in_channel_first, out_channel_first, scans)
+        
+        if one_by_one:
+            y = y.view(B, 4, -1, H, W) if in_channel_first else y.view(B, H, W, 4, -1)
+        else:
+            y = y.view(B, -1, H, W) if in_channel_first else y.view(B, H, W, -1)
+
+        return y, None, None, None, None
+
+
+class CrossMergeF(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, ys: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=0):
+        # x: (B, C, H, W) | (B, H, W, C) | (B, 4, C, H, W) | (B, H, W, 4, C)
+        # y: (B, 4, C, H * W) | (B, H * W, 4, C)
+        ctx.in_channel_first = in_channel_first
+        ctx.out_channel_first = out_channel_first
+        ctx.one_by_one = one_by_one
+        ctx.scans = scans
+
+        B, K, C, H, W = ys.shape
+        if not out_channel_first:
+            B, H, W, K, C = ys.shape
+        ctx.shape = (B, C, H, W)
+        
+        _fn = cross_merge1b1_fwd if one_by_one else cross_merge_fwd
+        y = _fn(ys, in_channel_first, out_channel_first, scans)
+
+        return y
+    
+    @staticmethod
+    def backward(ctx, x: torch.Tensor):
+        # B, D, L = x.shape
+        # out: (b, k, d, h, w)
+        in_channel_first = ctx.in_channel_first
+        out_channel_first = ctx.out_channel_first
+        one_by_one = ctx.one_by_one
+        scans = ctx.scans
+        B, C, H, W = ctx.shape
+    
+        if not one_by_one:
+            if in_channel_first:
+                x = x.view(B, C, H, W)
+            else:
+                x = x.view(B, H, W, C)
+        else:
+            if in_channel_first:
+                x = x.view(B, 4, C, H, W)
+            else:
+                x = x.view(B, H, W, 4, C)   
+                     
+        _fn = cross_scan1b1_fwd if one_by_one else cross_scan_fwd
+        x = _fn(x, in_channel_first, out_channel_first, scans)
+        x = x.view(B, 4, C, H, W) if out_channel_first else x.view(B, H, W, 4, C)
+    
+        return x, None, None, None, None
+
+
+# triton implements ========================================
+
+@triton.jit
+def triton_cross_scan_flex(
+    x: tl.tensor, # (B, C, H, W) | (B, H, W, C) | (B, 4, C, H, W) | (B, H, W, 4, C)
+    y: tl.tensor, # (B, 4, C, H, W) | (B, H, W, 4, C)
+    x_layout: tl.constexpr,
+    y_layout: tl.constexpr,
+    operation: tl.constexpr,
+    onebyone: tl.constexpr,
+    scans: tl.constexpr,
+    BC: tl.constexpr,
+    BH: tl.constexpr,
+    BW: tl.constexpr,
+    DC: tl.constexpr,
+    DH: tl.constexpr,
+    DW: tl.constexpr,
+    NH: tl.constexpr,
+    NW: tl.constexpr,
+):
+    # x_layout = 0
+    # y_layout = 1 # 0 BCHW, 1 BHWC
+    # operation = 0 # 0 scan, 1 merge
+    # onebyone = 0 # 0 false, 1 true
+    # scans = 0 # 0 cross scan, 1 unidirectional, 2 bidirectional
+
+    i_hw, i_c, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_h, i_w = (i_hw // NW), (i_hw % NW)
+    _mask_h = (i_h * BH + tl.arange(0, BH)) < DH
+    _mask_w = (i_w * BW + tl.arange(0, BW)) < DW
+    _mask_hw = _mask_h[:, None] & _mask_w[None, :]
+    _for_C = min(DC - i_c * BC, BC)
+
+    pos_h = (i_h * BH + tl.arange(0, BH)[:, None])
+    pos_w = (i_w * BW + tl.arange(0, BW)[None, :])
+    neg_h = (DH - i_h * BH - 1 - tl.arange(0, BH)[:, None])
+    neg_w = (DW - i_w * BW - 1 - tl.arange(0, BW)[None, :])
+    if scans == 0:
+        # none; trans; flip; trans + flip;
+        HWRoute0 = pos_h * DW + pos_w
+        HWRoute1 = pos_w * DH + pos_h # trans
+        HWRoute2 = neg_h * DW + neg_w # flip
+        HWRoute3 = neg_w * DH + neg_h # trans + flip
+    elif scans == 1:
+        # none; none; none; none;
+        HWRoute0 = pos_h * DW + pos_w
+        HWRoute1 = HWRoute0
+        HWRoute2 = HWRoute0
+        HWRoute3 = HWRoute0
+    elif scans == 2:
+        # none; none; flip; flip;
+        HWRoute0 = pos_h * DW + pos_w
+        HWRoute1 = HWRoute0
+        HWRoute2 = neg_h * DW + neg_w # flip
+        HWRoute3 = HWRoute2 
+    elif scans == 3:
+        # none; rot90; rot180==flip; rot270;
+        HWRoute0 = pos_h * DW + pos_w
+        HWRoute1 = neg_w * DH + pos_h
+        HWRoute2 = neg_h * DW + neg_w
+        HWRoute3 = pos_w * DH + neg_h
+
+    _tmp1 = DC * DH * DW
+
+    y_ptr_base = y + i_b * 4 * _tmp1 + (i_c * BC * DH * DW if y_layout == 0 else i_c * BC)
+    if y_layout == 0:
+        p_y1 = y_ptr_base + HWRoute0
+        p_y2 = y_ptr_base + _tmp1 + HWRoute1
+        p_y3 = y_ptr_base + 2 * _tmp1 + HWRoute2
+        p_y4 = y_ptr_base + 3 * _tmp1 + HWRoute3
+    else:
+        p_y1 = y_ptr_base + HWRoute0 * 4 * DC
+        p_y2 = y_ptr_base + DC + HWRoute1 * 4 * DC
+        p_y3 = y_ptr_base + 2 * DC + HWRoute2 * 4 * DC
+        p_y4 = y_ptr_base + 3 * DC + HWRoute3 * 4 * DC       
+    
+    if onebyone == 0:
+        x_ptr_base = x + i_b * _tmp1 + (i_c * BC * DH * DW if x_layout == 0 else i_c * BC)
+        if x_layout == 0:
+            p_x = x_ptr_base + HWRoute0
+        else:
+            p_x = x_ptr_base + HWRoute0 * DC
+
+        if operation == 0:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                _x = tl.load(p_x + _idx_x, mask=_mask_hw)
+                tl.store(p_y1 + _idx_y, _x, mask=_mask_hw)
+                tl.store(p_y2 + _idx_y, _x, mask=_mask_hw)
+                tl.store(p_y3 + _idx_y, _x, mask=_mask_hw)
+                tl.store(p_y4 + _idx_y, _x, mask=_mask_hw)
+        elif operation == 1:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                _y1 = tl.load(p_y1 + _idx_y, mask=_mask_hw)
+                _y2 = tl.load(p_y2 + _idx_y, mask=_mask_hw)
+                _y3 = tl.load(p_y3 + _idx_y, mask=_mask_hw)
+                _y4 = tl.load(p_y4 + _idx_y, mask=_mask_hw)
+                tl.store(p_x + _idx_x, _y1 + _y2 + _y3 + _y4, mask=_mask_hw)
+
+    else:
+        x_ptr_base = x + i_b * 4 * _tmp1 + (i_c * BC * DH * DW if x_layout == 0 else i_c * BC)
+        if x_layout == 0:
+            p_x1 = x_ptr_base + HWRoute0
+            p_x2 = p_x1 + _tmp1
+            p_x3 = p_x2 + _tmp1
+            p_x4 = p_x3 + _tmp1  
+        else:
+            p_x1 = x_ptr_base + HWRoute0 * 4 * DC
+            p_x2 = p_x1 + DC
+            p_x3 = p_x2 + DC
+            p_x4 = p_x3 + DC        
+    
+        if operation == 0:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                tl.store(p_y1 + _idx_y, tl.load(p_x1 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+                tl.store(p_y2 + _idx_y, tl.load(p_x2 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+                tl.store(p_y3 + _idx_y, tl.load(p_x3 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+                tl.store(p_y4 + _idx_y, tl.load(p_x4 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+        else:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                tl.store(p_x1 + _idx_x, tl.load(p_y1 + _idx_y), mask=_mask_hw)
+                tl.store(p_x2 + _idx_x, tl.load(p_y2 + _idx_y), mask=_mask_hw)
+                tl.store(p_x3 + _idx_x, tl.load(p_y3 + _idx_y), mask=_mask_hw)
+                tl.store(p_x4 + _idx_x, tl.load(p_y4 + _idx_y), mask=_mask_hw)
+
+
+class CrossScanTritonF(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=0):
+        if one_by_one:
+            if in_channel_first:
+                B, _, C, H, W = x.shape
+            else:
+                B, H, W, _, C = x.shape
+        else:
+            if in_channel_first:
+                B, C, H, W = x.shape
+            else:
+                B, H, W, C = x.shape
+        B, C, H, W = int(B), int(C), int(H), int(W)
+        BC, BH, BW = 1, 32, 32
+        NH, NW, NC = triton.cdiv(H, BH), triton.cdiv(W, BW), triton.cdiv(C, BC)
+        
+        ctx.in_channel_first = in_channel_first
+        ctx.out_channel_first = out_channel_first
+        ctx.one_by_one = one_by_one
+        ctx.scans = scans
+        ctx.shape = (B, C, H, W)
+        ctx.triton_shape = (BC, BH, BW, NC, NH, NW)
+
+        y = x.new_empty((B, 4, C, H * W)) if out_channel_first else x.new_empty((B, H * W, 4, C))
+        triton_cross_scan_flex[(NH * NW, NC, B)](
+            x.contiguous(), y, 
+            (0 if in_channel_first else 1), (0 if out_channel_first else 1), 0, (0 if not one_by_one else 1), scans, 
+            BC, BH, BW, C, H, W, NH, NW
+        )
+        return y
+        
+    @staticmethod
+    def backward(ctx, y: torch.Tensor):
+        in_channel_first = ctx.in_channel_first
+        out_channel_first = ctx.out_channel_first
+        one_by_one = ctx.one_by_one
+        scans = ctx.scans
+        B, C, H, W = ctx.shape
+        BC, BH, BW, NC, NH, NW = ctx.triton_shape
+        if one_by_one:
+            x = y.new_empty((B, 4, C, H, W)) if in_channel_first else y.new_empty((B, H, W, 4, C))
+        else:
+            x = y.new_empty((B, C, H, W)) if in_channel_first else y.new_empty((B, H, W, C))
+        
+        triton_cross_scan_flex[(NH * NW, NC, B)](
+            x, y.contiguous(), 
+            (0 if in_channel_first else 1), (0 if out_channel_first else 1), 1, (0 if not one_by_one else 1), scans,
+            BC, BH, BW, C, H, W, NH, NW
+        )
+        return x, None, None, None, None
+
+
+class CrossMergeTritonF(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, y: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=0):
+        if out_channel_first:
+            B, _, C, H, W = y.shape
+        else:
+            B, H, W, _, C = y.shape
+        B, C, H, W = int(B), int(C), int(H), int(W)
+        BC, BH, BW = 1, 32, 32
+        NH, NW, NC = triton.cdiv(H, BH), triton.cdiv(W, BW), triton.cdiv(C, BC)
+        ctx.in_channel_first = in_channel_first
+        ctx.out_channel_first = out_channel_first
+        ctx.one_by_one = one_by_one
+        ctx.scans = scans
+        ctx.shape = (B, C, H, W)
+        ctx.triton_shape = (BC, BH, BW, NC, NH, NW)
+        if one_by_one:
+            x = y.new_empty((B, 4, C, H * W)) if in_channel_first else y.new_empty((B, H * W, 4, C))
+        else:
+            x = y.new_empty((B, C, H * W)) if in_channel_first else y.new_empty((B, H * W, C))
+        triton_cross_scan_flex[(NH * NW, NC, B)](
+            x, y.contiguous(), 
+            (0 if in_channel_first else 1), (0 if out_channel_first else 1), 1, (0 if not one_by_one else 1), scans,
+            BC, BH, BW, C, H, W, NH, NW
+        )
+        return x
+        
+    @staticmethod
+    def backward(ctx, x: torch.Tensor):
+        in_channel_first = ctx.in_channel_first
+        out_channel_first = ctx.out_channel_first
+        one_by_one = ctx.one_by_one
+        scans = ctx.scans
+        B, C, H, W = ctx.shape
+        BC, BH, BW, NC, NH, NW = ctx.triton_shape
+        y = x.new_empty((B, 4, C, H, W)) if out_channel_first else x.new_empty((B, H, W, 4, C))
+        triton_cross_scan_flex[(NH * NW, NC, B)](
+            x.contiguous(), y, 
+            (0 if in_channel_first else 1), (0 if out_channel_first else 1), 0, (0 if not one_by_one else 1), scans,
+            BC, BH, BW, C, H, W, NH, NW
+        )
+        return y, None, None, None, None, None
+
+
+# @torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
+def cross_scan_fn(x: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=0, force_torch=False):
+    # x: (B, C, H, W) | (B, H, W, C) | (B, 4, C, H, W) | (B, H, W, 4, C)
+    # y: (B, 4, C, L) | (B, L, 4, C)
+    # scans: 0: cross scan; 1 unidirectional; 2: bidirectional;
+    CSF = CrossScanTritonF if WITH_TRITON and x.is_cuda and (not force_torch) else CrossScanF
+    if x.is_cuda:
+        with torch.cuda.device(x.device):
+            return CSF.apply(x, in_channel_first, out_channel_first, one_by_one, scans)
+    else:
+        return CrossScanF.apply(x, in_channel_first, out_channel_first, one_by_one, scans)
+
+
+# @torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
+def cross_merge_fn(y: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=0, force_torch=False):
+    # y: (B, 4, C, L) | (B, L, 4, C)
+    # x: (B, C, H * W) | (B, H * W, C) | (B, 4, C, H * W) | (B, H * W, 4, C)
+    # scans: 0: cross scan; 1 unidirectional; 2: bidirectional;
+    CMF = CrossMergeTritonF if WITH_TRITON and y.is_cuda and (not force_torch) else CrossMergeF
+    if y.is_cuda:
+        with torch.cuda.device(y.device):
+            return CMF.apply(y, in_channel_first, out_channel_first, one_by_one, scans)
+    else:
+        return CrossMergeF.apply(y, in_channel_first, out_channel_first, one_by_one, scans)
+
+
+##########################################################
+# csms6s.py
+##########################################################
+
+WITH_SELECTIVESCAN_MAMBA = True
+try:
+    import selective_scan_cuda
+except ImportError:
+    WITH_SELECTIVESCAN_MAMBA = False
+
+
+def selective_scan_torch(
+    u: torch.Tensor, # (B, K * C, L)
+    delta: torch.Tensor, # (B, K * C, L)
+    A: torch.Tensor, # (K * C, N)
+    B: torch.Tensor, # (B, K, N, L)
+    C: torch.Tensor, # (B, K, N, L)
+    D: torch.Tensor = None, # (K * C)
+    delta_bias: torch.Tensor = None, # (K * C)
+    delta_softplus=True, 
+    oflex=True, 
+    *args,
+    **kwargs
+):
+    dtype_in = u.dtype
+    Batch, K, N, L = B.shape
+    KCdim = u.shape[1]
+    Cdim = int(KCdim / K)
+    assert u.shape == (Batch, KCdim, L)
+    assert delta.shape == (Batch, KCdim, L)
+    assert A.shape == (KCdim, N)
+    assert C.shape == B.shape
+
+    if delta_bias is not None:
+        delta = delta + delta_bias[..., None]
+    if delta_softplus:
+        delta = torch.nn.functional.softplus(delta)
+            
+    u, delta, A, B, C = u.float(), delta.float(), A.float(), B.float(), C.float()
+    B = B.view(Batch, K, 1, N, L).repeat(1, 1, Cdim, 1, 1).view(Batch, KCdim, N, L)
+    C = C.view(Batch, K, 1, N, L).repeat(1, 1, Cdim, 1, 1).view(Batch, KCdim, N, L)
+    deltaA = torch.exp(torch.einsum('bdl,dn->bdln', delta, A))
+    deltaB_u = torch.einsum('bdl,bdnl,bdl->bdln', delta, B, u)
+    
+    if True:
+        x = A.new_zeros((Batch, KCdim, N))
+        ys = []
+        for i in range(L):
+            x = deltaA[:, :, i, :] * x + deltaB_u[:, :, i, :]
+            y = torch.einsum('bdn,bdn->bd', x, C[:, :, :, i])
+            ys.append(y)
+        y = torch.stack(ys, dim=2) # (B, C, L)
+    
+    out = y if D is None else y + u * D.unsqueeze(-1)
+    return out if oflex else out.to(dtype=dtype_in)
+
+
+class SelectiveScanCuda(torch.autograd.Function):
+    @staticmethod
+    @torch.cuda.amp.custom_fwd
+    def forward(ctx, u, delta, A, B, C, D=None, delta_bias=None, delta_softplus=False, oflex=True, backend=None):
+        ctx.delta_softplus = delta_softplus
+        # backend = "oflex" if WITH_SELECTIVESCAN_OFLEX and (backend is None) else backend
+        # backend = "core" if WITH_SELECTIVESCAN_CORE and (backend is None) else backend
+        backend = "mamba" if WITH_SELECTIVESCAN_MAMBA and (backend is None) else backend
+        ctx.backend = backend
+        if backend == "oflex":
+            out, x, *rest = selective_scan_cuda_oflex.fwd(u, delta, A, B, C, D, delta_bias, delta_softplus, 1, oflex)
+        elif backend == "mamba":
+            out, x, *rest = selective_scan_cuda.fwd(u, delta, A, B, C, D, None, delta_bias, delta_softplus)
+        ctx.save_for_backward(u, delta, A, B, C, D, delta_bias, x)
+        return out
+    
+    @staticmethod
+    @torch.cuda.amp.custom_bwd
+    def backward(ctx, dout, *args):
+        u, delta, A, B, C, D, delta_bias, x = ctx.saved_tensors
+        backend = ctx.backend
+        if dout.stride(-1) != 1:
+            dout = dout.contiguous()
+        if backend == "oflex":
+            du, ddelta, dA, dB, dC, dD, ddelta_bias, *rest = selective_scan_cuda_oflex.bwd(
+                u, delta, A, B, C, D, delta_bias, dout, x, ctx.delta_softplus, 1
+            )
+        elif backend == "mamba":
+            du, ddelta, dA, dB, dC, dD, ddelta_bias, *rest = selective_scan_cuda.bwd(
+                u, delta, A, B, C, D, None, delta_bias, dout, x, None, None, ctx.delta_softplus,
+                False
+            )
+        return du, ddelta, dA, dB, dC, dD, ddelta_bias, None, None, None
+
+
+def selective_scan_fn(
+    u: torch.Tensor, # (B, K * C, L)
+    delta: torch.Tensor, # (B, K * C, L)
+    A: torch.Tensor, # (K * C, N)
+    B: torch.Tensor, # (B, K, N, L)
+    C: torch.Tensor, # (B, K, N, L)
+    D: torch.Tensor = None, # (K * C)
+    delta_bias: torch.Tensor = None, # (K * C)
+    delta_softplus=True, 
+    oflex=True,
+    backend=None,
+):
+    fn = selective_scan_torch if backend == "torch" or (not WITH_SELECTIVESCAN_MAMBA) else SelectiveScanCuda.apply
+    return fn(u, delta, A, B, C, D, delta_bias, delta_softplus, oflex, backend)
+
+##########################################################
+############## HuggingFace modeling file #################
+##########################################################
+
+def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+    def extra_repr(self):
+        return f'drop_prob={round(self.drop_prob,3):0.3f}'
+
+class DASSLinear2d(nn.Linear):
+    def __init__(self, *args, groups=1, **kwargs):
+        nn.Linear.__init__(self, *args, **kwargs)
+        self.groups = groups
+    
+    def forward(self, x: torch.Tensor):
+        if len(x.shape) == 4:
+            return F.conv2d(x, self.weight[:, :, None, None], self.bias, groups=self.groups)
+        elif len(x.shape) == 3:
+            return F.conv1d(x, self.weight[:, :, None], self.bias, groups=self.groups)
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        self_state_dict = self.state_dict()
+        load_state_dict_keys = list(state_dict.keys())
+        if prefix + "weight" in load_state_dict_keys:
+            state_dict[prefix + "weight"] = state_dict[prefix + "weight"].view_as(self_state_dict["weight"])
+        return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+
+
+class DASSLayerNorm2d(nn.LayerNorm):
+    def __init__(self, *args, **kwargs):
+        nn.LayerNorm.__init__(self, *args, **kwargs)
+
+    def forward(self, x: torch.Tensor):
+        x = x.permute(0, 2, 3, 1)
+        x = nn.LayerNorm.forward(self, x)
+        x = x.permute(0, 3, 1, 2)
+        return x
+
+
+class DASSPatchEmbeddings(nn.Module):
+    """
+    This class turns `input_values` into the initial `hidden_states` (patch embeddings) of shape `(batch_size,
+    seq_length, hidden_size)` to be consumed by a State-space model.
+    """
+
+    def __init__(self, patch_size=4,embed_dim=96):
+        super().__init__()
+
+        stride = patch_size // 2
+        kernel_size = stride + 1
+        padding = 1
+
+        self.projection = nn.Sequential(
+            nn.Conv2d(1, embed_dim // 2, kernel_size=kernel_size, stride=stride, padding=padding),
+            DASSLayerNorm2d(embed_dim // 2),
+            nn.GELU(),
+            nn.Conv2d(embed_dim // 2, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding),
+            DASSLayerNorm2d(embed_dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.unsqueeze(1)
+        x = x.transpose(2, 3)
+        x = self.projection(x)
+        return x
+
+
+class DASSDowsample(nn.Module):
+    """
+    This class downsamples the input tensor using a convolutional layer followed by a layer normalization.
+    """
+    def __init__(self, dim, out_dim, use_norm=True):
+        super().__init__()
+        self.down = nn.Conv2d(dim, out_dim, kernel_size=3, stride=2, padding=1)
+        self.norm = DASSLayerNorm2d(out_dim) if use_norm else nn.Identity()
+
+    def forward(self, x):
+        x = self.down(x)
+        x = self.norm(x)
+        return x
+
+
+class DASSMlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = DASSLinear2d(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = DASSLinear2d(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class SS2D(nn.Module):
+    def __init__(
+        self,
+        # basic dims ===========
+        d_model=96,
+        d_state=16,
+        ssm_ratio=2.0,
+        dt_rank="auto",
+        act_layer=nn.SiLU,
+        # dwconv ===============
+        d_conv=3,
+        conv_bias=True,
+        # ======================
+        dropout=0.0,
+        bias=False,
+        # dt init ==============
+        dt_min=0.001,
+        dt_max=0.1,
+        dt_init="random",
+        dt_scale=1.0,
+        dt_init_floor=1e-4,
+        # forward_type="v05_noz" is always used
+        # ======================
+        **kwargs,
+    ):
+        super().__init__()
+        self.k_group = 4
+        self.d_model = int(d_model)
+        self.d_state = int(d_state)
+        self.d_inner = int(ssm_ratio * d_model)
+        self.dt_rank = int(math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank)
+        self.forward_core = partial(self.forward_corev2, force_fp32=False, no_einsum=True)
+        self.with_dconv = d_conv > 1
+
+        # In projection
+        self.in_proj = DASSLinear2d(self.d_model, self.d_inner, bias=bias)
+        self.act: nn.Module = act_layer()
+
+        # Convolution
+        if self.with_dconv:
+            self.conv2d = nn.Conv2d(
+                in_channels=self.d_inner,
+                out_channels=self.d_inner,
+                groups=self.d_inner,
+                bias=conv_bias,
+                kernel_size=d_conv,
+                padding=(d_conv - 1) // 2,
+            )
+
+        # x_proj and dt_proj
+        self.x_proj = DASSLinear2d(self.d_inner, self.k_group * (self.dt_rank + self.d_state * 2), groups=self.k_group, bias=False)
+        self.dt_projs = DASSLinear2d(self.dt_rank, self.k_group * self.d_inner, groups=self.k_group, bias=False)
+
+        # out projection
+        self.out_proj = DASSLinear2d(self.d_inner, self.d_model, bias=bias)
+        self.dropout = nn.Dropout(dropout) if dropout > 0. else nn.Identity()
+
+        # Initialization
+        self.A_logs, self.Ds, self.dt_projs_weight, self.dt_projs_bias = self.init_dt_A_D(
+            self.d_state, self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor, k_group=self.k_group,
+        )
+        self.dt_projs.weight.data = self.dt_projs_weight.data.view(self.dt_projs.weight.shape)
+        # self.dt_projs.bias.data = self.dt_projs_bias.data.view(self.dt_projs.bias.shape)
+        del self.dt_projs_weight
+        # del self.dt_projs_bias
+        # Define out_norm directly with "LN2D"
+        self.out_norm = DASSLayerNorm2d(self.d_inner)
+
+    @staticmethod
+    def dt_init(dt_rank, d_inner, dt_scale=1.0, dt_init="random", dt_min=0.001, dt_max=0.1, dt_init_floor=1e-4):
+        dt_proj = nn.Linear(dt_rank, d_inner, bias=True)
+
+        dt_init_std = dt_rank**-0.5 * dt_scale
+        if dt_init == "constant":
+            nn.init.constant_(dt_proj.weight, dt_init_std)
+        elif dt_init == "random":
+            nn.init.uniform_(dt_proj.weight, -dt_init_std, dt_init_std)
+        else:
+            raise NotImplementedError
+
+        dt = torch.exp(
+            torch.rand(d_inner) * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min)
+        ).clamp(min=dt_init_floor)
+
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        with torch.no_grad():
+            dt_proj.bias.copy_(inv_dt)
+        
+        return dt_proj
+
+    @staticmethod
+    def A_log_init(d_state, d_inner, copies=-1, device=None, merge=True):
+        A = torch.arange(1, d_state + 1, dtype=torch.float32, device=device).view(1, -1).repeat(d_inner, 1).contiguous()
+        A_log = torch.log(A)
+        if copies > 0:
+            A_log = A_log[None].repeat(copies, 1, 1).contiguous()
+            if merge:
+                A_log = A_log.flatten(0, 1)
+        A_log = nn.Parameter(A_log)
+        A_log._no_weight_decay = True
+        return A_log
+
+    @staticmethod
+    def D_init(d_inner, copies=-1, device=None, merge=True):
+        D = torch.ones(d_inner, device=device)
+        if copies > 0:
+            D = D[None].repeat(copies, 1).contiguous()
+            if merge:
+                D = D.flatten(0, 1)
+        D = nn.Parameter(D)
+        D._no_weight_decay = True
+        return D
+
+    @classmethod
+    def init_dt_A_D(cls, d_state, dt_rank, d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor, k_group=4):
+        dt_projs = [
+            cls.dt_init(dt_rank, d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor)
+            for _ in range(k_group)
+        ]
+        dt_projs_weight = nn.Parameter(torch.stack([t.weight for t in dt_projs], dim=0))
+        dt_projs_bias = nn.Parameter(torch.stack([t.bias for t in dt_projs], dim=0))
+        del dt_projs
+            
+        A_logs = cls.A_log_init(d_state, d_inner, copies=k_group, merge=True)
+        Ds = cls.D_init(d_inner, copies=k_group, merge=True)
+        return A_logs, Ds, dt_projs_weight, dt_projs_bias
+
+    def forward_corev2(
+        self,
+        x: torch.Tensor,
+        force_fp32=False,
+        no_einsum=True,
+    ):
+        B, D, H, W = x.shape
+        N = self.d_state
+        L = H * W
+
+        xs = cross_scan_fn(x, in_channel_first=True, out_channel_first=True)
+        x_dbl = self.x_proj(xs.view(B, -1, L))
+        dts, Bs, Cs = torch.split(x_dbl.view(B, self.k_group, -1, L), [self.dt_rank, N, N], dim=2)
+        dts = dts.contiguous().view(B, -1, L)
+        dts = self.dt_projs(dts)
+
+        xs = xs.view(B, -1, L)
+        dts = dts.contiguous().view(B, -1, L)
+        As = -self.A_logs.to(torch.float32).exp()
+        Ds = self.Ds.to(torch.float32)
+        Bs = Bs.contiguous().view(B, self.k_group, N, L)
+        Cs = Cs.contiguous().view(B, self.k_group, N, L)
+        delta_bias = self.dt_projs_bias.view(-1).to(torch.float32)
+        
+        ys = selective_scan_fn(
+            xs, dts, As, Bs, Cs, Ds, delta_bias, delta_softplus=True, backend="mamba"
+        ).view(B, self.k_group, -1, H, W)
+        
+        y = cross_merge_fn(ys, in_channel_first=True, out_channel_first=True)
+        y = y.view(B, -1, H, W)
+        y = self.out_norm(y)
+        return y.to(x.dtype)
+
+    def forward(self, x: torch.Tensor):
+        x = self.in_proj(x)
+        x = self.conv2d(x)
+            
+        x = self.act(x)
+        y = self.forward_core(x)
+        
+        out = self.dropout(self.out_proj(y))
+        return out
+
+
+class VSSBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int = 0,
+        drop_path: float = 0,
+        ssm_d_state: int = 1,
+        ssm_ratio=1.0,
+        ssm_dt_rank: Any = "auto",
+        ssm_act_layer=nn.SiLU,
+        ssm_conv: int = 3,
+        ssm_conv_bias=False,
+        ssm_drop_rate: float = 0,
+        mlp_ratio=4.0,
+        mlp_act_layer=nn.GELU,
+        mlp_drop_rate: float = 0.0,
+        use_checkpoint: bool = False,
+        post_norm: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.ssm_branch = ssm_ratio > 0
+        self.mlp_branch = mlp_ratio > 0
+        self.use_checkpoint = use_checkpoint
+        self.post_norm = post_norm
+
+        if self.ssm_branch:
+            self.norm = DASSLayerNorm2d(hidden_dim)
+            self.op = SS2D(
+                d_model=hidden_dim, 
+                d_state=ssm_d_state, 
+                ssm_ratio=ssm_ratio,
+                dt_rank=ssm_dt_rank,
+                act_layer=ssm_act_layer,
+                d_conv=ssm_conv,
+                conv_bias=ssm_conv_bias,
+                dropout=ssm_drop_rate,
+            )
+        
+        self.drop_path = DropPath(drop_path)
+        
+        if self.mlp_branch:
+            self.norm2 = DASSLayerNorm2d(hidden_dim)
+            mlp_hidden_dim = int(hidden_dim * mlp_ratio)
+            self.mlp = DASSMlp(in_features=hidden_dim, hidden_features=mlp_hidden_dim, act_layer=mlp_act_layer, drop=mlp_drop_rate)
+
+    def _forward(self, input: torch.Tensor):
+        x = input
+        if self.ssm_branch:
+            if self.post_norm:
+                x = x + self.drop_path(self.norm(self.op(x)))
+            else:
+                x = x + self.drop_path(self.op(self.norm(x)))
+        if self.mlp_branch:
+            if self.post_norm:
+                x = x + self.drop_path(self.norm2(self.mlp(x))) 
+            else:
+                x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+    def forward(self, input: torch.Tensor):
+        if self.use_checkpoint:
+            return checkpoint.checkpoint(self._forward, input)
+        else:
+            return self._forward(input)
+
+class DASSLayer(nn.Module):
+
+    def __init__(
+        self,
+        input_dim,
+        depth,
+        drop_path=0.0,
+        norm_layer=DASSLayerNorm2d,
+        downsample=nn.Identity(),
+        use_checkpoint=False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.use_checkpoint = use_checkpoint
+
+        self.blocks = nn.ModuleList()
+        for i in range(depth):
+            self.blocks.append(
+                VSSBlock(hidden_dim=input_dim,
+                    drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                    norm_layer=norm_layer,use_checkpoint=use_checkpoint,**kwargs,
+                )
+            )
+        
+        self.downsample = downsample
+
+    def forward(self, x):
+        for block in self.blocks:
+            x = block(x)
+
+        x = self.downsample(x)
+        return x
+
+class DASSPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and
+    a simple interface for downloading and loading pretrained models.
+    """
+
+    config_class = DASSConfig
+    base_model_prefix = "dass"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]) -> None:
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+            if isinstance(module, nn.Linear) and module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.constant_(module.bias, 0)
+            nn.init.constant_(module.weight, 1.0)
+
+
+class DASSModel(DASSPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+
+        dims = config.dims
+        if isinstance(dims, int):
+            dims = [int(dims * 2**i_layer) for i_layer in range(self.num_layers)]
+
+        self.dims = dims
+        self.patch_embeddings = DASSPatchEmbeddings(patch_size=config.patch_size,
+                                                     embed_dim=dims[0])
+        
+        self.num_layers = len(config.depths)
+        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths))]
+        self.num_features = dims[-1]
+
+        self.layers = nn.ModuleList()
+        for i in range(self.num_layers):
+            layer = DASSLayer(
+                input_dim=self.dims[i],
+                depth=config.depths[i],
+                drop_path=dpr[sum(config.depths[:i]):sum(config.depths[:i+1])],
+                downsample=DASSDowsample(self.dims[i], self.dims[i+1]) if i < self.num_layers - 1 else nn.Identity(),
+                use_checkpoint=config.use_checkpoint,
+            )
+            self.layers.append(layer)
+        
+        self.norm = DASSLayerNorm2d(self.num_features)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+
+    def get_input_embeddings(self) -> DASSPatchEmbeddings:
+        return self.patch_embeddings
+    
+    def forward(self, input_values: torch.Tensor):
+        x = self.patch_embeddings(input_values)
+        for layer in self.layers:
+            x = layer(x)
+        x = self.norm(x)
+        x = self.avgpool(x).flatten(1)
+        return x
+
+
+class DASSForAudioClassification(DASSPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.num_classes = config.num_classes
+        self.dass = DASSModel(config)
+        self.head = nn.Linear(self.dass.num_features, self.num_classes) if self.num_classes > 0 else nn.Identity()
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_values: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+    ):
+
+        outputs = self.dass(
+            input_values,
+        )
+
+        logits = self.head(outputs)
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            if self.config.loss_type == "ce":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "bce":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+
+        if return_dict:
+            output = (logits,) + (outputs,)
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs,
+        )
+
+__all__ = [
+    "DASSModel",
+    "DASSPreTrainedModel",
+    "DASSForAudioClassification",
+]