Eager to support the Embedding layer and embedding_lookup method make real

[mullerhai]

HI :
I want to build recommend system use deepfm model ,but it need a lots of layer ,most important layer is Embedding layer, and so many NLP work model need it ,would you like coding this part code. will appreciate about these work .and I has try to build embedding look up method ,but it is can not work , I just try to write the similar code like [from tensorflow.python.ops import embedding_ops ]

here is my scala code ,I think maybe it will give you some idea. now only clipByNorm is work

package org.tensorflow.ops

import org.tensorflow.Output
import org.tensorflow.op.{Ops, Scope}
import org.tensorflow.types.TInt32
import org.tensorflow.types.family.TNumber

/** Method for combining sparse embeddings. */
//T: TF : IsReal, I: TF : IsIntOrLong
sealed trait Combiner {
  @inline def combine[T <:TNumber, I <: TNumber](parameters: Output[T], indices: Output[I],
                                      segmentIndices: Output[TInt32]
                           )(implicit tf:Ops,scope:Scope): Output[T]

  @inline def combineWeighted[T <:TNumber, I <: TNumber]( parameters: Output[T],weights: Output[T],
                                     segmentIndices: Output[I]
                                   )(implicit tf:Ops,scope:Scope): Output[T]
}

/** Combines sparse embeddings by using a weighted sum. */
case object SumCombiner extends Combiner {
  //T: TF : IsReal, I: TF : IsIntOrLong
  @inline override def combine[T <:TNumber, I <: TNumber](parameters: Output[T],indices: Output[I],segmentIndices: Output[TInt32]
                                    )(implicit tf:Ops,scope:Scope): Output[T] = {
//    tf.math.segmentSum()
    tf.sparse.sparseSegmentSum(parameters,indices,segmentIndices).asOutput()
//    Math.sparseSegmentSum(parameters, indices, segmentIndices)
  }

  //T: TF : IsReal, I: TF : IsIntOrLong
  @inline def combineWeighted[T <:TNumber, I <: TNumber](parameters: Output[T],weights: Output[T],segmentIndices: Output[I]
                                   )(implicit tf:Ops,scope:Scope): Output[T] = {
        tf.math.segmentSum(parameters,segmentIndices).asOutput()
//    Math.segmentSum(parameters, segmentIndices)
  }
}

/** Combines sparse embeddings by using a weighted sum divided by the total weight. */
case object MeanCombiner extends Combiner {
  //T: TF : IsReal, I: TF : IsIntOrLong
  @inline override def combine[T <:TNumber, I <: TNumber](
                                      parameters: Output[T],
                                      indices: Output[I],
                                      segmentIndices: Output[TInt32]
                                    )(implicit tf:Ops,scope:Scope): Output[T] = {
//    Math.sparseSegmentMean(parameters, indices, segmentIndices)
    tf.sparse.sparseSegmentMean(parameters,indices,segmentIndices).asOutput()
  }

  //T: TF : IsReal, I: TF : IsIntOrLong
  @inline def combineWeighted[T <:TNumber, I <: TNumber](
                                     parameters: Output[T],
                                     weights: Output[T],
                                     segmentIndices: Output[I]
                                   )(implicit tf:Ops,scope:Scope): Output[T] = {
//    val embeddings = Math.segmentSum(parameters, segmentIndices)
//    val weightsSum = Math.segmentSum(weights, segmentIndices)
//    Math.divide(embeddings, weightsSum)
    val embeddings = tf.math.segmentSum(parameters,segmentIndices)
    val weightsSum = tf.math.segmentSum(weights,segmentIndices)
    val  outRes = tf.math.divNoNan(embeddings,weightsSum)
    outRes.asOutput()

  }
}

/** Combines sparse embeddings by using a weighted sum divided by the square root of the sum of the
 * squares of the weights. */
case object SumSqrtNCombiner extends Combiner {
  //T: TF : IsReal, I: TF : IsIntOrLong
  @inline override def combine[T <:TNumber, I <: TNumber](
                                      parameters: Output[T],
                                      indices: Output[I],
                                      segmentIndices: Output[TInt32]
                                    )(implicit tf:Ops,scope:Scope): Output[T] = {

    val  segmentIds:Output[TInt32] = tf.math.segmentSum(indices.asInstanceOf[Output[TInt32]],segmentIndices).asOutput()
    val res = tf.sparse.sparseSegmentSqrtNWithNumSegments(parameters,indices,segmentIds ,segmentIndices)
    res.asOutput()
//    Math.sparseSegmentSumSqrtN(parameters, indices, segmentIndices)
//
//      tf.sparse.sparseSegmentSqrtN()
//    )
  }

  //T: TF : IsReal, I: TF : IsIntOrLong
  @inline def combineWeighted[T <:TNumber, I <: TNumber](
                                     parameters: Output[T],
                                     weights: Output[T],
                                     segmentIndices: Output[I]
                                   )(implicit tf:Ops,scope:Scope): Output[T] = {
//    val embeddings = Math.segmentSum(parameters, segmentIndices)
//    val weightsSquaredSum = Math.segmentSum(weights.square, segmentIndices)
//    Math.divide(embeddings, weightsSquaredSum.sqrt)
    val embeddings = tf.math.segmentSum(parameters,segmentIndices)
    val weightsSquaredSum = tf.math.segmentSum(tf.math.sqrt(weights),segmentIndices)
    val resOut = tf.math.divNoNan(embeddings,tf.math.sqrt(weightsSquaredSum)).asOutput()
    resOut
  }
}

package org.tensorflow.ops


import org.tensorflow.{Operand, Output}
import org.tensorflow.op.{Ops, Scope}
import org.tensorflow.types.family.TNumber
import org.tensorflow.op.math.{Div, TruncateDiv}
import org.tensorflow.types.{TFloat32, TInt32}
import org.tensorflow.ndarray.Shape

/** Partitioning strategy for the embeddings map. */
sealed trait TFPartitionStrategy {
  /** Transforms the provided ids based on this partition strategy and returns the partition assignments and the
   * new/transformed ids. */
  //T: TF : IsNotQuantized, I: TF : IsIntOrLong
  def transformIds[T<: TNumber, I <: TNumber](ids: Output[I], parameters: Seq[TFEmbeddingParameters[T]], numPartitions: Output[I]
                        )(implicit  tf :Ops,scope:Scope): (Output[I], Output[I])
}

/** Each id is assigned to partition `p = id % parameters.numPartitions`. For instance, 13 ids are split across 5
 * partitions as: `[[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]]`. */
case object ModStrategy extends TFPartitionStrategy {
  //: TF : IsNotQuantized, I: TF : IsIntOrLong
  override def transformIds[T<: TNumber, I<: TNumber](ids: Output[I], parameters: Seq[TFEmbeddingParameters[T]],
                                  numPartitions: Output[I]
                                 )(implicit  tf :Ops,scope:Scope): (Output[I], Output[I]) = {
    val partitionAssignments = tf.math.mod(ids, numPartitions)
//    val partitionAssignments = ids % numPartitions
    val newIds = TruncateDiv.create(scope,ids,numPartitions)
    (partitionAssignments.asOutput(), newIds.asOutput())
  }
}

/** Ids are assigned to partitions in a contiguous manner. In this case, 13 ids are split across 5 partitions as:
 * `[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]]`. */
case object DivStrategy extends TFPartitionStrategy {

  //T: TF : IsNotQuantized, I: TF : IsIntOrLong
  override def transformIds[T<: TNumber, I<: TNumber](ids: Output[I], parameters: Seq[TFEmbeddingParameters[T]], numPartitions: Output[I]
                                 )(implicit  tf :Ops,scope:Scope): (Output[I], Output[I]) = {
    // We compute `numTotalIds` as the sum of the first dimension size of `parameters`, and then we assign to
    // partitions based on a constant number of ids per partition. We optimize if we already know the full shape
    // statically.

    val numTotalIds = {
      if (parameters.forall(p => p.colocationOp.rank() != -1 && p.staticShape.get(0) != -1)) {
        tf.constant(parameters.map(_.staticShape.get(0)).sum).asInstanceOf[I]
//        tf.constant(parameters.map(_.staticShape(0)).toTensor.sum()).castTo[I]
      } else {
        val axis0Sizes = parameters.map(p => {
          if (p.colocationOp.rank != -1 && p.staticShape.get(0) != -1)
            tf.constant(p.staticShape.get(0))
          else
            tf.constant(p.dynamicShape(0))
//            Op.colocateWith(Set(p.colocationOp), ignoreExisting = true)(p.dynamicShape(0))
        }).toIterable
        val axisCont =axis0Sizes.asInstanceOf[java.lang.Iterable[Operand[T]]] //.asInstanceOf[I]

        tf.sum(tf.stack(axisCont),tf.constant(Array(0,1))).asOutput()
      }
    }


    val idsPerPartition = TruncateDiv.create(scope, numTotalIds.asInstanceOf[Operand[I]],numPartitions)
    val extras = tf.math.div(numTotalIds.asInstanceOf[Operand[I]], numPartitions) //: Output[I]
    val shapez =idsPerPartition.shape().asArray()
    val one = tf.ones[TInt32](tf.constant(shapez ),classOf[TInt32] )// TFBasic.ones[I](Shape())
    val tempMax1 = TruncateDiv.create(scope,ids,tf.math.add(idsPerPartition,one.asInstanceOf[Operand[I]]))
    val tempMax2 = TruncateDiv.create(scope,tf.math.sub(ids,extras),idsPerPartition)
    val partitionAssignments = tf.math.maximum( tempMax1,tempMax2).asOutput() //: Output[I]
    val newIds = tf.select(
      tf.math.less(partitionAssignments, extras),
      tf.math.div(ids, tf.math.add(idsPerPartition,one.asInstanceOf[Operand[I]])),
      tf.math.div(tf.math.sub(ids,extras),idsPerPartition))
    (partitionAssignments.asInstanceOf[Output[I]], newIds.asOutput().asInstanceOf[Output[I]])
  }
}

package org.tensorflow.ops

import org.tensorflow.types.TInt32
import org.tensorflow.ndarray.Shape
import org.tensorflow.op.Ops
import org.tensorflow.{Operand, Output}
import org.tensorflow.types.family.TNumber
import org.tensorflow.op.core.Variable


case class TFEmbeddingMap[T<:TNumber](partitionParameters: Seq[TFEmbeddingParameters[T]]) {
  val numPartitions: Int = partitionParameters.size
}
/** Trait for specifying supported embedding parameter types. */
trait TFEmbeddingParameters[T <:TNumber] {
  /** Returns the op that generates these parameters (to be used for colocating other ops with it). */
  @inline def colocationOp: Operand[T]

  /** Returns the static shape of this parameters tensor. */
  @inline def staticShape: Shape

  /** Returns the dynamic shape of this parameters tensor. */
  @inline def dynamicShape: Array[Long]

  /** Gathers the embeddings corresponding to `indices` from `parameters`. */
  def gather[I<:TNumber](
                 indices: Output[I],
                 name: String = "Gather"
               ): Output[T]
}


//: TF : IsNotQuantized
case class OutputParameters[T<:TNumber](tf: Ops,parameters: Output[T] )extends TFEmbeddingParameters[T] {
  @inline override def colocationOp: Operand[T] = {
    parameters.op().asInstanceOf[Output[T]]
  }

  @inline override def staticShape: Shape = {
    parameters.shape
  }

  @inline override def dynamicShape: Array[Long] = {
//    tf.shape()
//    TFBasic.shape(parameters)
    parameters.shape().asArray()
  }

  //: TF : IsIntOrLong
  override def gather[I<:TNumber](indices: Output[I], name: String = "Gather"): Output[T] = {
    tf.gather(parameters,indices,tf.constant(0)).asOutput()
//    TFBasic.gather(parameters, indices, axis = 0, name = name)
  }
}


case class VariableParameters[T<:TNumber](tf: Ops,parameters: Variable[T]) extends TFEmbeddingParameters[T] {
  @inline override def colocationOp:  Operand[T] = {
    parameters.op().asInstanceOf[Variable[T]]
  }

  @inline override def staticShape: Shape = {
    parameters.shape
  }

  @inline override def dynamicShape: Array[Long] = {
//    TFBasic.shape(parameters.value)
    parameters.shape().asArray()
  }

  //IsIntOrLong
  override def gather[I<:TNumber](indices: Output[I], name: String = "Gather"): Output[T] = {
    tf.gather(parameters,indices,tf.constant(0)).asOutput()
//    parameters.gather(indices, name = name)
  }
}

package org.tensorflow.ops


import org.tensorflow.{Operand, Output}
import org.tensorflow.ndarray.Shape
import org.tensorflow.op.core.ReduceSum
import org.tensorflow.op.{Ops, Scope}
import org.tensorflow.types.{TFloat64, TInt32}
import org.tensorflow.types.family.TNumber
import scala.language.postfixOps

/** Contains functions for constructing ops related to embeddings.
 *
 * @author Emmanouil Antonios Platanios
 */
trait TFEmbedding {

