diff --git a/AGENTS.md b/AGENTS.md
index 89c94f0..fc70a29 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -758,8 +758,8 @@ Use **case classes** to group parameters. DimWit automatically derives `TensorTr
 
 ```scala
 import dimwit.*
-import dimwit.autodiff.{Autodiff}
-import dimwit.tensortree.{TensorTree, FloatTree}
+import dimwit.autodiff.Autodiff
+import dimwit.tensortree.TensorTree
 
 trait Feature derives Label
 trait Hidden derives Label
@@ -904,7 +904,7 @@ val lossFunc = mse(trainData, trainLabels)
 val gradFunc = Autodiff.grad(lossFunc)
 
 // Create optimizer
-val optimizer = GradientDescent.of(VType[Float32])(learningRate = 0.01f)
+val optimizer = GradientDescent(learningRate = 0.01f)
 
 // Training loop with iterator
 val trained = optimizer.iterate(initModelParams)(gradFunc)
@@ -922,7 +922,7 @@ import dimwit.optimizer.Lion
 import dimwit.Conversions.given // enables implicit conversion from Float to Tensor[V]
 
 // Lion optimizer with momentum
-val lionOptimizer = Lion.of(VType[Float32])(learningRate = 1e-3f, beta1 = 0.9f, beta2 = 0.99f, weightDecay = 0.0f)
+val lionOptimizer = Lion(learningRate = 1e-3f, beta1 = 0.9f, beta2 = 0.99f, weightDecay = 0.0f)
 
 // Training with Lion
 val trainedLion = lionOptimizer.iterate(initModelParams)(gradFunc)
@@ -967,7 +967,7 @@ val initRegressionParams = RegressionParams(initSlope, initIntercept)
 
 // Train
 val regressionGrad = Autodiff.grad(regressionLoss(xData, yData))
-val gdOptimizer = GradientDescent.of(VType[Float32])(learningRate = 0.1f)
+val gdOptimizer = GradientDescent(learningRate = 0.1f)
 
 val finalParams = gdOptimizer.iterate(initRegressionParams)(regressionGrad)
   .take(100)
diff --git a/core/src/main/scala/dimwit/autodiff/Grad.scala b/core/src/main/scala/dimwit/autodiff/Grad.scala
index 27f9964..f4d862c 100644
--- a/core/src/main/scala/dimwit/autodiff/Grad.scala
+++ b/core/src/main/scala/dimwit/autodiff/Grad.scala
@@ -51,14 +51,10 @@ object Grad:
 
     def fromPyTree(pyVal: Jax.PyAny): Grad[T] = Grad(ev.fromPyTree(pyVal))
 
-  // FloatTree witness for gradient math (++, --, scale, etc.)
-  // given [T, V: IsFloating](using FloatTree[T, V]): FloatTree[Grad[T], V] with {}
+  // TreeOf witness for gradient math (++, --, scale, etc.)
+  // given [T, V: IsFloating](using TreeOf[T, V]): TreeOf[Grad[T], V] with {}
 
   // Bridge extension so we can call .asFloats directly on Grad[Params[V]]
-  extension [F[_], V](g: Grad[F[V]])(using
-      tt: TensorTree[F[V]],
-      ft: FloatTree[F[V], V],
-      isF: IsFloating[V]
-  )
+  extension [F[_], V: IsFloating](g: Grad[F[V]])(using TensorTree[F[V]], TreeOf[F[V], V])
     def asFloats[NewV: IsFloating](vtype: VType[NewV])(using m: Mirror.ProductOf[F[NewV]]): Grad[F[NewV]] =
-      Grad(dimwit.FloatTree.ops.asFloats(g.value)(vtype))
+      Grad(TreeOf.ops.asFloats(g.value)(vtype))
diff --git a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala
index 26aff0c..ba038f2 100644
--- a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala
+++ b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala
@@ -4,8 +4,8 @@ import dimwit.*
 import dimwit.Conversions.given
 import dimwit.autodiff.*
 import dimwit.autodiff.Grad
-import dimwit.tensortree.*
-import dimwit.tensortree.FloatTree.ops.*
+import dimwit.tensortree.TreeOf
+import dimwit.tensortree.TreeOf.ops.*
 
 /** Gradient optimizer interface with functional state management.
   *
@@ -26,35 +26,33 @@ import dimwit.tensortree.FloatTree.ops.*
   *        optimizer.update(grads, params, state)
   *    }}}
   */
-trait GradientOptimizer[V: IsFloating, State0[_]]:
+trait GradientOptimizer[State0[_]]:
 
   type State[P] = State0[P]
 
   // Core API
-  def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): State[P]
-  def update[P](gradients: Grad[P], params: P, state: State[P])(using TensorTree[P], FloatTree[P, V]): (P, State[P])
+  def init[P: TensorTree, V](params: P)(using TreeOf[P, V])(using IsFloating[V]): State[P]
+
+  def update[P: TensorTree, V](gradients: Grad[P], params: P, state: State[P])(using TreeOf[P, V])(using IsFloating[V]): (P, State[P])
 
   // Convenience: iterator with fixed gradient function
-  def iterateWithState[P](init: P)(df: P => Grad[P])(using TensorTree[P], FloatTree[P, V]): Iterator[(P, State[P])] =
+  def iterateWithState[P: TensorTree, V](init: P)(df: P => Grad[P])(using TreeOf[P, V])(using IsFloating[V]): Iterator[(P, State[P])] =
     Iterator.iterate((init, this.init(init))): (params, state) =>
       val grads = df(params)
       update(grads, params, state)
 
