<commit_before>/* Copyright 2017-19, <NAME>. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */

package org.platanios.tensorflow.api.ops.training.optimizers

import org.platanios.tensorflow.api.core.{Graph, Shape}
import org.platanios.tensorflow.api.core.types._
import org.platanios.tensorflow.api.implicits.Implicits._
import org.platanios.tensorflow.api.ops._
import org.platanios.tensorflow.api.ops.basic.Basic
import org.platanios.tensorflow.api.ops.control_flow.ControlFlow
import org.platanios.tensorflow.api.ops.math.Math
import org.platanios.tensorflow.api.ops.training.optimizers.Optimizer._
import org.platanios.tensorflow.api.ops.variables.{ConstantInitializer, Initializer, Variable, VariableScope}
import org.platanios.tensorflow.api.tensors.Tensor
import org.platanios.tensorflow.api.utilities.DefaultsTo.{AnyDefault, LongDefault}

import scala.collection.mutable

/**
  * @author <NAME>
  */
trait Optimizer {
  /** Name of this optimizer. This name is used for the accumulators created for this optimizer. */
  val name: String

  /** Boolean value indicating whether to apply use locks to prevent concurrent updates to variables. */
  val useLocking: Boolean

  /** Boolean value indicating whether to ignore duplicate indices during sparse updates. */
  val ignoreDuplicateSparseIndices: Boolean

  // TODO: [OPTIMIZER] Slot variables make re-using optimizer objects a bit dirty.

  /** Some optimizer subclasses use additional variables. For example, `MomentumOptimizer` and `AdaGradOptimizer`
    * use variables to accumulate updates. This map is where these variables are stored. */
  protected final val slots: mutable.Map[String, mutable.Map[Variable[Any], Variable[Any]]] = {
    mutable.Map.empty
  }

  /** Returns the names of all slots used by this optimizer. */
  protected final def slotNames: Set[String] = {
    slots.keySet.toSet
  }

  /** Contains variables used by some optimizers that require no slots to be stored. */
  protected final val nonSlotVariables: mutable.Map[(String, Option[Graph]), Variable[Any]] = {
    mutable.Map.empty
  }

  /** Creates an op that makes a step towards minimizing `loss` by updating the values of the variables in `variables`.
    *
    * This method simply combines calls [[computeGradients]] and [[applyGradients]]. If you want to process the
    * gradients before applying them call [[computeGradients]] and [[applyGradients]] explicitly instead of using this
    * method.
    *
    * @param  loss                       Loss value whose gradients will be computed.
    * @param  lossGradients              Optional gradients to back-propagate for `loss`.
    * @param  variables                  Optional list of variables for which to compute the gradients. Defaults to the
    *                                    set of trainable variables in the graph where `loss` is defined.
    * @param  gradientsGatingMethod      Gating method for the gradients computation.
    * @param  gradientsAggregationMethod Aggregation method used to combine gradient terms.
    * @param  colocateGradientsWithOps   Boolean value indicating whether to colocate the gradient ops with the original
    *                                    ops.
    * @param  iteration                  Optional `Variable` to increment by one after the variables have been updated.
    * @param  name                       Name for the created op.
    * @return Created op.
    */
  @throws[IllegalArgumentException]
  def minimize[T: TF : IsFloatOrDouble, I: LongDefault : TF : IsIntOrLong](
      loss: Output[T],
      lossGradients: Seq[OutputLike[T]] = null,
      variables: Set[Variable[Any]] = null,
      gradientsGatingMethod: Gradients.GatingMethod = Gradients.OpGating,
      gradientsAggregationMethod: Gradients.AggregationMethod = Gradients.AddAggregationMethod,
      colocateGradientsWithOps: Boolean = false,
      iteration: Option[Variable[I]] = None,
      name: String = "Minimize"
  ): UntypedOp = {
    val gradientsAndVariables = computeGradients(
      loss, lossGradients, variables, gradientsGatingMethod,
      gradientsAggregationMethod, colocateGradientsWithOps)
    applyGradients(gradientsAndVariables, iteration, name)
  }