  def clipByNorm[T <: TNumber](someNums: Operand[T], clipNormBias: Operand[T], axis: Operand[TInt32])(implicit tf: Ops, scope: Scope): Output[T] = {
    try {
      //      val someNums =  tf.withName("c").constant(Array(Array(1.0,2.0,3.0,4.4,5.2)))
      //      val clipNormBias = tf.withName("d").constant(2.0) //2.0
      //      val axis = tf.constant(Array(1,0))
      //      tf.array(0,1)
      val l2sum = tf.reduceSum(tf.math.square(tf.dtypes.cast(someNums, classOf[TFloat64])), axis, ReduceSum.keepDims(true))
      val pred = tf.math.greater(l2sum, tf.constant(0.0))
      val l2sumSafe = tf.select(pred, l2sum, tf.onesLike(l2sum))
      val l2norm = tf.select(pred, tf.math.sqrt(l2sumSafe), l2sum)
      val intermediate = tf.math.mul(tf.dtypes.cast(someNums, classOf[TFloat64]), tf.dtypes.cast(clipNormBias, classOf[TFloat64]))
      val valuesClip = tf.identity(tf.math.divNoNan(intermediate, tf.math.maximum(l2norm, tf.dtypes.cast(clipNormBias, classOf[TFloat64]))))
      if (someNums.shape().isCompatibleWith(intermediate.shape())) {
        println(s"match the shape ${someNums.shape().get(0)} ${someNums.shape().get(1)}  || ${intermediate.shape().get(0)} ${intermediate.shape().get(1)}")
        valuesClip.asOutput().asInstanceOf[Output[T]]  //Float64
      } else {
        println(s"can not match the shape ${someNums.shape().get(0)} ${someNums.shape().get(1)}  || ${intermediate.shape().get(0)} ${intermediate.shape().get(1)}")
        throw new Exception("bad shape")
      }
    }
  }

