Skip to content

Add tanh as a first-class TensorOps activation primitive #630

Description

@michalharakal

Problem

tanh is one of the canonical activation functions in ML pedagogy and the default in micrograd, the most-tutorial'd autodiff library on the planet. SKaiNET 0.25.0 ships relu, gelu, elu, leakyRelu, silu, sigmoid, softmax, and logSoftmax on TensorOps — but not tanh.

The gap is visible inside this codebase already: skainet-compile/skainet-compile-dag/.../DefaultExecutionTape.kt:910-914 has the comment:

// Since we don't have tanh yet, we can use: tanh(x) = (exp(2x) - 1) / (exp(2x) + 1)
// or just use a simpler approximation if tanh is missing.
// Actually, let's implement a simple tanh using sigmoid if possible.
// tanh(x) = 2 * sigmoid(2x) - 1

— a local lambda named tanh synthesized inside the GELU backward formula because no real tanh op exists. The downstream skainet-notebook repo has a tanh polyfill doing the same 2*sigmoid(2x) - 1 composition for its micrograd-tribute notebooks. Every consumer re-derives the same expression — that's exactly the kind of duplication a core primitive should eliminate.

Proposal

Add tanh symmetrically with sigmoid. The wiring is the existing six-layer pattern (none of it is novel except the math):

  1. TensorOps interface — @Diff @ActivationDsl public fun <T : DType, V> tanh(tensor: Tensor<T, V>): Tensor<T, V>
  2. TanhOperation<T, V> : BaseOperation("tanh", "activation", ...) in TensorOperations.kt
  3. Tensor<T, V>.tanh() extension in TensorExtensions.kt
  4. VoidTensorOps.tanh shape-passthrough
  5. DefaultCpuOps.tanh delegating to kotlin.math.tanh (numerically stable across the full float range)
  6. RecordingTensorOpsDecorator.tanh records TanhOperation to the tape; stableInputName learns TanhOperation -> "input"
  7. Autograd backward: d/dx tanh(x) = 1 - tanh(x)^2, alongside the existing SigmoidOperation backward entry

Full PRD draft (with code shapes per module) is in the downstream repo at docs/upstream/tanh-activation.md.

Release shape

Purely additive — new method on TensorOps, new op class, new extension function. No existing call site changes behaviour. Safe for a 0.25.1 patch release; doesn't require a minor bump.

Downstream consumers waiting on this

  • skainet-notebook micrograd demos — the two existing notebooks polyfill tanh and the polyfill deletes entirely once this lands; their rendered graphs collapse from a mulScalar -> sigmoid -> mulScalar -> subScalar chain to a single tanh block.
  • skainet-notebook issue for the full micrograd/demo.ipynb port (training + decision boundary) — blocked on the polyfill's perf cost (4× op count per tanh call).
  • The GELU backward formula in DefaultExecutionTape.kt:898-919 would collapse from "synthesize a local tanh from sigmoid" to "call ops.tanh(...)" — separate cleanup PR after the primitive lands.

Non-goals

  • Don't approximate. kotlin.math.tanh is fast and IEEE-correct; no reason to ship the 2*sigmoid(2x) - 1 composition as the CPU primitive even though it's the polyfill shape.
  • Don't add Hardtanh/clamped variants in the same PR (separate op, separate use case).
  • Don't change GELU's backward in the same PR — focused cleanup commit after tanh is verified.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions