diff --git a/src/main/scala/scanet/estimators/package.scala b/src/main/scala/scanet/estimators/package.scala
index e624d63..28e067c 100644
--- a/src/main/scala/scanet/estimators/package.scala
+++ b/src/main/scala/scanet/estimators/package.scala
@@ -5,7 +5,7 @@ import scanet.core.{Expr, Numeric, Session, Shape}
 
 import scala.{math => m}
 import scanet.math.syntax._
-import scanet.models.{TrainedModel_}
+import scanet.models.TrainedModel
 import scanet.optimizers.Iterators.Partial
 import scanet.optimizers.syntax._
 import scanet.optimizers.Record
@@ -16,9 +16,9 @@ import scala.collection.immutable.Seq
 package object estimators {
 
   def accuracy[A: Numeric](
-      model: TrainedModel_[A],
-      ds: Dataset[Record[A]],
-      batch: Int = 1000): Float = {
+                            model: TrainedModel[A],
+                            ds: Dataset[Record[A]],
+                            batch: Int = 1000): Float = {
     import ds.sparkSession.implicits._
     val brModel = ds.sparkSession.sparkContext.broadcast(model)
     val (positives, total) = ds
@@ -47,31 +47,31 @@ package object estimators {
   }
 
   def RMSE[A: Numeric](
-      model: TrainedModel_[A],
-      ds: Dataset[Record[A]],
-      batch: Int = 1000): Float =
+                        model: TrainedModel[A],
+                        ds: Dataset[Record[A]],
+                        batch: Int = 1000): Float =
     m.sqrt(MSE(model, ds, batch)).toFloat
 
   def MSE[A: Numeric](
-      model: TrainedModel_[A],
-      ds: Dataset[Record[A]],
-      batch: Int = 1000): Float =
+                       model: TrainedModel[A],
+                       ds: Dataset[Record[A]],
+                       batch: Int = 1000): Float =
     meanError(model, ds, batch) {
       (predicted, expected) => (predicted - expected).sqr
     }
 
   def MAE[A: Numeric](
-      model: TrainedModel_[A],
-      ds: Dataset[Record[A]],
-      batch: Int = 1000): Float =
+                       model: TrainedModel[A],
+                       ds: Dataset[Record[A]],
+                       batch: Int = 1000): Float =
     meanError(model, ds, batch) {
       (predicted, expected) => (predicted - expected).abs
     }
 
   private def meanError[A: Numeric](
-      model: TrainedModel_[A],
-      ds: Dataset[Record[A]],
-      batch: Int)(
+                                     model: TrainedModel[A],
+                                     ds: Dataset[Record[A]],
+                                     batch: Int)(
       error: (Expr[A], Expr[A]) => Expr[A]): Float = {
     import ds.sparkSession.implicits._
     val brModel = ds.sparkSession.sparkContext.broadcast(model)
@@ -100,9 +100,9 @@ package object estimators {
   }
 
   def R2Score[A: Numeric](
-      model: TrainedModel_[A],
-      ds: Dataset[Record[A]],
-      batch: Int = 1000): Float = {
+                           model: TrainedModel[A],
+                           ds: Dataset[Record[A]],
+                           batch: Int = 1000): Float = {
     require(ds.labelsShape == Shape(1), "labels should have shape (1)")
     import ds.sparkSession.implicits._
     val brModel = ds.sparkSession.sparkContext.broadcast(model)
diff --git a/src/main/scala/scanet/models/Math.scala b/src/main/scala/scanet/models/Math.scala
index 99cc672..e2d3661 100644
--- a/src/main/scala/scanet/models/Math.scala
+++ b/src/main/scala/scanet/models/Math.scala
@@ -12,13 +12,13 @@ object Math {
 
   case object `x^2` extends StatelessLayer {
 
-    override def params_(input: Shape): Params[ParamDef] =
+    override def params(input: Shape): Params[ParamDef] =
       Params(Weights -> ParamDef(Shape(), Initializer.Zeros, Some(Avg), trainable = true))
 
     override def buildStateless_[E: Floating](input: Expr[E], params: Params[Expr[E]]): Expr[E] =
       pow(params(Weights), 2)
 
-    override def penalty_[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
+    override def penalty[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
       zeros[E](Shape())
 
     override def outputShape(input: Shape): Shape = input
diff --git a/src/main/scala/scanet/models/Model.scala b/src/main/scala/scanet/models/Model.scala
index a109af2..d947a57 100644
--- a/src/main/scala/scanet/models/Model.scala
+++ b/src/main/scala/scanet/models/Model.scala
@@ -15,7 +15,7 @@ abstract class Model extends Serializable {
     * @param input input shape
     * @return param definitions
     */
-  def params_(input: Shape): Params[ParamDef]
+  def params(input: Shape): Params[ParamDef]
 
   /** Build a model
     *
@@ -23,30 +23,29 @@ abstract class Model extends Serializable {
     * @param params initialized or calculated model params
     * @return tuple where the first element is model output and second is changed params
     */
-  def build_[E: Floating](input: Expr[E], params: Params[Expr[E]]): (Expr[E], Params[Expr[E]])
+  def build[E: Floating](input: Expr[E], params: Params[Expr[E]]): (Expr[E], Params[Expr[E]])
 
   /** Additional model penalty to be added to the loss
     *
     * @param params initialized or calculated model params
     * @return penalty
     */
-  def penalty_[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E]
+  def penalty[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E]
 
-  def result_[E: Floating]: (Expr[E], Params[Expr[E]]) => Expr[E] =
-    (input, params) => build_(input, params)._1
+  def result[E: Floating]: (Expr[E], Params[Expr[E]]) => Expr[E] =
+    (input, params) => build(input, params)._1
 
-  def resultStateful_[E: Floating]: (Expr[E], Params[Expr[E]]) => (Expr[E], Params[Expr[E]]) =
-    (input, params) => build_(input, params)
+  def resultStateful[E: Floating]: (Expr[E], Params[Expr[E]]) => (Expr[E], Params[Expr[E]]) =
+    (input, params) => build(input, params)
 
-  // do we really need that???
   def outputShape(input: Shape): Shape
 
   def withLoss(loss: Loss): LossModel = LossModel(this, loss)
 
   private def makeGraph[E: Floating](input: Shape): Expr[E] =
-    build_(
+    build(
       input = placeholder[E](input),
-      params = params_(input).mapValues(paramDef => placeholder[E](paramDef.shape)))
+      params = params(input).mapValues(paramDef => placeholder[E](paramDef.shape)))
       ._1
 
   def displayResult[E: Floating](input: Shape, dir: String = ""): Unit =
@@ -56,7 +55,7 @@ abstract class Model extends Serializable {
     println(makeGraph[E](input).as("result").toString)
 
   def info(input: Shape): Seq[LayerInfo] = {
-    val (weights, state) = params_(input).partitionValues(_.trainable)
+    val (weights, state) = params(input).partitionValues(_.trainable)
     Seq(LayerInfo(
       toString,
       weights.values.map(_.shape).toList,
@@ -79,47 +78,47 @@ abstract class Model extends Serializable {
 
 case class LossModel(model: Model, lossF: Loss) extends Serializable {
 
-  def build_[E: Floating](
+  def build[E: Floating](
       input: Expr[E],
       output: Expr[E],
       params: Params[Expr[E]]): Expr[E] =
-    buildStateful_(input, output, params)._1
+    buildStateful(input, output, params)._1
 
-  def buildStateful_[E: Floating](
+  def buildStateful[E: Floating](
       input: Expr[E],
       output: Expr[E],
       params: Params[Expr[E]]): (Expr[E], Params[Expr[E]]) = {
-    val (result, nextParams) = model.build_(input, params)
-    val loss = lossF.build(result, output) plus model.penalty_(input.shape, params)
+    val (result, nextParams) = model.build(input, params)
+    val loss = lossF.build(result, output) plus model.penalty(input.shape, params)
     (loss, nextParams)
   }
 
-  def loss_[E: Floating]: (Expr[E], Expr[E], Params[Expr[E]]) => Expr[E] =
-    (input, output, params) => buildStateful_(input, output, params)._1
+  def loss[E: Floating]: (Expr[E], Expr[E], Params[Expr[E]]) => Expr[E] =
+    (input, output, params) => buildStateful(input, output, params)._1
 
-  def lossStateful_[E: Floating]
+  def lossStateful[E: Floating]
       : (Expr[E], Expr[E], Params[Expr[E]]) => (Expr[E], Params[Expr[E]]) =
-    (input, output, params) => buildStateful_(input, output, params)
+    (input, output, params) => buildStateful(input, output, params)
 
-  def grad_[E: Floating]: (Expr[E], Expr[E], Params[Expr[E]]) => Params[Expr[E]] =
+  def grad[E: Floating]: (Expr[E], Expr[E], Params[Expr[E]]) => Params[Expr[E]] =
     (input, output, weights) => {
-      val loss = build_(input, output, weights)
+      val loss = build(input, output, weights)
       loss.grad(weights).returns[E]
     }
 
-  def gradStateful_[E: Floating]
+  def gradStateful[E: Floating]
       : (Expr[E], Expr[E], Params[Expr[E]], Params[Expr[E]]) => (Params[Expr[E]], Params[Expr[E]]) =
     (input, output, weights, state) => {
-      val (loss, nextState) = buildStateful_(input, output, weights ++ state)
+      val (loss, nextState) = buildStateful(input, output, weights ++ state)
       val grad = loss.grad(weights).returns[E]
       (grad, nextState)
     }
 
-  def trained_[E: Floating](params: Params[Tensor[E]]) = new TrainedModel_(this, params)
+  def trained[E: Floating](params: Params[Tensor[E]]) = new TrainedModel(this, params)
 
   def displayLoss[E: Floating](input: Shape, dir: String = ""): Unit = {
-    val params = model.params_(input)
-    build_(
+    val params = model.params(input)
+    build(
       input = placeholder[E](input),
       output = placeholder[E](model.outputShape(input)),
       params = params.mapValues(paramDef => placeholder[E](paramDef.shape)))
@@ -128,8 +127,8 @@ case class LossModel(model: Model, lossF: Loss) extends Serializable {
   }
 
   def displayGrad[E: Floating](input: Shape, dir: String = ""): Unit = {
-    val (weights, state) = model.params_(input).partitionValues(_.trainable)
-    val (grad, _) = gradStateful_[E].apply(
+    val (weights, state) = model.params(input).partitionValues(_.trainable)
+    val (grad, _) = gradStateful[E].apply(
       placeholder[E](input),
       placeholder[E](model.outputShape(input)),
       weights.mapValues(paramDef => placeholder[E](paramDef.shape)),
@@ -142,13 +141,13 @@ case class LossModel(model: Model, lossF: Loss) extends Serializable {
   override def toString: String = s"$lossF($model)"
 }
 
-class TrainedModel_[E: Floating](val lossModel: LossModel, val params: Params[Tensor[E]]) {
+class TrainedModel[E: Floating](val lossModel: LossModel, val params: Params[Tensor[E]]) {
 
   def buildResult(input: Expr[E]): Expr[E] =
     buildResultStateful(input)._1
 
   def buildResultStateful(input: Expr[E]): (Expr[E], Params[Expr[E]]) =
-    lossModel.model.build_(input, params.mapValues(_.const))
+    lossModel.model.build(input, params.mapValues(_.const))
 
   def result: Expr[E] => Expr[E] = (input: Expr[E]) => buildResult(input)
 
@@ -161,7 +160,7 @@ class TrainedModel_[E: Floating](val lossModel: LossModel, val params: Params[Te
   def buildLossStateful(
       input: Expr[E],
       output: Expr[E]): (Expr[E], Params[Expr[E]]) =
-    lossModel.buildStateful_(input, output, params.mapValues(_.const))
+    lossModel.buildStateful(input, output, params.mapValues(_.const))
 
   def loss: (Expr[E], Expr[E]) => Expr[E] = (input, output) => buildLoss(input, output)
 
diff --git a/src/main/scala/scanet/models/layer/Activate.scala b/src/main/scala/scanet/models/layer/Activate.scala
index aadde1c..dafeba5 100644
--- a/src/main/scala/scanet/models/layer/Activate.scala
+++ b/src/main/scala/scanet/models/layer/Activate.scala
@@ -13,7 +13,7 @@ case class Activate(activation: Activation) extends NotTrainableLayer {
 
   override def name: String = activation.toString
 
-  override def build_[E: Floating](input: Expr[E]): Expr[E] =
+  override def build[E: Floating](input: Expr[E]): Expr[E] =
     activation.build(input)
 
   override def outputShape(input: Shape): Shape = input
diff --git a/src/main/scala/scanet/models/layer/Bias.scala b/src/main/scala/scanet/models/layer/Bias.scala
index 6d7f745..92b2364 100644
--- a/src/main/scala/scanet/models/layer/Bias.scala
+++ b/src/main/scala/scanet/models/layer/Bias.scala
@@ -23,13 +23,13 @@ import scala.collection.immutable.Seq
 case class Bias(features: Int, reg: Regularization = Zero, initializer: Initializer = Zeros)
     extends StatelessLayer {
 
-  override def params_(input: Shape): Params[ParamDef] =
+  override def params(input: Shape): Params[ParamDef] =
     Params(Weights -> ParamDef(Shape(features), initializer, Some(Avg), trainable = true))
 
   override def buildStateless_[E: Floating](input: Expr[E], params: Params[Expr[E]]): Expr[E] =
     input + params.weights
 
-  override def penalty_[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
+  override def penalty[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
     reg.build(params.weights)
 
   override def outputShape(input: Shape): Shape = input
diff --git a/src/main/scala/scanet/models/layer/Composed.scala b/src/main/scala/scanet/models/layer/Composed.scala
index eb54cbc..10a0158 100644
--- a/src/main/scala/scanet/models/layer/Composed.scala
+++ b/src/main/scala/scanet/models/layer/Composed.scala
@@ -13,27 +13,27 @@ import scala.collection.immutable.Seq
   */
 case class Composed(left: Layer, right: Layer) extends Layer {
 
-  override def params_(input: Shape): Params[ParamDef] = {
+  override def params(input: Shape): Params[ParamDef] = {
     // todo: flatten
-    val leftParams = left.params_(input).prependPath("l")
-    val rightParams = right.params_(left.outputShape(input)).prependPath("r")
+    val leftParams = left.params(input).prependPath("l")
+    val rightParams = right.params(left.outputShape(input)).prependPath("r")
     leftParams ++ rightParams
   }
 
-  override def build_[E: Floating](
+  override def build[E: Floating](
       input: Expr[E],
       params: Params[Expr[E]]): (Expr[E], Params[Expr[E]]) = {
     val leftParams = params.children("l")
     val rightParams = params.children("r")
-    val (leftOutput, leftState) = left.build_(input, leftParams)
-    val (rightOutput, rightState) = right.build_(leftOutput, rightParams)
+    val (leftOutput, leftState) = left.build(input, leftParams)
+    val (rightOutput, rightState) = right.build(leftOutput, rightParams)
     (rightOutput, leftState.prependPath("l") ++ rightState.prependPath("r"))
   }
 
-  override def penalty_[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] = {
+  override def penalty[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] = {
     val leftParams = params.children("l")
     val rightParams = params.children("r")
-    left.penalty_(input, leftParams) plus right.penalty_(left.outputShape(input), rightParams)
+    left.penalty(input, leftParams) plus right.penalty(left.outputShape(input), rightParams)
   }
 
   override def outputShape(input: Shape): Shape =
diff --git a/src/main/scala/scanet/models/layer/Conv2D.scala b/src/main/scala/scanet/models/layer/Conv2D.scala
index 45d30a3..81d4dce 100644
--- a/src/main/scala/scanet/models/layer/Conv2D.scala
+++ b/src/main/scala/scanet/models/layer/Conv2D.scala
@@ -80,7 +80,7 @@ case class Conv2D private (
   def filterHeight: Int = kernel._1
   def filterWidth: Int = kernel._2
 
-  override def params_(input: Shape): Params[ParamDef] = {
+  override def params(input: Shape): Params[ParamDef] = {
     require(
       input.rank == 4,
       s"Conv2D input should have a shape (NHWC) or (NCHW) but was $input")
@@ -101,7 +101,7 @@ case class Conv2D private (
       format = format)
   }
 
-  override def penalty_[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
+  override def penalty[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
     zeros[E](Shape())
 
   override def outputShape(input: Shape): Shape = {
diff --git a/src/main/scala/scanet/models/layer/Dense.scala b/src/main/scala/scanet/models/layer/Dense.scala
index 14603c7..4776159 100644
--- a/src/main/scala/scanet/models/layer/Dense.scala
+++ b/src/main/scala/scanet/models/layer/Dense.scala
@@ -44,7 +44,7 @@ object Dense {
 case class Dense private (outputs: Int, reg: Regularization, initializer: Initializer)
     extends StatelessLayer {
 
-  override def params_(input: Shape): Params[ParamDef] =
+  override def params(input: Shape): Params[ParamDef] =
     Params(Weights -> ParamDef(Shape(input(1), outputs), initializer, Some(Avg), trainable = true))
 
   override def buildStateless_[E: Floating](input: Expr[E], params: Params[Expr[E]]): Expr[E] =
@@ -53,7 +53,7 @@ case class Dense private (outputs: Int, reg: Regularization, initializer: Initia
     // x * w -> (samples, features) * (features, outputs) -> (samples, outputs)
     input matmul params.weights
 
-  override def penalty_[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
+  override def penalty[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
     reg.build(params.weights)
 
   override def outputShape(input: Shape): Shape = Shape(input.head, outputs)
diff --git a/src/main/scala/scanet/models/layer/Flatten.scala b/src/main/scala/scanet/models/layer/Flatten.scala
index 9961f53..cd43106 100644
--- a/src/main/scala/scanet/models/layer/Flatten.scala
+++ b/src/main/scala/scanet/models/layer/Flatten.scala
@@ -9,7 +9,7 @@ import scanet.core.{Expr, Floating, Shape}
   */
 case object Flatten extends NotTrainableLayer {
 
-  override def build_[E: Floating](input: Expr[E]): Expr[E] = {
+  override def build[E: Floating](input: Expr[E]): Expr[E] = {
     val shape = input.shape
     require(shape.rank >= 2, s"rank should be >= 2, but was ${shape.rank}")
     val batch = shape(0)
diff --git a/src/main/scala/scanet/models/layer/Layer.scala b/src/main/scala/scanet/models/layer/Layer.scala
index 9e32573..fe1ae61 100644
--- a/src/main/scala/scanet/models/layer/Layer.scala
+++ b/src/main/scala/scanet/models/layer/Layer.scala
@@ -35,7 +35,7 @@ trait StatelessLayer extends Layer {
 
   def buildStateless_[E: Floating](input: Expr[E], params: Params[Expr[E]]): Expr[E]
 
-  override def build_[E: Floating](
+  override def build[E: Floating](
       input: Expr[E],
       params: Params[Expr[E]]): (Expr[E], Params[Expr[E]]) = {
     (buildStateless_(input, params), Params.empty)
@@ -46,14 +46,14 @@ trait NotTrainableLayer extends StatelessLayer {
 
   override def trainable: Boolean = false
 
-  override def params_(input: Shape): Params[ParamDef] = Params.empty
-  override def penalty_[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
+  override def params(input: Shape): Params[ParamDef] = Params.empty
+  override def penalty[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
     zeros[E](Shape())
 
-  def build_[E: Floating](input: Expr[E]): Expr[E]
+  def build[E: Floating](input: Expr[E]): Expr[E]
 
   override def buildStateless_[E: Floating](input: Expr[E], params: Params[Expr[E]]): Expr[E] = {
     require(params.isEmpty, s"$this layer does not require params")
-    build_(input)
+    build(input)
   }
 }
diff --git a/src/main/scala/scanet/models/layer/Pool2D.scala b/src/main/scala/scanet/models/layer/Pool2D.scala
index c47efa6..2398230 100644
--- a/src/main/scala/scanet/models/layer/Pool2D.scala
+++ b/src/main/scala/scanet/models/layer/Pool2D.scala
@@ -33,7 +33,7 @@ case class Pool2D(
     reduce: Reduce = Reduce.Max)
     extends NotTrainableLayer {
 
-  override def build_[E: Floating](input: Expr[E]): Expr[E] =
+  override def build[E: Floating](input: Expr[E]): Expr[E] =
     pool2D[E](
       input = input,
       window = Seq(window._1, window._2),
diff --git a/src/main/scala/scanet/models/layer/RNN.scala b/src/main/scala/scanet/models/layer/RNN.scala
index c5e643a..57d58e4 100644
--- a/src/main/scala/scanet/models/layer/RNN.scala
+++ b/src/main/scala/scanet/models/layer/RNN.scala
@@ -29,15 +29,15 @@ case class RNN(cell: Layer, returnSequence: Boolean = false, stateful: Boolean =
 
   // todo: better params management
 
-  override def params_(input: Shape): Params[ParamDef] = {
+  override def params(input: Shape): Params[ParamDef] = {
     val (weights, state) = paramsPartitioned(input)
     if (stateful) state ++ weights else weights
   }
 
   private def paramsPartitioned(input: Shape): (Params[ParamDef], Params[ParamDef]) =
-    cell.params_(dropTime(input)).partitionValues(_.trainable)
+    cell.params(dropTime(input)).partitionValues(_.trainable)
 
-  override def build_[E: Floating](
+  override def build[E: Floating](
       input: Expr[E],
       params: Params[Expr[E]]): (Expr[E], Params[Expr[E]]) = {
     require(input.rank == 3, "RNN requires input to have Shape(batch, time, features)")
@@ -54,7 +54,7 @@ case class RNN(cell: Layer, returnSequence: Boolean = false, stateful: Boolean =
         step: Int,
         outputs: Seq[Expr[E]],
         state: Params[Expr[E]]): (Seq[Expr[E]], Params[Expr[E]]) = {
-      val (output, outputState) = cell.build_(stepsInput.slice(step), state ++ Params(weightParams))
+      val (output, outputState) = cell.build(stepsInput.slice(step), state ++ Params(weightParams))
       if (step < timeSteps - 1) {
         stackCells(step + 1, outputs :+ output, outputState)
       } else {
@@ -71,8 +71,8 @@ case class RNN(cell: Layer, returnSequence: Boolean = false, stateful: Boolean =
     (output, lastOutputState)
   }
 
-  override def penalty_[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
-    cell.penalty_(dropTime(input), params)
+  override def penalty[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
+    cell.penalty(dropTime(input), params)
 
   private def dropTime(input: Shape): Shape = input.remove(1)
 
@@ -163,7 +163,7 @@ case class SimpleRNNCell(
 
   override def stateful: Boolean = true
 
-  override def params_(input: Shape): Params[ParamDef] = Params(
+  override def params(input: Shape): Params[ParamDef] = Params(
     // (features, units)
     Kernel -> ParamDef(Shape(input(1), units), kernelInitializer, Some(Avg), trainable = true),
     // (units, units)
@@ -171,7 +171,7 @@ case class SimpleRNNCell(
     // state, keeps previous output
     State -> ParamDef(outputShape(input), Initializer.Zeros))
 
-  override def build_[E: Floating](
+  override def build[E: Floating](
       input: Expr[E],
       params: Params[Expr[E]]): (Expr[E], Params[Expr[E]]) = {
     require(input.rank >= 2, "SimpleRNNCell requires input Seq(batch, features)")
@@ -179,7 +179,7 @@ case class SimpleRNNCell(
     (result, Params(State -> result))
   }
 
-  override def penalty_[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
+  override def penalty[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
     kernelReg.build(params(Kernel)) + recurrentReg.build(params(Recurrent))
 
   override def outputShape(input: Shape): Shape = Shape(input.head, units)
@@ -265,16 +265,15 @@ case class LSTMCell(
   private val iCell = cell("input", recurrentActivation, useBias = biasInitializer)
   private val gCell = cell("gate", activation, useBias = biasInitializer)
   private val oCell = cell("output", recurrentActivation, useBias = biasInitializer)
-  private val cells = Seq(fCell, iCell, gCell, oCell).map(_._2)
-  private val cells_ = Map(fCell, iCell, gCell, oCell)
+  private val cells = Map(fCell, iCell, gCell, oCell)
 
   override def stateful: Boolean = true
 
-  override def params_(input: Shape): Params[ParamDef] = {
-    val weights = cells_
+  override def params(input: Shape): Params[ParamDef] = {
+    val weights = cells
       .map {
         case (name, cell) =>
-          val params = cell.params_(input)
+          val params = cell.params(input)
           val onlyWeights = params.filterPaths(path => !path.endsWith(SimpleRNNCell.State))
           onlyWeights.prependPath(name)
       }
@@ -296,33 +295,33 @@ case class LSTMCell(
     *  - output h: (batch, units)
     *  - output y: (batch, units)
     */
-  override def build_[E: Floating](
+  override def build[E: Floating](
       input: Expr[E],
       params: Params[Expr[E]]): (Expr[E], Params[Expr[E]]) = {
     require(input.rank >= 2, "LSTMCell requires input Seq(batch, features)")
     val cPrev = params(CellState)
     val hPrev = params(HiddenState)
-    val List(f, i, g, o) = cells_.toList.map {
+    val List(f, i, g, o) = cells.toList.map {
       case (name, cell) =>
-        val cellState = cell.params_(input.shape)
+        val cellState = cell.params(input.shape)
           .filter {
             case (path, param) =>
               path.endsWith(Params.State) && param.nonTrainable && param.shape == hPrev.shape
           }
           .mapValues(_ => hPrev)
         val cellParams = params.children(name) ++ cellState
-        cell.build_(input, cellParams)._1
+        cell.build(input, cellParams)._1
     }
     val c = cPrev * f + i * g
     val h = o * activation.build(c)
     (h, Params(CellState -> c, HiddenState -> h))
   }
 
-  override def penalty_[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
-    cells_.foldLeft(Floating[E].zero.const) {
+  override def penalty[E: Floating](input: Shape, params: Params[Expr[E]]): Expr[E] =
+    cells.foldLeft(Floating[E].zero.const) {
       case (sum, (name, cell)) =>
         val cellParams = params.children(name)
-        sum + cell.penalty_(input, cellParams)
+        sum + cell.penalty(input, cellParams)
     }
 
   override def outputShape(input: Shape): Shape =
diff --git a/src/main/scala/scanet/optimizers/Effect.scala b/src/main/scala/scanet/optimizers/Effect.scala
index 6f0f325..19cb719 100644
--- a/src/main/scala/scanet/optimizers/Effect.scala
+++ b/src/main/scala/scanet/optimizers/Effect.scala
@@ -55,7 +55,7 @@ object Effect {
       tensorboard: Boolean = false)
       extends Effect[E] {
     override def apply(state: State, next: StepContext[E]): State = {
-      val trained = next.lossModel.trained_(next.result.params)
+      val trained = next.lossModel.trained(next.result.params)
       val a = accuracy(trained, ds, next.step.batch)
       if (tensorboard)
         state.board.addScalar("accuracy", a, next.step.iter)
diff --git a/src/main/scala/scanet/optimizers/KryoSerializers.scala b/src/main/scala/scanet/optimizers/KryoSerializers.scala
index 850cdec..e35cf99 100644
--- a/src/main/scala/scanet/optimizers/KryoSerializers.scala
+++ b/src/main/scala/scanet/optimizers/KryoSerializers.scala
@@ -4,7 +4,7 @@ import com.esotericsoftware.kryo.Kryo
 import com.twitter.chill._
 import org.apache.spark.serializer.KryoRegistrator
 import scanet.core.{Shape, Tensor, TensorType}
-import scanet.models.TrainedModel_
+import scanet.models.TrainedModel
 import org.tensorflow.proto.framework.DataType.DT_STRING
 import scanet.core.syntax._
 
@@ -12,7 +12,7 @@ class KryoSerializers extends KryoRegistrator {
 
   override def registerClasses(kryo: Kryo): Unit = {
     kryo.forClass[Tensor[_]](new TensorSerializer())
-    kryo.register(classOf[TrainedModel_[_]])
+    kryo.register(classOf[TrainedModel[_]])
   }
 }
 
diff --git a/src/main/scala/scanet/optimizers/Optimizer.scala b/src/main/scala/scanet/optimizers/Optimizer.scala
index 11eeec8..b942f47 100644
--- a/src/main/scala/scanet/optimizers/Optimizer.scala
+++ b/src/main/scala/scanet/optimizers/Optimizer.scala
@@ -55,7 +55,7 @@ case class Optimizer[A: Floating](
 
   private val lossModel = model.withLoss(loss)
 
-  def run(): TrainedModel_[A] = {
+  def run(): TrainedModel[A] = {
     val jobId = UUID.randomUUID().toString
     val ds: Dataset[Record[A]] = dataset.cache()
     val sc = ds.sparkSession.sparkContext
@@ -93,7 +93,7 @@ case class Optimizer[A: Floating](
     }
     val params = optimize(Step(batchSize), Params.empty, Params.empty)
 
-    lossModel.trained_(params)
+    lossModel.trained(params)
   }
 
   private def optimizeOnPartition(
@@ -107,7 +107,7 @@ case class Optimizer[A: Floating](
     val result = resource.sessionPool.within(session => {
       val batches = TensorIterator(it, shapes, batchSize)
       val batchedInputShape = batchSize +: shapes._1
-      val paramsDef = model.params_(batchedInputShape)
+      val paramsDef = model.params(batchedInputShape)
       val weightNames = paramsDef.filterValues(_.trainable).paths
       val (paramsInitialized, metaInitialized) =
         if (globalIter == 0) {
@@ -147,7 +147,7 @@ case class Optimizer[A: Floating](
   private def compileLoss(cache: Optimizer.Cache, session: Session) = {
     cache.getOrCompute(
       s"$lossModel:loss[${TensorType[A].classTag}]",
-      lossModel.lossStateful_[A].tf) compileWith session
+      lossModel.lossStateful[A].tf) compileWith session
   }
 
   private def compileBackprop(cache: Optimizer.Cache, session: Session) = {
@@ -161,7 +161,7 @@ case class Optimizer[A: Floating](
             metas: Params[Expr[A]],
             state: Params[Expr[A]],
             iter: Expr[Int]) => {
-          val (grads, nextState) = lossModel.gradStateful_[A].apply(x, y, weights, state)
+          val (grads, nextState) = lossModel.gradStateful[A].apply(x, y, weights, state)
           val (nextAcc, nextMeta) = weights.join(grads).join(metas).params
             .foldLeft((newOutputSeq, newOutputSeq)) { (acc, next) =>
               val (gAcc, metaAcc) = acc
@@ -288,7 +288,7 @@ object Optimizer {
 
     @nowarn def build(implicit ev: State =:= Complete): Optimizer[A] = optimizer
 
-    def run()(implicit ev: State =:= Complete): TrainedModel_[A] = build.run()
+    def run()(implicit ev: State =:= Complete): TrainedModel[A] = build.run()
   }
 
   def minimize[R: Floating](model: Model)(
diff --git a/src/main/scala/scanet/research/Neuron.scala b/src/main/scala/scanet/research/Neuron.scala
index 44c03a9..e2f23d1 100644
--- a/src/main/scala/scanet/research/Neuron.scala
+++ b/src/main/scala/scanet/research/Neuron.scala
@@ -26,12 +26,12 @@ object Neuron {
     val w = Tensor.matrix(Array(0.6f, 0.35f, 0.9f))
     val params = Params(Weights -> w)
     // to make a prediction we need to run a forward pass
-    val result = neuron.result_[Float].compile
+    val result = neuron.result[Float].compile
     // (0.7 * 50 + 0.5 * 25 + 1 * 1)
     println(result(x, params))
     val paramsExpr = params.mapValues(_.const)
     // let's calculate prediction error (loss)
-    val loss = MeanSquaredError.build(neuron.build_(x.const, paramsExpr)._1, y.const)
+    val loss = MeanSquaredError.build(neuron.build(x.const, paramsExpr)._1, y.const)
     println(loss.eval)
     // let's calculate a gradient
     val grads = loss.grad(paramsExpr).returns[Float]
diff --git a/src/test/scala/scanet/models/RegressionSpec.scala b/src/test/scala/scanet/models/RegressionSpec.scala
index 19d0461..355f2f9 100644
--- a/src/test/scala/scanet/models/RegressionSpec.scala
+++ b/src/test/scala/scanet/models/RegressionSpec.scala
@@ -11,7 +11,7 @@ import scanet.test.CustomMatchers
 class RegressionSpec extends AnyFlatSpec with CustomMatchers {
 
   "linear regression" should "calculate loss with Float precision" in {
-    val loss = LinearRegression().withLoss(MeanSquaredError).loss_[Float].compile
+    val loss = LinearRegression().withLoss(MeanSquaredError).loss[Float].compile
     val x = Tensor.matrix(Array(1.0f, 2.0f), Array(2.0f, 4.0f))
     val y = Tensor.matrix(Array(6.0f), Array(12.0f))
     val weights = Tensor.matrix(Array(2.0f), Array(3.0f))
@@ -23,7 +23,7 @@ class RegressionSpec extends AnyFlatSpec with CustomMatchers {
   }
 
   it should "calculate loss with Double precision" in {
-    val loss = LinearRegression().withLoss(MeanSquaredError).loss_[Double].compile
+    val loss = LinearRegression().withLoss(MeanSquaredError).loss[Double].compile
     val x = Tensor.matrix(Array(1.0, 2.0), Array(2.0, 4.0))
     val y = Tensor.matrix(Array(6.0), Array(12.0))
     val weights = Tensor.matrix(Array(2.0), Array(3.0))
@@ -35,7 +35,7 @@ class RegressionSpec extends AnyFlatSpec with CustomMatchers {
   }
 
   it should "calculate result" in {
-    val result = LinearRegression().result_[Float].compile
+    val result = LinearRegression().result[Float].compile
     val x = Tensor.matrix(Array(1.0f, 2.0f), Array(2.0f, 4.0f))
     val y = Tensor.matrix(Array(9.0f), Array(17.0f))
     val weights = Tensor.matrix(Array(2.0f), Array(3.0f))
@@ -47,7 +47,7 @@ class RegressionSpec extends AnyFlatSpec with CustomMatchers {
   }
 
   it should "calculate gradient" in {
-    val grad = LinearRegression().withLoss(MeanSquaredError).grad_[Float].compile
+    val grad = LinearRegression().withLoss(MeanSquaredError).grad[Float].compile
     val x = Tensor.matrix(Array(1.0f, 2.0f), Array(2.0f, 4.0f))
     val y = Tensor.matrix(Array(6.0f), Array(12.0f))
     val weights = Tensor.matrix(Array(0.0f), Array(0.0f))
@@ -65,7 +65,7 @@ class RegressionSpec extends AnyFlatSpec with CustomMatchers {
   }
 
   "logistic regression" should "calculate loss" in {
-    val regression = LogisticRegression().withLoss(BinaryCrossentropy).loss_[Float].compile
+    val regression = LogisticRegression().withLoss(BinaryCrossentropy).loss[Float].compile
     val x = Tensor.matrix(Array(0.34f, 0.78f), Array(0.6f, 0.86f))
     val y = Tensor.matrix(Array(0.402f), Array(0.47800002f))
     val weights = Tensor.matrix(Array(0.2f), Array(0.3f))
@@ -77,7 +77,7 @@ class RegressionSpec extends AnyFlatSpec with CustomMatchers {
   }
 
   it should "calculate result" in {
-    val result = LogisticRegression().result_[Float].compile
+    val result = LogisticRegression().result[Float].compile
     val x = Tensor.matrix(Array(0.34f, 0.78f), Array(0.6f, 0.86f))
     val y = Tensor.matrix(Array(0.599168f), Array(0.617276f))
     val weights = Tensor.matrix(Array(0.2f), Array(0.3f))
@@ -90,7 +90,7 @@ class RegressionSpec extends AnyFlatSpec with CustomMatchers {
   }
 
   it should "calculate gradient " in {
-    val grad = LogisticRegression().withLoss(BinaryCrossentropy).grad_[Float].compile
+    val grad = LogisticRegression().withLoss(BinaryCrossentropy).grad[Float].compile
     val x = Tensor.matrix(Array(0.34f, 0.78f), Array(0.6f, 0.86f))
     val y = Tensor.matrix(Array(0.402f), Array(0.47800002f))
     val weights = Tensor.matrix(Array(0.2f), Array(0.3f))
diff --git a/src/test/scala/scanet/models/layer/ActivateLayerSpec.scala b/src/test/scala/scanet/models/layer/ActivateLayerSpec.scala
index d2c4446..659d466 100644
--- a/src/test/scala/scanet/models/layer/ActivateLayerSpec.scala
+++ b/src/test/scala/scanet/models/layer/ActivateLayerSpec.scala
@@ -19,7 +19,7 @@ class ActivateLayerSpec extends AnyWordSpec with CustomMatchers {
         Array(1f, 0f, 1f),
         Array(1f, 1f, 1f))
       val model = Sigmoid.layer
-      val result = model.result_[Float].compile
+      val result = model.result[Float].compile
       val y = Tensor.matrix(
         Array(0.5f, 0.5f, 0.7310586f),
         Array(0.5f, 0.7310586f, 0.7310586f),
diff --git a/src/test/scala/scanet/models/layer/BiasLayerSpec.scala b/src/test/scala/scanet/models/layer/BiasLayerSpec.scala
index cb7f95d..399fa20 100644
--- a/src/test/scala/scanet/models/layer/BiasLayerSpec.scala
+++ b/src/test/scala/scanet/models/layer/BiasLayerSpec.scala
@@ -20,7 +20,7 @@ class BiasLayerSpec extends AnyWordSpec with CustomMatchers {
         Array(1f, 1f, 1f))
       val b = Tensor.vector(1f, 2f, 3f)
       val model = Bias(3)
-      val result = model.result_[Float].compile
+      val result = model.result[Float].compile
       val y = Tensor.matrix(
         Array(1f, 2f, 4f),
         Array(1f, 3f, 4f),
diff --git a/src/test/scala/scanet/models/layer/ComposedLayerSpec.scala b/src/test/scala/scanet/models/layer/ComposedLayerSpec.scala
index e00fa1a..7f44dda 100644
--- a/src/test/scala/scanet/models/layer/ComposedLayerSpec.scala
+++ b/src/test/scala/scanet/models/layer/ComposedLayerSpec.scala
@@ -34,7 +34,7 @@ class ComposedLayerSpec extends AnyWordSpec with CustomMatchers with SharedSpark
         Array(1f),
         Array(0f))
       val b2 = Tensor.vector(0f)
-      val forward = model.result_[Float].compile
+      val forward = model.result[Float].compile
       val expected = Tensor.matrix(
         Array(0.705357f),
         Array(0.770136f),
@@ -66,7 +66,7 @@ class ComposedLayerSpec extends AnyWordSpec with CustomMatchers with SharedSpark
       "l" / "l" / "r" / Weights -> b1,
       "r" / "l" / "l" / Weights -> w2,
       "r" / "l" / "r" / Weights -> b2)
-    model.penalty_(Shape(1, 4), params.mapValues(_.const)).eval should be(
+    model.penalty(Shape(1, 4), params.mapValues(_.const)).eval should be(
       Tensor.scalar(3.83f))
   }
 
@@ -91,7 +91,7 @@ class ComposedLayerSpec extends AnyWordSpec with CustomMatchers with SharedSpark
       "l" / "l" / "r" / Weights -> b1,
       "r" / "l" / "l" / Weights -> w2,
       "r" / "l" / "r" / Weights -> b2)
-    val loss = model.withLoss(MeanSquaredError).loss_[Float].compile
+    val loss = model.withLoss(MeanSquaredError).loss[Float].compile
     val result = loss(x, y, params).const.roundAt(6).eval
     result should be(Tensor.scalar(0.339962f))
   }
@@ -117,7 +117,7 @@ class ComposedLayerSpec extends AnyWordSpec with CustomMatchers with SharedSpark
       "l" / "l" / "r" / Weights -> b1,
       "r" / "l" / "l" / Weights -> w2,
       "r" / "l" / "r" / Weights -> b2)
-    val loss = model.withLoss(MeanSquaredError).loss_[Float].compile
+    val loss = model.withLoss(MeanSquaredError).loss[Float].compile
     loss(x, y, params) should be(Tensor.scalar(3.6199622f))
   }
 }
diff --git a/src/test/scala/scanet/models/layer/Conv2DLayerSpec.scala b/src/test/scala/scanet/models/layer/Conv2DLayerSpec.scala
index 825d84b..b36ce80 100644
--- a/src/test/scala/scanet/models/layer/Conv2DLayerSpec.scala
+++ b/src/test/scala/scanet/models/layer/Conv2DLayerSpec.scala
@@ -33,7 +33,7 @@ class Conv2DLayerSpec extends AnyWordSpec with CustomMatchers {
         Array(7.0, 11.0, 16.0, 7.0),
         Array(10.0, 7.0, 4.0, 7.0))
         .reshape(1, 4, 4, 1)
-      val result = model.result_[Double].compile
+      val result = model.result[Double].compile
       result(input, Params(Weights -> filters)).const.roundAt(2).eval shouldBe output
     }
 
diff --git a/src/test/scala/scanet/models/layer/DenseLayerSpec.scala b/src/test/scala/scanet/models/layer/DenseLayerSpec.scala
index 6d40bbb..c56c92c 100644
--- a/src/test/scala/scanet/models/layer/DenseLayerSpec.scala
+++ b/src/test/scala/scanet/models/layer/DenseLayerSpec.scala
@@ -36,7 +36,7 @@ class DenseLayerSpec extends AnyWordSpec with CustomMatchers {
         Array(0.880797f, 0.622459f, 0.768525f, 0.598688f),
         Array(0.890903f, 0.817574f, 0.900250f, 0.802184f))
       val model = Dense(4, Sigmoid)
-      val forward = model.result_[Float].compile
+      val forward = model.result[Float].compile
       val params = Params("l" / "l" / Weights -> w, "l" / "r" / Weights -> b)
       val y = forward(x, params).const.roundAt(6).eval
       y should be(yExpected)
@@ -51,12 +51,12 @@ class DenseLayerSpec extends AnyWordSpec with CustomMatchers {
       val b = Tensor.vector(0f, 0f)
       val model = Dense(4, Sigmoid, reg = L2(lambda = 1))
       val params = Params("l" / "l" / Weights -> w, "l" / "r" / Weights -> b)
-      model.penalty_(Shape(1, 4), params.mapValues(_.const)).eval should be(Tensor.scalar(1.63f))
+      model.penalty(Shape(1, 4), params.mapValues(_.const)).eval should be(Tensor.scalar(1.63f))
     }
 
     "produce gradient when combined with loss function" in {
       val loss = Dense(4, Sigmoid).withLoss(BinaryCrossentropy)
-      val grad = loss.grad_[Float].compile
+      val grad = loss.grad[Float].compile
       val x = Tensor.matrix(
         Array(0f, 0f, 1f),
         Array(0f, 1f, 1f),
diff --git a/src/test/scala/scanet/models/layer/FlattenLayerSpec.scala b/src/test/scala/scanet/models/layer/FlattenLayerSpec.scala
index 38e5980..6909823 100644
--- a/src/test/scala/scanet/models/layer/FlattenLayerSpec.scala
+++ b/src/test/scala/scanet/models/layer/FlattenLayerSpec.scala
@@ -13,7 +13,7 @@ class FlattenLayerSpec extends AnyWordSpec with CustomMatchers {
     "flatten the output from (batch, features_1, ... features_n) into (batch, features) tensor" in {
       val x = Tensor.ones[Float](10, 5, 5)
       val model = Flatten
-      val result = model.result_[Float].compile
+      val result = model.result[Float].compile
       result(x, Params.empty) shouldBe x.reshape(10, 25)
     }
   }
diff --git a/src/test/scala/scanet/models/layer/Pool2DLayerSpec.scala b/src/test/scala/scanet/models/layer/Pool2DLayerSpec.scala
index 013e6dc..4f8316f 100644
--- a/src/test/scala/scanet/models/layer/Pool2DLayerSpec.scala
+++ b/src/test/scala/scanet/models/layer/Pool2DLayerSpec.scala
@@ -27,7 +27,7 @@ class Pool2DLayerSpec extends AnyWordSpec with CustomMatchers {
         Array(2.0, 3.0, 3.0, 3.0),
         Array(2.0, 3.0, 3.0, 2.0))
         .reshape(1, 4, 4, 1)
-      val result = model.result_[Double].compile
+      val result = model.result[Double].compile
       result(input, Params.empty).const.roundAt(2).eval shouldBe output
     }
 
diff --git a/src/test/scala/scanet/models/layer/RNNLayerSpec.scala b/src/test/scala/scanet/models/layer/RNNLayerSpec.scala
index e5dca1b..e9a95f6 100644
--- a/src/test/scala/scanet/models/layer/RNNLayerSpec.scala
+++ b/src/test/scala/scanet/models/layer/RNNLayerSpec.scala
@@ -24,7 +24,7 @@ class RNNLayerSpec extends AnyWordSpec with CustomMatchers {
         Array(-0.6313778f, 0.7754754f))
       val b = Tensor.vector(0f, 0f)
       val expected = Tensor.matrix(Array(0.222901f, -6.019066f))
-      val result = layer.result_[Float].compile
+      val result = layer.result[Float].compile
       val params = Params(
         "l" / "kernel_weights" -> wx,
         "l" / "recurrent_weights" -> wh,
@@ -74,7 +74,7 @@ class RNNLayerSpec extends AnyWordSpec with CustomMatchers {
         "output" / "l" / "l" / "recurrent_weights" -> wo(1),
         "output" / "l" / "r" / Weights -> wo(2),
        )
-      val result = layer.result_[Float].compile
+      val result = layer.result[Float].compile
       val prediction = result(input, params).const.roundAt(6).eval
       prediction shouldBe Tensor.matrix(Array(0.382158f, 0.029766f))
     }