  /** $OpDocEmbeddingEmbeddingLookup
   *
   * @group EmbeddingOps
   * @param parameters        Embedding map, which is either a single tensor, a list of `P` tensors with the same
   *                          shape, except for their first dimension, representing sharded embedding tensors, or a
   *                          `PartitionedVariable`, created by partitioning along the first dimension.
   * @param ids               Tensor to be looked up in `parameters`.
   * @param partitionStrategy Partitioning strategy to use if `parameters.numPartitions > 1`.
   * @param transformFn       If provided, this function is applied to each partitioned tensor of retrieved
   *                          embeddings, colocated with the embeddings. The shape of the argument to this function
   *                          will be the same as that of `parameters`, except for the size of the first dimension.
   *                          The first dimension of the result's shape must have the same size as that of the
   *                          argument's. Note that, if `maxNorm` is provided, then norm-based clipping is performed
   *                          before the `transformFn` is applied.
   * @param maxNorm           If provided, embedding values are l2-normalized to this value.
   * @param name              Name prefix used for the created op.
   * @return Obtained embeddings for the provided `ids`.
   */
  //T: TF : IsNotQuantized, I: TF : IsIntOrLong
  def embeddingLookup[T <: TNumber, I <: TNumber](parameters: TFEmbeddingMap[T], ids: Output[I], partitionStrategy: TFPartitionStrategy = ModStrategy,
                                                  transformFn: Output[T] => Output[T] = null,
                                                  maxNorm: Operand[T] = null,
                                                  name: String = "EmbeddingLookup"
                                                 )(implicit tf: Ops, scope: Scope): Output[TFloat64] = {
    tf.withName(name)
    //    Op.nameScope(name) {
    if (parameters.numPartitions == 1 && (ids.rank == 1 || transformFn == null)) {
      var result = parameters.partitionParameters.head.gather(ids)
      if (maxNorm != null) {
        val axis = tf.dtypes.cast(tf.constant(Array(1, 0)), classOf[TInt32])
        result = TFEmbedding.clipByNorm(result, maxNorm, axis) //.asInstanceOf[Output[T]]
      }
      if (transformFn != null)
        result = transformFn(result)
      result
    } else {

      val numPartitions = tf.constant(parameters.numPartitions) //.asInstanceOf[I]
      // val numPartitions = TFBasic.constant(parameters.numPartitions).castTo[I]
      // Flatten the ids. There are two cases where we need to do this:
      //   - There is more than one parameter tensors.
      //   - There is a `transformFn` and ids is not statically known to be 1-D.
      // In this case, we must flatten because `transformFn` expects a flat tensor of embeddings.
      val flattenedIds = tf.reshape(ids, tf.constant(Array(-1)))
      //        val flattenedIds = ids.reshape(Shape(-1))
      val originalIds = tf.range(tf.constant(0), tf.constant(flattenedIds.size()), tf.constant(1))
      //        val originalIds = Math.range(0, flattenedIds.size.toInt)
      // Create `partitionAssignments` and set `newIds` depending on the strategy.
      val transformedIds = partitionStrategy.transformIds(
        flattenedIds.output(), parameters.partitionParameters, numPartitions.asOutput())
      // Cast partition assignments to integers for use in `dynamicPartition`.
      // There really should not be more than 2^32 partitions.
      //        val partitionAssignments =  transformedIds._1.castTo[Int]
      val partitionAssignments = transformedIds._1.asInstanceOf[Output[TInt32]]
      val newIds = transformedIds._2

      // Partition list of ids based on assignments into `parameters.numPartitions` separate lists.
      val gatherIds = tf.dynamicPartition(
        newIds, partitionAssignments, parameters.numPartitions)
      // Similarly, partition the original indices.
      val partitionIndices = tf.dynamicPartition(
        originalIds, partitionAssignments, parameters.numPartitions)

      val gatherIdsList = gatherIds.outputs().asInstanceOf[Iterable[Output[I]]]
      // Do `parameters.numPartitions` separate lookups, finding embeddings for `plist(p)` in `parameters(p)`.
      val partitionedResult = parameters.partitionParameters.zip(gatherIdsList).map {
        case (params, paramIds: Output[I]) =>
          //            Op.colocateWith(Set(params.colocationOp), ignoreExisting = true) {
          //          }
          var result = params.gather(paramIds)
          // If `transformFn` is provided, the `clipByNorm` precedes the transform and hence must be co-located.
          // See below for the counterpart if `transformFn` is not provided.
          val axis2 = tf.dtypes.cast(tf.constant(Array(1, 0)), classOf[TInt32])
          if (maxNorm != null) {
            //result = TFEmbedding.clipByNorm(result, paramIds, maxNorm)
            result = TFEmbedding.clipByNorm(paramIds, maxNorm, axis2).asInstanceOf[Output[T]]
          }
          if (transformFn != null && maxNorm != null)
          //              result = transformFn(TFEmbedding.clipByNorm(result, paramIds, maxNorm))
            result = transformFn(TFEmbedding.clipByNorm(paramIds, maxNorm, axis2).asInstanceOf[Output[T]])
          else if (transformFn != null)
            result = transformFn(result)
          result
      }.asInstanceOf[java.lang.Iterable[Operand[T]]]

      // Stitch these back together.
      val partIndices = partitionIndices.outputs().asInstanceOf[java.lang.Iterable[Operand[TInt32]]]
      var result = tf.dynamicStitch(partIndices, partitionedResult)

      // Determine the static element shape.
      val elementStaticShape = {
        if (transformFn == null) {
          val shapeArr = parameters.partitionParameters.head.staticShape.asArray()
          var shape = shapeArr.slice(1, shapeArr.length) //.t(1 ::)
//          parameters.partitionParameters.tail.foreach(p => {
//            shape = shape.mergeWith(p.staticShape(1 ::))
//
//          })
          shape
        } else {
          result.shape.asArray().slice(1, result.shape.asArray().length) //(1 ::)
        }
      }

      // Compute the dynamic element shape.
      val elementDynamicShape = {
        if (elementStaticShape.length>1) {
           tf.constant(elementStaticShape )
        } else if (transformFn == null) {
          parameters.partitionParameters.head.dynamicShape.slice(1,parameters.partitionParameters.head.dynamicShape.length)   //(1 ::)
        } else {
          tf.shape(result).shape().asArray().slice(1,tf.shape(result).shape().asArray().length)  //.slice(1 ::)
        }
      }

      // Reshape to reverse the flattening of the ids.

      val idsShape = tf.shape(ids)
      //        val resultShape = tf.concatenate(Seq(idsShape, elementDynamicShape))
      //        result = result.reshape(resultShape)
      val resultShape = tf.concat(Seq(idsShape, elementDynamicShape).asInstanceOf[java.lang.Iterable[Operand[I]]], tf.constant(Array(1, 0)))
      val resultz = tf.reshape(result, resultShape)
      // Normally the reshape is sufficient, but setting shape explicitly teaches shape inference that
      // `parameters.partitionParameters(1 ::).shape` matters (in the case that `transformFn` is `null`).

      //        tf.math.squaredDifference()

      val eleStaticShape = Shape.of(elementStaticShape: _*)
      ids.shape().append(eleStaticShape)
      var resultzs = tf.reshape(resultz, tf.constant(ids.shape().asArray())).asOutput()
      //        resultz.setShape(ids.shape.concatenateWith(elementStaticShape))
      val axis3 = tf.dtypes.cast(tf.constant(Array(1, 0)), classOf[TInt32])
      if (transformFn == null) {
        // If `transformFn` was provided, the `clipByNorm` was done above.
        //          resultz = TFEmbedding.clipByNorm(result, ids, maxNorm)
        return TFEmbedding.clipByNorm(ids, maxNorm, axis3)

      }else{
        return resultzs.asInstanceOf[Output[TFloat64]]
      }


    }
    //    }
  }

