Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
41 lines
1.1 KiB
Rust
41 lines
1.1 KiB
Rust
//! Activation operator code generation.
|
|
|
|
use proc_macro2::TokenStream;
|
|
use quote::quote;
|
|
|
|
use super::sanitize_name;
|
|
use crate::error::Result;
|
|
use crate::ir::Node;
|
|
|
|
/// Generate softmax operation.
|
|
pub fn generate_softmax(node: &Node) -> Result<TokenStream> {
|
|
let input = &node.inputs[0];
|
|
let output = &node.outputs[0];
|
|
|
|
let in_ident = quote::format_ident!("{}", sanitize_name(input));
|
|
let out_ident = quote::format_ident!("{}", sanitize_name(output));
|
|
|
|
// Get axis attribute (default -1)
|
|
let axis = node.get_int("axis").unwrap_or(-1);
|
|
|
|
Ok(quote! {
|
|
let #out_ident = #in_ident.softmax(#axis)?;
|
|
})
|
|
}
|
|
|
|
/// Generate leaky ReLU operation.
|
|
pub fn generate_leaky_relu(node: &Node) -> Result<TokenStream> {
|
|
let input = &node.inputs[0];
|
|
let output = &node.outputs[0];
|
|
|
|
let in_ident = quote::format_ident!("{}", sanitize_name(input));
|
|
let out_ident = quote::format_ident!("{}", sanitize_name(output));
|
|
|
|
// Get alpha attribute (default 0.01)
|
|
let alpha = node.get_float("alpha").unwrap_or(0.01);
|
|
|
|
Ok(quote! {
|
|
let #out_ident = #in_ident.leaky_relu(#alpha)?;
|
|
})
|
|
}
|