Files
rustytorch/crates/production/rtx-onnx-codegen/src/ops/activation.rs
T
osobhandClaude Opus 4.6 02d382d5f6 style: apply rustfmt across all crates and demos
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]>
2026-04-12 07:01:58 -07:00

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)?;
})
}