pulsatrix::SoftmaxModule Class Reference

Softmax over the tensor's last dimension, applied independently to every "row" (every fixed combination of all leading dimensions). More...

#include <softmax_module.hpp>

Inheritance diagram for pulsatrix::SoftmaxModule:

Collaboration diagram for pulsatrix::SoftmaxModule:

Public Member Functions
	SoftmaxModule (DeviceBackend *backend)
	Constructs a softmax module.
Tensor	backward (const Tensor &grad_output) override
	Computes the gradient w.r.t. this module's input.
OpType	op_type () const override
	Activation per charter's closed OpType set (same category as ReluModule).
Tensor	propagate_relevance (const Tensor &relevance_out, const LRPRuleConfig &config) override
	AttnLRP's softmax relevance rule (Achtibat et al. 2024, Eq. 13 – Deep Taylor Decomposition), applied per row: R_in[i] = x[i] * (R_out[i] - s[i] * sum_j(R_out[j])), with x the cached forward input and s the cached forward output.
std::optional< DeviceType >	compute_device () const override
	Where this layer computes, so forward() rejects an input on another device (FND-8).
Public Member Functions inherited from pulsatrix::Module
virtual	~Module ()=default
Tensor	forward (const Tensor &input)
	Runs this module's forward computation.
std::pair< Tensor, NodeId >	forward_traced (const Tensor &input, NodeId input_node, ComputationGraph &graph, Autograd &autograd)
	Runs forward() while also registering a ComputationGraph node (tagged with this module's op_type(), parented to input_node) and wiring an Autograd backward function that reuses this module's own backward() – the opt-in traced/explainable path, per Phase 2 Mission 0.
virtual bool	supports_lrp_rule (LRPRule rule) const
	Whether propagate_relevance() implements rule (no silent fallback: callers such as ExplainerContext::relevance_pass() throw rather than run a module on a rule it does not implement).
virtual std::vector< NamedParamRef >	named_parameters ()
	This module's trainable parameters, each with its hierarchical name – the one place a module declares its parameters (roadmap FND-1).
virtual std::vector< ParamRef >	parameters ()
	This module's trainable parameters and their gradients, for an optimizer to update uniformly across module types.
void	set_requires_grad (bool requires_grad, const std::string &prefix="")
	Freezes (false) or unfreezes (true) parameters by name (roadmap FND-2).
virtual void	set_training (bool training)
	Sets this module's training/eval mode. Defaults to training (matches every mainstream framework's Module default).
bool	is_training () const
	Whether this module is currently in training mode.

Protected Member Functions
Tensor	forward_impl (const Tensor &input) override
	Numerically stable softmax over the last axis, per row (subtract the row max before exponentiating – same convention as CrossEntropyLoss::forward).

Detailed Description

Softmax over the tensor's last dimension, applied independently to every "row" (every fixed combination of all leading dimensions).

Note

Rank-agnostic by design – Phase 3's MultiHeadAttentionModule softmaxes attention scores of shape (N, num_heads, L, L) over the final L axis, so this module must not assume rank 1 the way CrossEntropyLoss's private inline softmax does.

No learnable parameters (parameters() returns empty, matching ReluModule).

Constructor & Destructor Documentation

SoftmaxModule()

pulsatrix::SoftmaxModule::SoftmaxModule ( DeviceBackend *  backend )

explicit

Constructs a softmax module.

Parameters

backend

Backend to allocate through. Not owned; must outlive this module.

Note

Takes only a backend – the softmax axis is always the last one, so there is no further configuration (matches ReluModule's constructor shape).

Member Function Documentation

backward()

Tensor pulsatrix::SoftmaxModule::backward ( const Tensor &  grad_output )

overridevirtual

Computes the gradient w.r.t. this module's input.

Parameters

grad_output

Gradient w.r.t. this module's output. Must match the shape of the most recent forward() call's output.

Returns

Gradient w.r.t. this module's input, per row: grad_in[i] = s[i] * (grad_out[i] - sum_j(s[j] * grad_out[j])), the standard softmax Jacobian-vector product with s the cached forward output.

Note

Must be called after forward() – uses the output cached from that call.

Device-generic: runs on Cpu, Cuda or Hip tensors (GPU-native-kernels Mission 1).

Implements pulsatrix::Module.

compute_device()

std::optional< DeviceType > pulsatrix::SoftmaxModule::compute_device ( ) const

inlineoverridevirtual

Where this layer computes, so forward() rejects an input on another device (FND-8).

Reimplemented from pulsatrix::Module.

forward_impl()

Tensor pulsatrix::SoftmaxModule::forward_impl ( const Tensor &  input )

overrideprotectedvirtual

Numerically stable softmax over the last axis, per row (subtract the row max before exponentiating – same convention as CrossEntropyLoss::forward).

Parameters

input

Input tensor. Must be rank >= 1; any device.

Returns

Softmax probabilities, same shape as input; each row sums to 1.

Implements pulsatrix::Module.

op_type()

OpType pulsatrix::SoftmaxModule::op_type ( ) const

inlineoverridevirtual

Activation per charter's closed OpType set (same category as ReluModule).

Implements pulsatrix::Module.

propagate_relevance()

Tensor pulsatrix::SoftmaxModule::propagate_relevance ( const Tensor &  relevance_out, const LRPRuleConfig &  config  )

overridevirtual

AttnLRP's softmax relevance rule (Achtibat et al. 2024, Eq. 13 – Deep Taylor Decomposition), applied per row: R_in[i] = x[i] * (R_out[i] - s[i] * sum_j(R_out[j])), with x the cached forward input and s the cached forward output.

Parameters

relevance_out	Relevance at this module's output. Must match forward()'s shape.
config	Unused – Eq. 13 takes no epsilon/gamma parameter. Deliberately NOT epsilon-stabilized (see the conservation note below).

Returns

Relevance at this module's input.

Note

This rule does not conserve relevance – sum(R_in) != sum(R_out) in general, unlike every other propagate_relevance in this codebase. Eq. 13 is a first-order Taylor/DTD approximation around a nonzero reference point and the paper's own text describes the residual "hidden bias term". Do NOT add a stabilizer or rescale to force conservation: that would silently diverge from the cited formula. This module is intentionally excluded from tests/lrp_conservation_test.cpp's AllModuleTypeCases(); its conservation gap is measured and reported in tests/softmax_module_test.cpp instead.

See also

cpp_engineering.aDNA's what/context/cpp_tdd/context_tdd_lrp_rule_pattern_taxonomy.md, "Known-Non-Conserving-by-Design Note" – this is the cited exception to rule shape 2 (bilinear split), not a fifth rule shape.

Note

Must be called after forward() – uses both the input and the output cached from that call (the input is needed by the x[i] factor, which is why this module caches both).

Device-generic: runs on Cpu, Cuda or Hip tensors (GPU-native-kernels Mission 3).

Implements pulsatrix::Module.

---

The documentation for this class was generated from the following file:

• include/pulsatrix/softmax_module.hpp