  /** $OpDocEmbeddingSparseEmbeddingLookup
   *
   * @group EmbeddingOps
   * @param parameters        Embedding map, which is either a single tensor, a list of `P` tensors with the same
   *                          shape, except for their first dimension, representing sharded embedding tensors, or a
   *                          `PartitionedVariable`, created by partitioning along the first dimension.
   * @param sparseIds         `NxM` sparse tensor containing ids, where `N` typically corresponds to the batch
   *                          size and `M` is arbitrary.
   * @param sparseWeights     Either a sparse tensor containing weight values, or None `null`
   *                          to indicate all weights should be taken to be equal to 1. If specified, `sparseWeights`
   *                          must have exactly the same shape and indices as `sparseIds`.
   * @param partitionStrategy Partitioning strategy to use if `parameters.numPartitions > 1`.
   * @param combiner          Combination/reduction strategy to use for the obtained embeddings.
   * @param maxNorm           If provided, embedding values are l2-normalized to this value.
   * @param name              Name prefix used for the created op.
   * @return Obtained embeddings for the provided `ids`.
   */
  //T: TF : IsReal, I: TF : IsIntOrLong
  //  def sparseEmbeddingLookup[T <:TNumber, I <:TNumber](parameters: TFEmbeddingMap[T], sparseIdz: Output[I],
  //                                  sparseWeightz: Output[T] = null,
  //                                  partitionStrategy: TFPartitionStrategy = ModStrategy,
  //                                  combiner: Combiner = SumSqrtNCombiner,
  //                                  maxNorm: Output[T] = null,
  //                                  name: String = "SparseEmbeddingLookup"
  //                                 )(implicit tf:Ops,scope: Scope): Output[T] = {
  //
  //    val ignoreWeights = sparseWeightz == null
  //    if (!ignoreWeights) {
  //      val sparseWeights = sparseWeightz.asTensor().asInstanceOf[SparseTensor]
  //      val sparseIds = sparseIdz.asTensor().asInstanceOf[SparseTensor]
  //       val idFlag = sparseIds.indices.shape.isCompatibleWith(sparseWeights.indices.shape)
  //       val valueFlag = sparseIds.values.shape.isCompatibleWith(sparseWeights.values.shape)
  ////      sparseIds.indices.shape.assertIsCompatibleWith(sparseWeights.indices.shape)
  ////      sparseIds.values.shape.assertIsCompatibleWith(sparseWeights.values.shape)
  //      if (idFlag==true &&valueFlag==true&& sparseIds.denseShape != null && sparseWeights.denseShape != null) {
  ////        sparseIds.denseShape.shape.assertIsCompatibleWith(sparseWeights.denseShape.shape)
  //        val spareFlag = sparseIds.denseShape.shape.isCompatibleWith(sparseWeights.denseShape.shape)
  ////        val segmentIds = sparseIds.indices(::, 0).castTo[Int]
  //        val segmentIds = tf.dtypes.cast(sparseIds.indices(::, 0), classOf[TInt64])
  //        val (ids, idx) = tf.unique(sparseIds.values,tf.constant(Array(0)))
  //        var embeddings = embeddingLookup(parameters, ids, partitionStrategy, maxNorm = maxNorm)
  //        if (ignoreWeights) {
  //          combiner.combine(embeddings, idx, segmentIds)
  //        } else {
  //          embeddings = embeddings.gather(idx)
  //          val weights = sparseWeights.values.castTo[T]
  //          // Reshape weights to allow broadcasting.
  //          val weightsStaticShape = weights.shape
  //          val weightsDynamicShape = tf.shape(weights)
  //          val ones = Output.ones[Int](tf.expandDims(tf.rank(embeddings) - 1, 0))
  //          val broadcastedWeightsShape = tf.concatenate(
  //            Seq(weightsDynamicShape, ones), axis = 0)
  //          val reshapedWeights = weights.reshape(broadcastedWeightsShape)
  //          // Set the weight shape, since after reshaping to `broadcastedWeightsShape`, the shape becomes unknown.
  //          if (embeddings.shape.rank != -1) {
  //            val onesShape = Shape.fromSeq((0 until embeddings.shape.rank - 1).map(_ => 1))
  //            reshapedWeights.setShape(weightsStaticShape.concatenateWith(onesShape))
  //          }
  //          val weightedEmbeddings = embeddings * reshapedWeights
  //          combiner.combineWeighted(weightedEmbeddings, reshapedWeights, segmentIds)
  //        }
  //
  //      } else if (sparseIds.denseShape != sparseWeights.denseShape) {
  //        throw  new  Exception("The dense shapes of 'sparseIds' and 'sparseWeights' must be compatible.")
  //      }
  //    }
  //  }

}

object TFEmbedding extends TFEmbedding {