-  def iterate[P](init: P)(df: P => Grad[P])(using TensorTree[P], FloatTree[P, V]): Iterator[P] =
+  def iterate[P: TensorTree, V](init: P)(df: P => Grad[P])(using TreeOf[P, V])(using IsFloating[V]): Iterator[P] =
     iterateWithState(init)(df).map(_._1)
 
-object GradientDescent:
-
-  def of[V: IsFloating](vtype: VType[V])(learningRate: Tensor0[V]): GradientDescent[V] = new GradientDescent(learningRate)
-
 type GradientDescentState[P] = Unit // empty state
 
-class GradientDescent[V: IsFloating](val learningRate: Tensor0[V]) extends GradientOptimizer[V, GradientDescentState]:
+class GradientDescent(val learningRate: Tensor0[Float32]) extends GradientOptimizer[GradientDescentState]:
 
-  def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): Unit = ()
+  def init[P: TensorTree, V](params: P)(using TreeOf[P, V])(using IsFloating[V]): Unit = ()
 
-  def update[P](gradients: Grad[P], params: P, state: Unit)(using TensorTree[P], FloatTree[P, V]): (P, Unit) =
-    val newParams = params -- gradients.value.scale(learningRate)
+  def update[P: TensorTree, V](gradients: Grad[P], params: P, state: Unit)(using ft: TreeOf[P, V])(using IsFloating[V]): (P, Unit) =
+    val α = learningRate.asFloat(VType[V])
+    val newParams = params -- gradients.value.scale(α)
     (newParams, ())
 
 case class LionState[P](
@@ -62,68 +60,64 @@ case class LionState[P](
     step: Tensor0[Int32]
 )
 
-object Lion:
-
-  def of[V](vtype: VType[V])(using IsFloating[V])(learningRate: Tensor0[V], weightDecay: Tensor0[V] = Tensor0(vtype)(0.0), beta1: Tensor0[V] = Tensor0(vtype)(0.9), beta2: Tensor0[V] = Tensor0(vtype)(0.99)): Lion[V] = new Lion(learningRate, weightDecay, beta1, beta2)
+class Lion(val learningRate: Tensor0[Float32], val weightDecay: Tensor0[Float32] = Tensor0(0.0f), val beta1: Tensor0[Float32] = Tensor0(0.9f), val beta2: Tensor0[Float32] = Tensor0(0.99f)) extends GradientOptimizer[LionState]:
 
-class Lion[V: IsFloating](val learningRate: Tensor0[V], val weightDecay: Tensor0[V] = Tensor0(0.0), val beta1: Tensor0[V] = Tensor0(0.9), val beta2: Tensor0[V] = Tensor0(0.99)) extends GradientOptimizer[V, LionState]:
-
-  def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): LionState[P] =
+  def init[P: TensorTree, V](params: P)(using TreeOf[P, V])(using IsFloating[V]): LionState[P] =
     LionState(params.fillCopy(0f), step = 1)
 
-  def update[P](gradients: Grad[P], params: P, state: LionState[P])(using TensorTree[P], FloatTree[P, V]): (P, LionState[P]) =
+  def update[P: TensorTree, V](gradients: Grad[P], params: P, state: LionState[P])(using TreeOf[P, V])(using IsFloating[V]): (P, LionState[P]) =
+    val α = learningRate.asFloat(VType[V])
+    val β1 = beta1.asFloat(VType[V])
+    val β2 = beta2.asFloat(VType[V])
+    val λ = weightDecay.asFloat(VType[V])
+
     // the direction (1 or -1)
     // is determined by the sign of the momentum + gradient
-    val updateDirection = (state.momentums **! beta1 ++ gradients.value **! (1f - beta1)).sign
+    val updateDirection = (state.momentums **! β1 ++ gradients.value **! (1f - β1)).sign
 
-    val updatedParams = params -- updateDirection.scale(learningRate) -- params.scale(weightDecay)
-    val newMomentums = state.momentums **! beta2 ++ gradients.value **! (1f - beta2)
+    val updatedParams = params -- updateDirection.scale(α) -- params.scale(λ)
+    val newMomentums = state.momentums **! β2 ++ gradients.value **! (1f - β2)
 
     (updatedParams, LionState(newMomentums, state.step + 1))
 
-case class AdamState[P, V](
-    momentums: P, // momentums
-    velocities: P, // velocities
-    b1: Tensor0[V], // decay rate for momentums mᵗ
-    b2: Tensor0[V] // decay rate for velocities vᵗ
+case class AdamState[P](
+    momentums: P,
+    velocities: P,
+    beta1t: Tensor0[Float32], // decay rate for momentums mᵗ, hard-coded precision to make State independent of V, making persisting and restoring easier
+    beta2t: Tensor0[Float32] // decay rate for velocities vᵗ, hard-coded precision to make State independent of V, making persisting and restoring easier
 )
 
-object Adam:
-
-  def of[V](vtype: VType[V])(using IsFloating[V])(learningRate: Tensor0[V], b1: Tensor0[V] = Tensor0(vtype)(0.9), b2: Tensor0[V] = Tensor0(vtype)(0.999), epsilon: Tensor0[V] = Tensor0(vtype)(1e-8)): Adam[V] = new Adam(learningRate, b1, b2, epsilon)
-
 /** Implements the Adam optimization algorithm.
   *
   * @see [[https://arxiv.org/abs/1412.6980 Adam: A Method for Stochastic Optimization]]
   */
-class Adam[V: IsFloating](
-    val learningRate: Tensor0[V],
-    b1: Tensor0[V] = Tensor0(0.9), // decay rate for momentums mᵗ
-    b2: Tensor0[V] = Tensor0(0.999), // decay rate for velocities vᵗ
-    epsilon: Tensor0[V] = Tensor0(1e-8) // small constant to prevent division by zero
-) extends GradientOptimizer[V, [P] =>> AdamState[P, V]]:
-
-  private val β1 = b1
-  private val β2 = b2
-  private val ε = epsilon
-
-  def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): AdamState[P, V] =
+class Adam(
+    val learningRate: Tensor0[Float32],
+    val beta1: Tensor0[Float32] = Tensor0(0.9f), // decay rate for momentums mᵗ
+    val beta2: Tensor0[Float32] = Tensor0(0.999f), // decay rate for velocities vᵗ
+    val epsilon: Tensor0[Float32] = Tensor0(1e-8f) // small constant to prevent division by zero
+) extends GradientOptimizer[AdamState]:
+
+  def init[P: TensorTree, V](params: P)(using TreeOf[P, V])(using IsFloating[V]): AdamState[P] =
     def zeros = params.fillCopy(0f)
