Module: Toy::LLM::Primitives::DepthScale

Defined in:
lib/toy/llm/primitives/depth_scale.rb

Constant Summary collapse

NAME =
:depth_scale

Class Method Summary collapse

Class Method Details

.apply(sess, x, inv_sqrt_depth) ⇒ Object

Scale a normalised tensor by the depth constant 1/sqrt(ell). x is the block’s RMSNorm output (the sublayer input); the block passes inv_sqrt_depth = 1.0/Math.sqrt(layer) as a Float constant. Returns the depth-scaled handle.



27
28
29
# File 'lib/toy/llm/primitives/depth_scale.rb', line 27

def self.apply(sess, x, inv_sqrt_depth)
  TinyNN.tnn_scale(sess, x, inv_sqrt_depth)
end