  /** If `maxNorm` is not `null`, this method clips `parameters` to a maximum l2-norm of `maxNorm`. */
  //: TF : IsNotQuantized  : TF : IsIntOrLong
  //  private[TFEmbedding] def clipByNorm[T <:TNumber, I <:TNumber](parameters: Output[T], indices: Output[I], maxNorm: Output[T] = null, name: String = "ClipNorm"): Output[T] = {
  //    if (maxNorm == null) {
  //      parameters
  //    } else if (parameters.rank != -1 && indices.rank != -1) {
  //      parameters.clipByNorm(
  //        maxNorm,
  //        indices.rank until parameters.rank)
  //
  //      input: Output[T],
  //      clipNorm: Output[T],
  //      axes: Output[I] = null,
  //      name: String = "ClipByNorm"
  //      Clip.clipByNorm(output, clipNorm, axes, name)
  //
  //      // Calculate the l2-norm and clip elements by the ratio of `clipNorm` to that l2-norm.
  //      val l2Norm = Math.sum(Math.square(input), axes, keepDims = true).sqrt
  //      val intermediate = input * clipNorm
  //      // Assert that the result shape is compatible with the initial shape, to prevent unintentional broadcasting.
  //      input.shape.assertIsCompatibleWith(intermediate.shape)
  //      TFBasic.identity(intermediate / Math.maximum(l2Norm, clipNorm))
  //
  //    } else {
  //      parameters.clipByNorm(
  //        maxNorm,
  //        Math.range(TFBasic.rank(indices), TFBasic.rank(parameters)))
  //    }
  //  }


  //: TF : IsNotQuantized


  /** @define OpDocEmbeddingEmbeddingLookup
   *   The `embeddingLookup` op looks up `ids` in a list of embedding tensors.
   *
   *   This function is used to perform parallel lookups on the embedding map in `parameters`. It is a generalization
   *   of the `gather` op, where `parameters` is interpreted as a partitioning of a large embedding tensor.
   *   `parameters` may be a `PartitionedVariable` as returned when creating a variable with a partitioner.
   *
   *   If `parameters` consists of more than 1 partition, each element `id` of `ids` is partitioned between the
   *   elements of `parameters` according to the `partitionStrategy`. In all strategies, if the id space does not
   *   evenly divide the number of partitions, each of the first `(maxId + 1) % parameters.numPartitions` partitions
   *   will be assigned one more id.
   *
   *   If `partitionStrategy` is [[TFEmbedding.ModStrategy]], we assign each id to partition
   *   `p = id % parameters.numPartitions`. For instance, 13 ids are split across 5 partitions as:
   *   `[[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]]`.
   *
   *   If `partitionStrategy` is [[TFEmbedding.DivStrategy]], we assign ids to partitions in a contiguous manner. In this
   *   case, 13 ids are split across 5 partitions as:
   *   `[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]]`.
   *
   *   The results of the lookup are concatenated into a dense tensor. The returned tensor has shape
   *   `ids.shape + parameters.partitionParameters(0).shape(1 ::)`.
   * @define OpDocEmbeddingSparseEmbeddingLookup
   *   The `sparseEmbeddingLookup` op looks up and computes embeddings for the given sparse ids and weights.
   *
   *   The op assumes that there is at least one id for each row in the dense tensor represented by `sparseIds` (i.e.,
   *   there are no rows with empty features), and that all the indices of `sparseIds` are in canonical row-major
   *   order. It also assumes that all id values lie in the range `[0, p0)`, where `p0` is the sum of the size of
   *   `parameters` along dimension 0.
   *
   *   The op returns a dense tensor representing the combined embeddings for the provided sparse ids. For each row in
   *   the dense tensor represented by `sparseIds`, the op looks up the embeddings for all ids in that row, multiplies
   *   them by the corresponding weight, and combines them using the provided `combiner`.
   *
   *   In other words, if `shape(combinedParameters) = [p0, p1, ..., pm]` and
   *   `shape(sparseIds) = shape(sparseWeights) = [d0, d1, ..., dn]`, then
   *   `shape(output) = [d0, d1, ..., dn-1, p1, ..., pm]`.
   *
   *   For instance, if `parameters` is a `10x20` matrix, and `sparseIds` and `sparseWeights` are as follows:
   *
   *     - [0, 0]: id 1, weight 2.0
   *     - [0, 1]: id 3, weight 0.5
   *     - [1, 0]: id 0, weight 1.0
   *     - [2, 3]: id 1, weight 3.0
   *
   *   and we are using the `MeanCombiner`, then the output will be a `3x20` matrix, where:
   *
   *     - output(0, ::) = (parameters(1, ::) * 2.0 + parameters(3, ::) * 0.5) / (2.0 + 0.5)
   *     - output(1, ::) = parameters(0, ::) * 1.0
   *     - output(2, ::) = parameters(1, ::) * 3.0
   */
  //  private[ops] trait Documentation
}