-    AdamState(zeros, zeros, b1 = Tensor0(VType[V])(1f), b2 = Tensor0(VType[V])(1f))
+    AdamState(zeros, zeros, beta1t = Tensor0(1f), beta2t = Tensor0(1f))
 
-  def update[P](
+  def update[P: TensorTree, V](
       gradients: Grad[P],
       params: P,
-      state: AdamState[P, V]
-  )(using TensorTree[P], FloatTree[P, V]): (P, AdamState[P, V]) =
+      state: AdamState[P]
+  )(using TreeOf[P, V])(using IsFloating[V]): (P, AdamState[P]) =
+    // rename parameters for internal clarity
+    val α = learningRate.asFloat(VType[V])
+    val β1 = beta1.asFloat(VType[V])
+    val β2 = beta2.asFloat(VType[V])
+    val ε = epsilon.asFloat(VType[V])
+
     // rename state variables to last time step for clarity
     val `mₜ₋₁` = state.momentums
     val `vₜ₋₁` = state.velocities
-    val `β1ₜ₋₁` = state.b1
-    val `β2ₜ₋₁` = state.b2
-
-    // rename parameters for internal clarity
-    val α = learningRate
+    val `β1ₜ₋₁` = state.beta1t.asFloat(VType[V])
+    val `β2ₜ₋₁` = state.beta2t.asFloat(VType[V])
 
     val `θₜ₋₁` = params
 
@@ -137,7 +131,7 @@ class Adam[V: IsFloating](
     val v̂ = vᵗ `//!` (1f - `β2ₜ`)
     val θₜ = `θₜ₋₁` -- (α **! m̂) `//` (v̂.sqrt ++! ε)
 
-    (θₜ, AdamState(mᵗ, vᵗ, β1ₜ, β2ₜ))
+    (θₜ, AdamState(mᵗ, vᵗ, β1ₜ.asFloat32, β2ₜ.asFloat32))
 
 /** Implements the AdamW algorithm (Adam with decoupled weight decay).
   *
@@ -149,52 +143,23 @@ class Adam[V: IsFloating](
   * @param learningRate The step size.
   * @param weightDecayFactor The coefficient for weight decay (lambda).
   */
-class AdamW[V: IsFloating](
-    val adam: Adam[V],
-    val weightDecayFactor: Tensor0[V]
-) extends GradientOptimizer[V, [P] =>> AdamState[P, V]]:
+class AdamW(
+    val adam: Adam,
+    val weightDecayFactor: Tensor0[Float32]
+) extends GradientOptimizer[AdamState]:
 
-  def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): AdamState[P, V] = adam.init(params)
+  def init[P: TensorTree, V](params: P)(using TreeOf[P, V])(using IsFloating[V]): AdamState[P] = adam.init(params)
 
-  def update[P](
+  def update[P: TensorTree, V](
       gradients: Grad[P],
       params: P,
-      state: AdamState[P, V]
-  )(using TensorTree[P], FloatTree[P, V]): (P, AdamState[P, V]) =
-    val α = adam.learningRate
+      state: AdamState[P]
+  )(using TreeOf[P, V])(using IsFloating[V]): (P, AdamState[P]) =
+    val α = adam.learningRate.asFloat(VType[V])
+    val `λ'` = weightDecayFactor.asFloat(VType[V])
+
     val `θₜ₋₁` = params
-    val `λ'` = weightDecayFactor
     val λ = `λ'` * α // Tie weight decay to learning rate
     val decayedParams = `θₜ₋₁` -- λ **! `θₜ₋₁`
     val (θₜ, adamState) = adam.update(gradients, decayedParams, state)
     (θₜ, adamState)
-
-case class LearningRateScheduleState[P, State[_]](
-    step: Tensor0[Int32],
-    optState: State[P]
-)
-type LearningRateScheduleStateFor[State[_]] = [P] =>> LearningRateScheduleState[P, State]
-
-object LearningRateSchedule:
-
-  def of[V: IsFloating, State[_]](
-      vtype: VType[V]
-  )(
-      optF: Tensor0[V] => GradientOptimizer[V, State],
-      schedule: Tensor0[Int32] => Tensor0[V]
-  ): LearningRateSchedule[V, State] =
-    new LearningRateSchedule(optF, schedule)
-
-class LearningRateSchedule[V: IsFloating, State[_]](val optF: Tensor0[V] => GradientOptimizer[V, State], schedule: Tensor0[Int32] => Tensor0[V]) extends GradientOptimizer[V, LearningRateScheduleStateFor[State]]:
-
-  def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): LearningRateScheduleState[P, State] =
-    val step = Tensor0(1)
-    val opt = optF(schedule(step))
-    LearningRateScheduleState(step, opt.init(params))
-
-  def update[P](gradients: Grad[P], params: P, state: LearningRateScheduleState[P, State])(using TensorTree[P], FloatTree[P, V]): (P, LearningRateScheduleState[P, State]) =
-    val step = state.step
-    val optState = state.optState
-    val opt = optF(schedule(step))
-    val (newParams, newOptState) = opt.update(gradients, params, optState)
-    (newParams, LearningRateScheduleState(step + 1, newOptState))
diff --git a/core/src/main/scala/dimwit/package.scala b/core/src/main/scala/dimwit/package.scala
index 56a807f..67191f7 100644
--- a/core/src/main/scala/dimwit/package.scala
+++ b/core/src/main/scala/dimwit/package.scala
@@ -86,7 +86,7 @@ package object dimwit:
   // Export automatic differentiation
   export dimwit.autodiff.{Autodiff, Grad}
   // Export tensor trees