  /** Computes the gradients of `loss` with respect to the variables in `variables`, if provided, otherwise with respect
    * to all the trainable variables in the graph where `loss` is defined.
    *
    * @param  loss                       Loss value whose gradients will be computed.
    * @param  lossGradients              Optional gradients to back-propagate for `loss`.
    * @param  variables                  Optional list of variables for which to compute the gradients. Defaults to the
    *                                    set of trainable variables in the graph where `loss` is defined.
    * @param  gradientsGatingMethod      Gating method for the gradients computation.
    * @param  gradientsAggregationMethod Aggregation method used to combine gradient terms.
    * @param  colocateGradientsWithOps   Boolean value indicating whether to colocate the gradient ops with the original
    *                                    ops.
    * @return Sequence of gradient-variable pairs.
    * @throws IllegalArgumentException If there are no variables to optimize.
    */
  @throws[IllegalArgumentException]
  def computeGradients[T: TF : IsFloatOrDouble](
      loss: Output[T],
      lossGradients: Seq[OutputLike[T]] = null,
      variables: Set[Variable[Any]] = null,
      gradientsGatingMethod: Gradients.GatingMethod = Gradients.OpGating,
      gradientsAggregationMethod: Gradients.AggregationMethod = Gradients.AddAggregationMethod,
      colocateGradientsWithOps: Boolean = false
  ): Seq[(OutputLike[T], Variable[Any])] = {
    // TODO: [VARIABLES] Settle on what keys to use for variables.
    val collectedVariables: Seq[Variable[Any]] = {
      {
        if (variables == null)
          loss.graph.trainableVariables
        else
          variables
      } ++ loss.graph.getCollection(Graph.Keys.STREAMING_MODEL_PORTS)
    }.toSeq
    if (collectedVariables.isEmpty)
      throw new IllegalArgumentException("There are no variables to optimize.")

    // TODO: [TYPES] !!! Super hacky. Remove in the future.
    implicit val ev: IsFloatOrDouble[Any] = null

    val variableProcessors = collectedVariables.map(v => {
      getVariableProcessor(v)(TF.fromDataType(v.dataType), ev)
    })
    val variableTargets = variableProcessors.map(_.target)
    val gradients = {
      // TODO: [TYPES] !!! What if the variable targets have different data types?
      val gradients = Gradients.gradients(
        ys = Seq(loss),
        xs = variableTargets,
        dataType = TF[T].dataType,
        dys = lossGradients,
        gateGradients = gradientsGatingMethod == Gradients.OpGating,
        aggregationMethod = gradientsAggregationMethod,
        colocateGradientsWithOps = colocateGradientsWithOps)
      if (gradientsGatingMethod == Gradients.GraphGating) {
        ControlFlow.tuple(gradients)
      } else {
        gradients
      }
    }
    gradients.zip(collectedVariables)
  }