here is use python keras api to build the DeepFM model

from keras.layers import *
from keras import backend as K
from keras.models import Model
from HwSumLayer import HwSumLayer
from HwMeanPooling import HwMeanPooling
from HwFlatten import HwFlatten
import tensorflow as tf
from tensorflow.python.ops import  data_flow_ops
from tensorflow.python.ops import embedding_ops,array_ops
# from keras.utils import
# from keras.layers.embeddings import

in_score = Input(shape=[1],name ="score")
in_sales = Input(shape=[1],name="sales")


# tf.realdiv
# tf.truncatediv
# data_flow_ops.
tf.range(tf.shape(x)[0])
in_gender = Input(shape=[1],name="gender")
in_age = Input(shape=[1],name="age")

in_interest = Input(shape=[3],name="interest")
in_topic = Input(shape=[4],name="topic")

numeric = Concatenate()([in_score,in_sales])
dense_numeric = Dense(1)(numeric)

emb_gender_id = Reshape([1])(Embedding(3,1)(in_gender))
emb_age_1d = Reshape([1])(Embedding(3,1)(in_age))

emb_interest_1d =Embedding(11,1,mask_zero=True)(in_interest)
emb_interest_1d = HwMeanPooling(axis=1)(emb_interest_1d)

emb_topic_1d = Embedding(22,1,mask_zero=True)(in_topic)
emb_topic_1d = HwMeanPooling(axis=1)(emb_topic_1d)

y_first_order  = Add()([dense_numeric,emb_gender_id,emb_age_1d,emb_interest_1d,emb_topic_1d])

latent =8

emb_score_Kd =RepeatVector(1)(Dense(latent)(in_score))

emb_sales_Kd = RepeatVector(1)(Dense(latent)(in_sales))

emb_gender_Kd = Embedding(3,latent)(in_gender)
emb_age_Kd= Embedding(10,latent)(in_age)

emb_interest_Kd = Embedding(11,latent,mask_zero=True)(in_interest)
emb_interest_Kd = RepeatVector(1)(HwMeanPooling(axis=1)(emb_interest_Kd))

emb_topic_Kd = Embedding(22,latent,mask_zero=True)(in_topic)
emb_topic_Kd = RepeatVector(1)(HwMeanPooling(axis=1)(emb_topic_Kd))

emb = Concatenate(axis=1)([emb_score_Kd,emb_sales_Kd,emb_gender_Kd,emb_age_Kd,
                           emb_interest_Kd,emb_topic_Kd])
summed_features_emb = HwSumLayer(axis=1)(emb)

summed_features_emb_square= Multiply()([summed_features_emb,summed_features_emb])

squared_features_emb = Multiply()([emb,emb])
squared_sum_features_emb = HwSumLayer(axis=1)(squared_features_emb)

sub = Subtract()([summed_features_emb_square,squared_sum_features_emb])
sub = Lambda(lambda  x: x*0.5)(sub)

y_second_order = HwSumLayer(axis=1)(sub)

y_deep = HwFlatten()(emb)
y_deep = Dropout(0.5)(Dense(128,activation='relu')(y_deep))
y_deep = Dropout(0.5)(Dense(64,activation='relu')(y_deep))
y_deep = Dropout(0.5)(Dense(32,activation='relu')(y_deep))
y_deep = Dropout(0.5)(Dense(1,activation='relu')(y_deep))
y = Concatenate(axis=1)([y_first_order,y_second_order,y_deep])
y = Dense(1,activation='sigmoid')(y)

model = Model(inputs=[in_score,in_sales,in_gender,in_age,in_interest,in_topic],outputs=[y])
model.compile(optimizer="rmsProp",loss="logloss")

model.compile(optimizer="adam",
              loss="binary_crossentropy",
              metrics=["binary_crossentropy", tf.keras.metrics.AUC(name='auc')])

train_data = total_data.loc[:500000-1]
valid_data = total_data.loc[500000:]

train_dense_x = [train_data[f].values for f in dense_feats]
train_sparse_x = [train_data[f].values for f in sparse_feats]
train_label = [train_data['label'].values]

val_dense_x = [valid_data[f].values for f in dense_feats]
val_sparse_x = [valid_data[f].values for f in sparse_feats]
val_label = [valid_data['label'].values]

model.fit(train_dense_x+train_sparse_x,
          train_label, epochs=5, batch_size=256,
          validation_data=(val_dense_x+val_sparse_x, val_label),
         )