-  export dimwit.tensortree.{TensorTree, TensorTreeIO, TensorTreeFormat, FloatTree}
+  export dimwit.tensortree.{TensorTree, TensorTreeIO, TensorTreeFormat, TreeOf}
   // Export Just-in-Time compilation
   export dimwit.jax.Jit.{jit, jitDonating, jitDonatingUnsafe}
   export dimwit.jax.EagerCleanup.eagerCleanup
diff --git a/core/src/main/scala/dimwit/tensortree/FloatTree.scala b/core/src/main/scala/dimwit/tensortree/TreeOf.scala
similarity index 60%
rename from core/src/main/scala/dimwit/tensortree/FloatTree.scala
rename to core/src/main/scala/dimwit/tensortree/TreeOf.scala
index e68663d..d71b75a 100644
--- a/core/src/main/scala/dimwit/tensortree/FloatTree.scala
+++ b/core/src/main/scala/dimwit/tensortree/TreeOf.scala
@@ -11,51 +11,52 @@ import scala.util.NotGiven
   * The given instances give evidence that the tensors are
   * of type V, constrained by IsFloating.
   */
-trait FloatTree[P, V]
+trait TreeOf[P, V]
 
-object FloatTree:
+object TreeOf:
 
   // 1. Base case for Tensors
-  given [Q <: Tuple, V: IsFloating]: FloatTree[Tensor[Q, V], V] with {}
+  given [Q <: Tuple, V](using TensorTree[Tensor[Q, V]]): TreeOf[Tensor[Q, V], V] with {}
 
   // 2. Inductive base cases for Tuples
   // This allows the compiler to step through the case class fields and lock in V.
-  given emptyTuple[V]: FloatTree[EmptyTuple, V] with {}
+  given emptyTuple[V]: TreeOf[EmptyTuple, V] with {}
 
   given consTuple[H, T <: Tuple, V](using
-      h: FloatTree[H, V],
-      t: FloatTree[T, V]
-  ): FloatTree[H *: T, V] with {}
+      h: TreeOf[H, V],
+      t: TreeOf[T, V]
+  )(using TensorTree[H *: T]): TreeOf[H *: T, V] with {}
 
   // 3. Standard collections
-  given listInstance[A, V](using FloatTree[A, V]): FloatTree[List[A], V] with {}
+  given listInstance[A: TensorTree, V](using TreeOf[A, V]): TreeOf[List[A], V] with {}
 
-  given mapInstance[K, A, V](using FloatTree[A, V]): FloatTree[Map[K, A], V] with {}
+  // given mapInstance[K, A, V](using TreeOf[A, V]): TreeOf[Map[K, A], V] with {}
 
-  // 4. Named tuples, delegating to the FloatTree instance of the underlying value tuple
-  given namedTupleInstance[N <: Tuple, V <: Tuple, Fl](using FloatTree[V, Fl]): FloatTree[NamedTuple[N, V], Fl] with {}
+  // 4. Named tuples, delegating to the TreeOf instance of the underlying value tuple
+  given namedTupleInstance[N <: Tuple, V <: Tuple: TensorTree, Fl](using TreeOf[V, Fl]): TreeOf[NamedTuple[N, V], Fl] with {}
 
-  inline given derived[P <: Product, V](using
+  inline given derived[P <: Product: TensorTree, V](using
       evNotTuple: NotGiven[P <:< Tuple],
       m: Mirror.ProductOf[P],
-      evElems: FloatTree[m.MirroredElemTypes, V]
-  ): FloatTree[P, V] =
-    FloatTreeImpl[P, V]()
+      evElems: TreeOf[m.MirroredElemTypes, V]
+  ): TreeOf[P, V] =
+    new TreeOfImpl[P, V]()
 
-  class FloatTreeImpl[P, V] extends FloatTree[P, V]
+  class TreeOfImpl[P: TensorTree, V]() extends TreeOf[P, V]
+
+  extension [P: TensorTree, V](p: P)(using TreeOf[P, V])
 
-  extension [P, V](p: P)(using tt: TensorTree[P], ft: FloatTree[P, V], isF: IsFloating[V])
     /** Maps a function over the TensorTree, as for a regular tensor tree,
       * but provides knowledge that tensors are of type V
       */
     def map[NewV](f: [T <: Tuple] => Labels[T] ?=> (Tensor[T, V] => Tensor[T, NewV])): P =
-      tt.map(p, [T <: Tuple, V0] => (n: Labels[T]) ?=> (t: Tensor[T, V0]) => f[T](using n)(t.asInstanceOf[Tensor[T, V]]).asInstanceOf[Tensor[T, V0]])
+      TensorTree[P].map(p, [T <: Tuple, V0] => (n: Labels[T]) ?=> (t: Tensor[T, V0]) => f[T](using n)(t.asInstanceOf[Tensor[T, V]]).asInstanceOf[Tensor[T, V0]])
 
     /** Maps a function over the TensorTree along with the structural path,
       * providing knowledge that tensors are of type V
       */
     def mapWithName[NewV](f: [T <: Tuple] => Labels[T] ?=> ((String, Tensor[T, V]) => Tensor[T, NewV]), path: String = ""): P =