  /** Creates an op that applies the provided gradients to the provided variables.
    *
    * @param  gradientsAndVariables Sequence with gradient-variable pairs.
    * @param  iteration             Optional `Variable` to increment by one after the variables have been updated.
    * @param  name                  Name for the created op.
    * @return Created op.
    */
  @throws[IllegalArgumentException]
  def applyGradients[T: TF, I: LongDefault : TF : IsIntOrLong](
      gradientsAndVariables: Seq[(OutputLike[T], Variable[Any])],
      iteration: Option[Variable[I]] = None,
      name: String = this.name
  ): UntypedOp = {
    // This is a default implementation of `applyGradients` that is shared by most optimizers. It relies on the subclass
    // implementing the following methods: `createSlots`, `prepare`, `finish`, `applyDense`, and `applySparse`.
    val variables: Seq[Variable[Any]] = gradientsAndVariables.filter(_._1 != null).map(_._2)
    if (variables.isEmpty) {
      throw new IllegalArgumentException(
        "No gradients were provided for any of the variables: " +
            s"${gradientsAndVariables.map(_._2).mkString(", ")}.")
    }

    Op.nameScope(name) {
      // Create the slots needed by the variables.
      Op.initializationScope {
        VariableScope.scope(name) {
          createSlots(variables)
        }
      }

      prepare(iteration)

      // Collect the update ops for all variables.
      val updateOps = mutable.Set.empty[UntypedOp]
      for ((g, v) <- gradientsAndVariables if g != null) {

        // TODO: [TYPES] !!! Super hacky. Remove in the future.
        implicit val ev: IsFloatOrDouble[Any] = null

        val p = getVariableProcessor(v)(TF.fromDataType(v.dataType), ev)
        // We colocate all ops created for variable application on the same device as the variable.
        Op.createWith(nameScope = s"Update/${v.op.name}") {
          Op.colocateWith(Set(v.op), ignoreExisting = true) {
            implicit val evRTF: TF[Any] = TF.fromDataType(v.dataType)
            val castedGradient = Cast.cast[T, Any, OutputLike](g)
            updateOps.add(p.updateOp(this, castedGradient, iteration))
          }
        }
      }

      // Create the op that applies the gradient updates to all variables.
      val applyUpdates = {
        iteration match {
          case Some(i) =>
            val finishOp = finish(updateOps.toSet, "Finish")<commit_msg>[OPS] Bug fix for optimizers.<commit_after>/* Copyright 2017-19, <NAME>. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */

package org.platanios.tensorflow.api.ops.training.optimizers

import org.platanios.tensorflow.api.core.{Graph, Shape}
import org.platanios.tensorflow.api.core.types._
import org.platanios.tensorflow.api.implicits.Implicits._
import org.platanios.tensorflow.api.ops._
import org.platanios.tensorflow.api.ops.basic.Basic
import org.platanios.tensorflow.api.ops.control_flow.ControlFlow
import org.platanios.tensorflow.api.ops.math.Math
import org.platanios.tensorflow.api.ops.training.optimizers.Optimizer._
import org.platanios.tensorflow.api.ops.variables.{ConstantInitializer, Initializer, Variable, VariableScope}
import org.platanios.tensorflow.api.tensors.Tensor
import org.platanios.tensorflow.api.utilities.DefaultsTo.{AnyDefault, LongDefault}

import scala.collection.mutable

/**
  * @author <NAME>
  */
trait Optimizer {
  /** Name of this optimizer. This name is used for the accumulators created for this optimizer. */
  val name: String

  /** Boolean value indicating whether to apply use locks to prevent concurrent updates to variables. */
  val useLocking: Boolean

  /** Boolean value indicating whether to ignore duplicate indices during sparse updates. */
  val ignoreDuplicateSparseIndices: Boolean

  // TODO: [OPTIMIZER] Slot variables make re-using optimizer objects a bit dirty.

  /** Some optimizer subclasses use additional variables. For example, `MomentumOptimizer` and `AdaGradOptimizer`
    * use variables to accumulate updates. This map is where these variables are stored. */
  protected final val slots: mutable.Map[String, mutable.Map[Variable[Any], Variable[Any]]] = {
    mutable.Map.empty
  }

  /** Returns the names of all slots used by this optimizer. */
  protected final def slotNames: Set[String] = {
    slots.keySet.toSet
  }

  /** Contains variables used by some optimizers that require no slots to be stored. */
  protected final val nonSlotVariables: mutable.Map[(String, Option[Graph]), Variable[Any]] = {
    mutable.Map.empty
  }

  /** Creates an op that makes a step towards minimizing `loss` by updating the values of the variables in `variables`.
    *
    * This method simply combines calls [[computeGradients]] and [[applyGradients]]. If you want to process the
    * gradients before applying them call [[computeGradients]] and [[applyGradients]] explicitly instead of using this
    * method.
    *
    * @param  loss                       Loss value whose gradients will be computed.
    * @param  lossGradients              Optional gradients to back-propagate for `loss`.
    * @param  variables                  Optional list of variables for which to compute the gradients. Defaults to the
    *                                    set of trainable variables in the graph where `loss` is defined.
    * @param  gradientsGatingMethod      Gating method for the gradients computation.
    * @param  gradientsAggregationMethod Aggregation method used to combine gradient terms.
    * @param  colocateGradientsWithOps   Boolean value indicating whether to colocate the gradient ops with the original
    *                                    ops.
    * @param  iteration                  Optional `Variable` to increment by one after the variables have been updated.
    * @param  name                       Name for the created op.
    * @return Created op.
    */
  @throws[IllegalArgumentException]
  def minimize[T: TF : IsFloatOrDouble, I: LongDefault : TF : IsIntOrLong](
      loss: Output[T],
      lossGradients: Seq[OutputLike[T]] = null,
      variables: Set[Variable[Any]] = null,
      gradientsGatingMethod: Gradients.GatingMethod = Gradients.OpGating,
      gradientsAggregationMethod: Gradients.AggregationMethod = Gradients.AddAggregationMethod,
      colocateGradientsWithOps: Boolean = false,
      iteration: Option[Variable[I]] = None,
      name: String = "Minimize"
  ): UntypedOp = {
    val gradientsAndVariables = computeGradients(
      loss, lossGradients, variables, gradientsGatingMethod,
      gradientsAggregationMethod, colocateGradientsWithOps)
    applyGradients(gradientsAndVariables, iteration, name)
  }

  /** Computes the gradients of `loss` with respect to the variables in `variables`, if provided, otherwise with respect
    * to all the trainable variables in the graph where `loss` is defined.
    *
    * @param  loss                       Loss value whose gradients will be computed.
    * @param  lossGradients              Optional gradients to back-propagate for `loss`.
    * @param  variables                  Optional list of variables for which to compute the gradients. Defaults to the
    *                                    set of trainable variables in the graph where `loss` is defined.
    * @param  gradientsGatingMethod      Gating method for the gradients computation.
    * @param  gradientsAggregationMethod Aggregation method used to combine gradient terms.
    * @param  colocateGradientsWithOps   Boolean value indicating whether to colocate the gradient ops with the original
    *                                    ops.
    * @return Sequence of gradient-variable pairs.
    * @throws IllegalArgumentException If there are no variables to optimize.
    */
  @throws[IllegalArgumentException]
  def computeGradients[T: TF : IsFloatOrDouble](
      loss: Output[T],
      lossGradients: Seq[OutputLike[T]] = null,
      variables: Set[Variable[Any]] = null,
      gradientsGatingMethod: Gradients.GatingMethod = Gradients.OpGating,
      gradientsAggregationMethod: Gradients.AggregationMethod = Gradients.AddAggregationMethod,
      colocateGradientsWithOps: Boolean = false
  ): Seq[(OutputLike[T], Variable[Any])] = {
    // TODO: [VARIABLES] Settle on what keys to use for variables.
    val collectedVariables: Seq[Variable[Any]] = {
      {
        if (variables == null)
          loss.graph.trainableVariables
        else
          variables
      } ++ loss.graph.getCollection(Graph.Keys.STREAMING_MODEL_PORTS)
    }.toSeq
    if (collectedVariables.isEmpty)
      throw new IllegalArgumentException("There are no variables to optimize.")

    // TODO: [TYPES] !!! Super hacky. Remove in the future.
    implicit val ev: IsFloatOrDouble[Any] = null

    val variableProcessors = collectedVariables.map(v => {
      getVariableProcessor(v)(TF.fromDataType(v.dataType), ev)
    })
    val variableTargets = variableProcessors.map(_.target)
    val gradients = {
      // TODO: [TYPES] !!! What if the variable targets have different data types?
      val gradients = Gradients.gradients(
        ys = Seq(loss),
        xs = variableTargets,
        dataType = TF[T].dataType,
        dys = lossGradients,
        gateGradients = gradientsGatingMethod == Gradients.OpGating,
        aggregationMethod = gradientsAggregationMethod,
        colocateGradientsWithOps = colocateGradientsWithOps)
      if (gradientsGatingMethod == Gradients.GraphGating) {
        ControlFlow.tuple(gradients)
      } else {
        gradients
      }
    }
    gradients.zip(collectedVariables)
  }

  /** Creates an op that applies the provided gradients to the provided variables.
    *
    * @param  gradientsAndVariables Sequence with gradient-variable pairs.
    * @param  iteration             Optional `Variable` to increment by one after the variables have been updated.
    * @param  name                  Name for the created op.
    * @return Created op.
    */
  @throws[IllegalArgumentException]
  def applyGradients[T: TF, I: LongDefault : TF : IsIntOrLong](
      gradientsAndVariables: Seq[(OutputLike[T], Variable[Any])],
      iteration: Option[Variable[I]] = None,
      name: String = this.name
  ): UntypedOp = {
    // This is a default implementation of `applyGradients` that is shared by most optimizers. It relies on the subclass
    // implementing the following methods: `createSlots`, `prepare`, `finish`, `applyDense`, and `applySparse`.
    val variables: Seq[Variable[Any]] = gradientsAndVariables.filter(_._1 != null).map(_._2)
    if (variables.isEmpty) {
      throw new IllegalArgumentException(
        "No gradients were provided for any of the variables: " +
            s"${gradientsAndVariables.map(_._2).mkString(", ")}.")
    }

    Op.nameScope(name) {
      // Create the slots needed by the variables.
      Op.initializationScope {
        VariableScope.scope(name) {
          createSlots(gradientsAndVariables.map(_._2))
        }
      }

      prepare(iteration)

      // Collect the update ops for all variables.
      val updateOps = mutable.Set.empty[UntypedOp]
      for ((g, v) <- gradientsAndVariables if g != null) {

        // TODO: [TYPES] !!! Super hacky. Remove in the future.
        implicit val ev: IsFloatOrDouble[Any] = null

        val p = getVariableProcessor(v)(TF.fromDataType(v.dataType), ev)
        // We colocate all ops created for variable application on the same device as the variable.
        Op.createWith(nameScope = s"Update/${v.op.name}") {
          Op.colocateWith(Set(v.op), ignoreExisting = true) {
            implicit val evRTF: TF[Any] = TF.fromDataType(v.dataType)
            val castedGradient = Cast.cast[T, Any, OutputLike](g)
            updateOps.add(p.updateOp(this, castedGradient, iteration))
          }
        }
      }

      // Create the op that applies the gradient updates to all variables.
      val applyUpdates = {
        iteration match {
          case Some(i) =>
            val finishOp = finish(updateOps.toSet, "Finish")
{'fraction_non_alphanumeric': 0.08590810061398296, 'fraction_numerical': 0.001336898395721925, 'mean_word_length': 2.9946598101265822, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1571054', 'n_tokens_mistral': 5507, 'n_tokens_neox': 5356, 'n_words': 1763}