Livt.ML

src/activation/ZeroIdentityLookup.lvt

namespace Livt.ML.Activation

using Livt.Math.Arithmetic
using Livt.ML.Storage

/**
 * Rounded lookup with a zero lower tail and identity upper tail, e.g. GELU-new.
 */
component ZeroIdentityLookup<LOW: int, HIGH: int, STEP_BITS: int, CENTER: int, TABLE: ILookupTable>
{
	static assert LOW >= -1073741823 && LOW < 0 && HIGH > 0 && HIGH <= 1073741823, "lookup bounds must straddle zero"

	static assert STEP_BITS >= 0 && STEP_BITS <= 20, "invalid lookup step"

	static assert CENTER >= 0 && CENTER <= 1048576, "invalid lookup center"

	// Borrowed approximation table; keep stable throughout Compute.
	table: TABLE

	// Owned arithmetic worker used by serialized operations.
	math: WideArithmetic<31>

	// Invalid-index status from the latest Compute call.
	failed: bool

	/**
	 * Borrows the initialized lookup table and owns a wide rounding worker; serialize Compute
	 * calls.
	 */
	new(table: TABLE)
	{
		this.table = table
		math = new WideArithmetic<31>()
	}

	/**
	 * Returns zero at or below LOW and value at or above HIGH. Interior values select a ties-
	 * even rounded table entry around CENTER; an invalid index returns zero and sets Failed.
	 */
	public fn Compute(value: int) int
	{
		failed = false

		if (value <= LOW) {
			return 0
		}

		if (value >= HIGH) {
			return value
		}

		var index = (math.DivideEven(value as logic[64], (1 << STEP_BITS) as logic[64]) as int) + CENTER

		if (index < 0 || index >= table.Length()) {
			failed = true
			return 0
		}

		return table.Read(index)
	}

	/**
	 * Reports an invalid table index in the latest Compute call. The next Compute clears this
	 * flag before evaluating its input.
	 */
	public fn Failed() bool
	{
		return failed
	}
}