-      tt.mapWithName(
+      TensorTree[P].mapWithName(
         p,
         [T <: Tuple, V0] => (n: Labels[T]) ?=> (pth: String, t: Tensor[T, V0]) => f[T](using n)(pth, t.asInstanceOf[Tensor[T, V]]).asInstanceOf[Tensor[T, V0]],
         path
@@ -64,7 +65,7 @@ object FloatTree:
     /** Foreach over the TensorTree, providing knowledge that tensors are of type V
       */
     def foreach(f: [T <: Tuple] => Labels[T] ?=> (Tensor[T, V] => Unit)): Unit =
-      tt.foreach(
+      TensorTree[P].foreach(
         p,
         [T <: Tuple, V0] => (n: Labels[T]) ?=> (t: Tensor[T, V0]) => f[T](using n)(t.asInstanceOf[Tensor[T, V]])
       )
@@ -73,7 +74,7 @@ object FloatTree:
       * providing knowledge that tensors are of type V
       */
     def foreachWithName(f: [T <: Tuple] => Labels[T] ?=> ((String, Tensor[T, V]) => Unit), path: String = ""): Unit =
-      tt.foreachWithName(
+      TensorTree[P].foreachWithName(
         p,
         [T <: Tuple, V0] => (n: Labels[T]) ?=> (pth: String, t: Tensor[T, V0]) => f[T](using n)(pth, t.asInstanceOf[Tensor[T, V]]),
         path
@@ -83,14 +84,14 @@ object FloatTree:
       * but provides knowledge that tensors are of type V
       */
     def zipMap(p2: P, f: [T <: Tuple] => Labels[T] ?=> ((Tensor[T, V], Tensor[T, V]) => Tensor[T, V])): P =
-      tt.zipMap(
+      TensorTree[P].zipMap(
         p,
         p2,
         [T <: Tuple, V0] => (n: Labels[T]) ?=> (t1: Tensor[T, V0], t2: Tensor[T, V0]) => f[T](using n)(t1.asInstanceOf[Tensor[T, V]], t2.asInstanceOf[Tensor[T, V]]).asInstanceOf[Tensor[T, V0]]
       )
 
     def mapLeaves[A](f: [T <: Tuple] => Labels[T] ?=> (Tensor[T, V] => A)): Iterator[A] =
-      tt.mapLeaves(p, [T <: Tuple, V0] => (n: Labels[T]) ?=> (t: Tensor[T, V0]) => f[T](using n)(t.asInstanceOf[Tensor[T, V]]))
+      TensorTree[P].mapLeaves(p, [T <: Tuple, V0] => (n: Labels[T]) ?=> (t: Tensor[T, V0]) => f[T](using n)(t.asInstanceOf[Tensor[T, V]]))
 
   /** Arithmetic and math operations for tensor trees of floating-point types.
     */
@@ -103,21 +104,21 @@ object FloatTree:
 
     // Scalar broadcast extensions (Tensor0 op Tree)
     extension [V: IsFloating](p2: Tensor0[V])
-      def ++![P](p1: P)(using TensorTree[P], FloatTree[P, V]): P = p1.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a +! p2)
-      def --![P](p1: P)(using TensorTree[P], FloatTree[P, V]): P = p1.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a -! p2)
-      def **![P](p1: P)(using TensorTree[P], FloatTree[P, V]): P = p1.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a *! p2)
-      def `//!`[P](p1: P)(using TensorTree[P], FloatTree[P, V]): P = p1.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a /! p2)
+      def ++![P: TensorTree](p1: P)(using TreeOf[P, V]): P = p1.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a +! p2)
+      def --![P: TensorTree](p1: P)(using TreeOf[P, V]): P = p1.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a -! p2)
+      def **![P: TensorTree](p1: P)(using TreeOf[P, V]): P = p1.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a *! p2)
+      def `//!`[P: TensorTree](p1: P)(using TreeOf[P, V]): P = p1.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a /! p2)
 
     // Scalar broadcast extensions (Tensor0 op Tree)
     extension [V: IsFloating](p2: Double)
-      def ++![P](p1: P)(using TensorTree[P], FloatTree[P, V]): P = Tensor0(VType[V])(p2) ++! p1
-      def --![P](p1: P)(using TensorTree[P], FloatTree[P, V]): P = Tensor0(VType[V])(p2) --! p1
-      def **![P](p1: P)(using TensorTree[P], FloatTree[P, V]): P = Tensor0(VType[V])(p2) **! p1
-      def `//!`[P](p1: P)(using TensorTree[P], FloatTree[P, V]): P = Tensor0(VType[V])(p2) `//!` p1
+      def ++![P: TensorTree](p1: P)(using TreeOf[P, V]): P = Tensor0(VType[V])(p2) ++! p1
+      def --![P: TensorTree](p1: P)(using TreeOf[P, V]): P = Tensor0(VType[V])(p2) --! p1
+      def **![P: TensorTree](p1: P)(using TreeOf[P, V]): P = Tensor0(VType[V])(p2) **! p1
+      def `//!`[P: TensorTree](p1: P)(using TreeOf[P, V]): P = Tensor0(VType[V])(p2) `//!` p1
 
     // Tree extensions (Tree op Tree, Tree op Scalar, and math ops)
     // Excluded for bare Tensor[T, V] to avoid conflicts with tensor's own operators
-    extension [P, V](p1: P)(using tt: TensorTree[P], ft: FloatTree[P, V], isF: IsFloating[V], ev: NotGiven[IsFloatingTensor[P, V]])
+    extension [P: TensorTree, V](p1: P)(using TreeOf[P, V])(using IsFloating[V])
       def ++(p2: P): P = p1.zipMap(p2, [T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V], b: Tensor[T, V]) => a + b)
       def ++!(p2: Tensor0[V]): P = p1.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a +! p2)
       def --(p2: P): P = p1.zipMap(p2, [T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V], b: Tensor[T, V]) => a - b)
@@ -127,6 +128,7 @@ object FloatTree:
       def `//`(p2: P): P = p1.zipMap(p2, [T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V], b: Tensor[T, V]) => a / b)
       def `//!`(p2: Tensor0[V]): P = p1.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a /! p2)
 
+    extension [P: TensorTree, V](p1: P)(using TreeOf[P, V], NotGiven[P <:< Tensor[?, ?]])(using IsFloating[V])
       def sqrt: P = p1.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => TensorOps.sqrt(a))
 
       def pow(exponent: Float): P = pow(Tensor0(VType[V])(exponent))
@@ -136,30 +138,7 @@ object FloatTree:
 
       def fillCopy(value: Float): P = p1.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => Tensor(a.shape, VType[V]).fill(value))
 
-    extension [F[_], V](p: F[V])(using tt: TensorTree[F[V]], ft: FloatTree[F[V], V], isF: IsFloating[V])
+    extension [F[_], V](p: F[V])(using tt: TensorTree[F[V]], ft: TreeOf[F[V], V], isF: IsFloating[V])
 
       def asFloats[NewV: IsFloating](vtype: VType[NewV])(using m: Mirror.ProductOf[F[NewV]]): F[NewV] =
         p.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a.asFloat(vtype)).asInstanceOf[F[NewV]]
-
-type FloatTreeFor[V] = [P] =>> FloatTree[P, V]
-
-/** A typeclass that proves P is a FloatTree, hiding the specific float type V
-  * from method signatures while keeping the evidence available.
-  */
-trait IsFloatTree[P]:
-  type V
-  given isFloating: IsFloating[V]
-  given floatTree: FloatTree[P, V]
-  given tensorTree: TensorTree[P]
-
-object IsFloatTree:
-  // The compiler automatically packages a FloatTree[P, V] into an IsFloatTree[P]
-  given pack[P, V0](using ft: FloatTree[P, V0], tt: TensorTree[P], isF: IsFloating[V0]): IsFloatTree[P] with
-    type V = V0
-    val isFloating = isF
-    val floatTree = ft
-    val tensorTree = tt
-
-  given unpackFloatTree[P](using isFT: IsFloatTree[P]): FloatTree[P, isFT.V] = isFT.floatTree
-  given unpackIsFloating[P](using isFT: IsFloatTree[P]): IsFloating[isFT.V] = isFT.isFloating
-  given unpackTensorTree[P](using isFT: IsFloatTree[P]): TensorTree[P] = isFT.tensorTree
diff --git a/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala b/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala
index 02a87e4..e2b03c5 100644
--- a/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala
+++ b/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala
@@ -2,20 +2,17 @@ package dimwit.optimizer
 
 import dimwit.*
 import dimwit.Conversions.given
-import dimwit.tensortree.FloatTree.*
-import dimwit.tensortree.FloatTree.ops.*
-import dimwit.autodiff.*
 
 class GradientOptimizerSuite extends DimwitTest:
 
   describe("GradientDescent"):
     it("should converge towards the minimum of f(x) = (x+1)^2 at x = -1"):
-      val optimizer = GradientDescent.of(VType[Float32])(learningRate = 0.1)
+      val optimizer = GradientDescent(learningRate = 0.1)
       val minX = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(1000).next()
       minX.item shouldBe -1.0f +- 0.1f
 
     it("should compute the step (single step)"):
-      val optimizer = GradientDescent.of(VType[Float32])(learningRate = 0.1)
+      val optimizer = GradientDescent(learningRate = 0.1)
       val initParams = Tensor0(2.0f)
       val initState = optimizer.init(initParams)
 
@@ -24,25 +21,14 @@ class GradientOptimizerSuite extends DimwitTest:
 
       nextParams.item shouldBe 1.4f +- 1e-5f
 
-  describe("LearningRateSchedule"):
-    it("basic"):
-      def learningRateSchedule(step: Tensor0[Int32]): Tensor0[Float32] =
-        if step.item <= 5 then Tensor0(0.0f)
-        else Tensor0(0.1f)
-      val optimizer = LearningRateSchedule.of(VType[Float32])(lr => GradientDescent(lr), learningRateSchedule)
-      val x1 = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(5).next()
-      x1.item shouldBe 2.0f +- 0.1
-      val x2 = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(1000).next()
-      x2.item shouldBe -1.0f +- 0.1f
-
   describe("Adam"):
     it("should converge towards the minimum of f(x) = (x+1)^2 at x = -1"):
-      val optimizer = Adam.of(VType[Float32])(learningRate = 0.1f)
+      val optimizer = Adam(learningRate = 0.1f)
       val minX = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(1000).next()
       minX.item shouldBe -1.0f +- 0.1f
 
     it("should compute the exact momentum and velocity updates (single step)"):
-      val optimizer = Adam.of(VType[Float32])(learningRate = 0.1f, b1 = 0.9f, b2 = 0.999f)
+      val optimizer = Adam(learningRate = 0.1f, beta1 = 0.9f, beta2 = 0.999f)
       val initParams = Tensor0(2.0f)
       val initState = optimizer.init(initParams)
 
@@ -52,17 +38,17 @@ class GradientOptimizerSuite extends DimwitTest:
       nextParams.item shouldBe 1.9f +- 1e-5f
       nextState.momentums.item shouldBe 0.6f +- 1e-5f
       nextState.velocities.item shouldBe 0.036f +- 1e-5f
-      nextState.b1.item shouldBe 0.9f +- 1e-5f
-      nextState.b2.item shouldBe 0.999f +- 1e-5
+      nextState.beta1t.item shouldBe 0.9f +- 1e-5f
+      nextState.beta2t.item shouldBe 0.999f +- 1e-5
 
   describe("AdamW"):
     it("should converge towards the minimum of f(x) = (x+1)^2 at x = -1"):
-      val optimizer = AdamW(Adam.of(VType[Float32])(learningRate = 0.1f), weightDecayFactor = 0.1f)
+      val optimizer = AdamW(Adam(learningRate = 0.1f), weightDecayFactor = 0.1f)
       val minX = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(1000).next()
       minX.item shouldBe -1.0f +- 0.1f
 
     it("should apply decoupled weight decay (single step)"):
-      val adam = Adam.of(VType[Float32])(learningRate = 0.1f)
+      val adam = Adam(learningRate = 0.1f)
       val adamW = AdamW(adam, weightDecayFactor = 0.1f)
 
       val initParams = Tensor0(2.0f)
@@ -74,12 +60,12 @@ class GradientOptimizerSuite extends DimwitTest:
 
   describe("Lion"):
     it("should converge towards the minimum of f(x) = (x+1)^2"):
-      val optimizer = Lion.of(VType[Float32])(learningRate = Tensor0(0.1f))
+      val optimizer = Lion(learningRate = 0.1f)
       val minX = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(1000).next()
       minX.item shouldBe -1.0f +- 0.1f
 
     it("should compute the exact sign-based update and momentum (single step)"):
-      val optimizer = Lion.of(VType[Float32])(learningRate = Tensor0(0.1f), beta1 = Tensor0(0.9f), beta2 = Tensor0(0.99f))
+      val optimizer = Lion(learningRate = 0.1f, beta1 = 0.9f, beta2 = 0.99f)
       val initParams = Tensor0(2.0f)
       val initMomentum = optimizer.init(initParams)
 
diff --git a/core/src/test/scala/dimwit/tensortree/FloatTensorTreeSuite.scala b/core/src/test/scala/dimwit/tensortree/TreeOfSuite.scala
similarity index 91%
rename from core/src/test/scala/dimwit/tensortree/FloatTensorTreeSuite.scala
rename to core/src/test/scala/dimwit/tensortree/TreeOfSuite.scala
index 418baa4..0b521bf 100644
--- a/core/src/test/scala/dimwit/tensortree/FloatTensorTreeSuite.scala
+++ b/core/src/test/scala/dimwit/tensortree/TreeOfSuite.scala
@@ -2,13 +2,34 @@ package dimwit.tensortree
 
 import dimwit.*
 import dimwit.Conversions.given
-import dimwit.tensortree.FloatTree.*
-import dimwit.tensortree.FloatTree.ops.*
+import dimwit.tensortree.TreeOf.*
+import dimwit.tensortree.TreeOf.given
+import dimwit.tensortree.TreeOf.ops.*
+import dimwit.tensor.DType.float32IsFloating
 
-class FloatTensorTreeSuite extends DimwitTest:
+class TreeOfSuite extends DimwitTest:
 
   describe("map"):
-    it("1-level case class"):
+    it("1-level case class (int32)"):
+      case class Params(
+          val w1: Tensor1[A, Int32],
+          val b1: Tensor0[Int32],
+          val w2: Tensor2[A, B, Int32],
+          val b2: Tensor0[Int32]
+      )
+      val params = Params(
+        Tensor1(Axis[A]).fromArray(Array(1, 2, 3)),
+        Tensor0(5),
+        Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1, 2), Array(3, 4), Array(5, 6))),
+        Tensor0(25)
+      )
+      val res = params.map([T <: Tuple] => (labels: Labels[T]) ?=> (x: Tensor[T, Int32]) => x +! 5)
+      res.w1 should equal(params.w1 +! 5)
+      res.b1 should equal(params.b1 + 5)
+      res.w2 should equal(params.w2 +! 5)
+      res.b2 should equal(params.b2 + 5)
+
+    it("1-level case class (float32)"):
       case class Params(
           val w1: Tensor1[A, Float32],
           val b1: Tensor0[Float32],
diff --git a/docs/quickstart.md b/docs/quickstart.md
index fa4b058..cd9b8a8 100644
--- a/docs/quickstart.md
+++ b/docs/quickstart.md
@@ -44,7 +44,7 @@ def fit(x: Tensor2[Batch, Feature, Float32], y: Tensor1[Batch, Float32]): Iterat
   val gradFn = grad(loss(x, y))
 
   // gradient based optimization
-  val gd = GradientDescent.of(VType[Float32])(learningRate = 0.1f) // this is wrong, should be 0.1f not Tensor0
+  val gd = GradientDescent(learningRate = 0.1f)
   gd.iterate(p0)(gradFn)
 ```
 
diff --git a/examples/src/main/scala/dimwit/basic/LogisticRegression.scala b/examples/src/main/scala/dimwit/basic/LogisticRegression.scala
index e7c7c88..827f881 100644
--- a/examples/src/main/scala/dimwit/basic/LogisticRegression.scala
+++ b/examples/src/main/scala/dimwit/basic/LogisticRegression.scala
@@ -34,7 +34,7 @@ object LogisticRegression:
     case class Params(
         weights: Tensor1[Feature, Float32],
         bias: Tensor0[Float32]
-    ) derives TensorTree
+    )
 
     // The loss is a simple binary cross-entropy loss
     def loss(data: Tensor2[Sample, Feature, Float32], labels: Tensor1[Sample, Bool])(params: BinaryLogisticRegression.Params)
@@ -125,7 +125,7 @@ object LogisticRegression:
     val trainLoss = jit(BinaryLogisticRegression.loss(trainingData, trainLabels))
     val valLoss = jit(BinaryLogisticRegression.loss(valData, valLabels))
     val learningRate = 5e-1f
-    val gd = GradientDescent.of(VType[Float32])(learningRate)
+    val gd = GradientDescent(learningRate)
 
     // Training loop
     val numiterations = 1000
diff --git a/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala b/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala
index ea72c1f..39a7e8c 100644
--- a/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala
+++ b/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala
@@ -2,7 +2,7 @@ package dimwit.examples.complex.vae
 
 import dimwit.Conversions.given
 import dimwit.*
-import dimwit.tensortree.FloatTree.*
+import dimwit.tensortree.TreeOf.*
 import dimwit.autodiff.*
 import dimwit.nn.ActivationFunctions.relu
 import dimwit.nn.ActivationFunctions.sigmoid
@@ -209,15 +209,15 @@ object VariationalAutoencoderExample:
       losses.sum / batchSize.toFloat
 
     val batches = trainImages.chunk(Axis[TrainSample], numSamples / batchSize)
-    val optimizer = GradientDescent.of(VType[Float32])(learningRate = learningRate)
-    def trainBatch(trainKey: Random.Key, batch: Tensor3[TrainSample, Height, Width, Float32], params: Params, state: optimizer.State[Params]): (Params, optimizer.State[Params]) =
+    val optimizer = GradientDescent(learningRate = learningRate)
+    def trainBatch(trainKey: Random.Key, batch: Tensor3[TrainSample, Height, Width, Float32], params: Params, state: Unit): (Params, Unit) =
       val grads = grad(batchLoss(trainKey, batch))(params)
       val (newParams, newState) = optimizer.update(grads, params, state)
       (newParams, newState)
 
     val (jitDonate, jitStep, jitReclaim) = jitDonating(trainBatch)
 
-    def trainEpoch(key: Random.Key, epoch: Int, params: Params, state: optimizer.State[Params]): (Params, optimizer.State[Params]) =
+    def trainEpoch(key: Random.Key, epoch: Int, params: Params, state: Unit): (Params, Unit) =
       val batchKeys = key.split(batches.size)
       jitReclaim(
         batches.zip(batchKeys).foldLeft(jitDonate(params, state)):
diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md
index e5e6496..57dc3cd 100644
--- a/mdocs/AGENTS.md
+++ b/mdocs/AGENTS.md
@@ -573,8 +573,8 @@ Use **case classes** to group parameters. DimWit automatically derives `TensorTr
 
 ```scala mdoc:reset:silent
 import dimwit.*
-import dimwit.autodiff.{Autodiff}
-import dimwit.tensortree.{TensorTree, FloatTree}
+import dimwit.autodiff.Autodiff
+import dimwit.tensortree.TensorTree
 
 trait Feature derives Label
 trait Hidden derives Label
@@ -711,7 +711,7 @@ val lossFunc = mse(trainData, trainLabels)
 val gradFunc = Autodiff.grad(lossFunc)
 
 // Create optimizer
-val optimizer = GradientDescent.of(VType[Float32])(learningRate = 0.01f)
+val optimizer = GradientDescent(learningRate = 0.01f)
 
 // Training loop with iterator
 val trained = optimizer.iterate(initModelParams)(gradFunc)
@@ -729,7 +729,7 @@ import dimwit.optimizer.Lion
 import dimwit.Conversions.given // enables implicit conversion from Float to Tensor[V]
 
 // Lion optimizer with momentum
-val lionOptimizer = Lion.of(VType[Float32])(learningRate = 1e-3f, beta1 = 0.9f, beta2 = 0.99f, weightDecay = 0.0f)
+val lionOptimizer = Lion(learningRate = 1e-3f, beta1 = 0.9f, beta2 = 0.99f, weightDecay = 0.0f)
 
 // Training with Lion
 val trainedLion = lionOptimizer.iterate(initModelParams)(gradFunc)
@@ -774,7 +774,7 @@ val initRegressionParams = RegressionParams(initSlope, initIntercept)
 
 // Train
 val regressionGrad = Autodiff.grad(regressionLoss(xData, yData))
-val gdOptimizer = GradientDescent.of(VType[Float32])(learningRate = 0.1f)
+val gdOptimizer = GradientDescent(learningRate = 0.1f)
 
 val finalParams = gdOptimizer.iterate(initRegressionParams)(regressionGrad)
   .take(100)
diff --git a/mdocs/docs/quickstart.md b/mdocs/docs/quickstart.md
index ea4c445..05e7dbd 100644
--- a/mdocs/docs/quickstart.md
+++ b/mdocs/docs/quickstart.md
@@ -44,7 +44,7 @@ def fit(x: Tensor2[Batch, Feature, Float32], y: Tensor1[Batch, Float32]): Iterat
   val gradFn = grad(loss(x, y))
 
   // gradient based optimization
-  val gd = GradientDescent.of(VType[Float32])(learningRate = 0.1f) // this is wrong, should be 0.1f not Tensor0
+  val gd = GradientDescent(learningRate = 0.1f)
   gd.iterate(p0)(gradFn)
 ```