diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala b/scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala index f9f2dbe42a90..837125908674 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala @@ -40,6 +40,8 @@ object NDArray extends NDArrayBase { val api = NDArrayAPI + val random = NDArrayRandomAPI + private def addDependency(froms: Array[NDArray], tos: Array[NDArray]): Unit = { froms.foreach { from => val weakRef = new WeakReference(from) diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayAPI.scala b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayAPI.scala index 1d8551c1b1e5..cb507877520d 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayAPI.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayAPI.scala @@ -23,3 +23,13 @@ package org.apache.mxnet object NDArrayAPI extends NDArrayAPIBase { // TODO: Implement CustomOp for NDArray } + +@AddNDArrayRandomAPIs(false) +/** + * typesafe NDArray random module: NDArray.random._ + * Main code will be generated during compile time through Macros + */ +object NDArrayRandomAPI extends NDArrayRandomAPIBase { + +} + diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/Symbol.scala b/scala-package/core/src/main/scala/org/apache/mxnet/Symbol.scala index 4472a8426f9f..17f3636c2e27 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/Symbol.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/Symbol.scala @@ -840,8 +840,12 @@ object Symbol extends SymbolBase { private val functions: Map[String, SymbolFunction] = initSymbolModule() private val bindReqMap = Map("null" -> 0, "write" -> 1, "add" -> 3) + type SymbolOrFloat = Any + val api = SymbolAPI + val random = SymbolRandomAPI + def pow(sym1: Symbol, sym2: Symbol): Symbol = { Symbol.createFromListedSymbols("_Power")(Array(sym1, sym2)) } diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/SymbolAPI.scala b/scala-package/core/src/main/scala/org/apache/mxnet/SymbolAPI.scala index 1bfb0559cf96..e5b55f6e1eec 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/SymbolAPI.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/SymbolAPI.scala @@ -32,3 +32,13 @@ object SymbolAPI extends SymbolAPIBase { Symbol.createSymbolGeneral("Custom", name, attr, Seq(), map.toMap) } } + +@AddSymbolRandomAPIs(false) +/** + * typesafe Symbol random module: Symbol.random._ + * Main code will be generated during compile time through Macros + */ +object SymbolRandomAPI extends SymbolRandomAPIBase { + +} + diff --git a/scala-package/core/src/test/scala/org/apache/mxnet/NDArraySuite.scala b/scala-package/core/src/test/scala/org/apache/mxnet/NDArraySuite.scala index 5d88bb39e502..f90487739f34 100644 --- a/scala-package/core/src/test/scala/org/apache/mxnet/NDArraySuite.scala +++ b/scala-package/core/src/test/scala/org/apache/mxnet/NDArraySuite.scala @@ -576,4 +576,22 @@ class NDArraySuite extends FunSuite with BeforeAndAfterAll with Matchers { assert(arr.internal.toDoubleArray === Array(2d, 2d)) assert(arr.internal.toByteArray === Array(2.toByte, 2.toByte)) } + + test("random module is generated properly") { + val lam = NDArray.ones(1, 2) + val rnd = NDArray.random.poisson(lam = Some(lam), shape = Some(Shape(3, 4))) + val rnd2 = NDArray.random.poisson(lam = Some(1f), shape = Some(Shape(3, 4))) + assert(rnd.shape === Shape(1, 2, 3, 4)) + assert(rnd2.shape === Shape(3, 4)) + } + + test("random module is generated properly - special case of 'normal'") { + val mu = NDArray.ones(1, 2) + val sigma = NDArray.ones(1, 2) * 2 + val rnd = NDArray.random.normal(mu = Some(mu), sigma = Some(sigma), shape = Some(Shape(3, 4))) + val rnd2 = NDArray.random.normal(mu = Some(1f), sigma = Some(2f), + shape = Some(Shape(3, 4))) + assert(rnd.shape === Shape(1, 2, 3, 4)) + assert(rnd2.shape === Shape(3, 4)) + } } diff --git a/scala-package/core/src/test/scala/org/apache/mxnet/SymbolSuite.scala b/scala-package/core/src/test/scala/org/apache/mxnet/SymbolSuite.scala index ebb61d7d4bfb..4ca8a4564452 100644 --- a/scala-package/core/src/test/scala/org/apache/mxnet/SymbolSuite.scala +++ b/scala-package/core/src/test/scala/org/apache/mxnet/SymbolSuite.scala @@ -19,7 +19,9 @@ package org.apache.mxnet import org.scalatest.{BeforeAndAfterAll, FunSuite} + class SymbolSuite extends FunSuite with BeforeAndAfterAll { + test("symbol compose") { val data = Symbol.Variable("data") @@ -71,4 +73,27 @@ class SymbolSuite extends FunSuite with BeforeAndAfterAll { val data2 = data.clone() assert(data.toJson === data2.toJson) } + + test("random module is generated properly") { + val lam = Symbol.Variable("lam") + val rnd = Symbol.random.poisson(lam = Some(lam), shape = Some(Shape(2, 2))) + val rnd2 = Symbol.random.poisson(lam = Some(1f), shape = Some(Shape(2, 2))) + // scalastyle:off println + println(s"Symbol.random.poisson debug info: ${rnd.debugStr}") + println(s"Symbol.random.poisson debug info: ${rnd2.debugStr}") + // scalastyle:on println + } + + test("random module is generated properly - special case of 'normal'") { + val loc = Symbol.Variable("loc") + val scale = Symbol.Variable("scale") + val rnd = Symbol.random.normal(mu = Some(loc), sigma = Some(scale), + shape = Some(Shape(2, 2))) + val rnd2 = Symbol.random.normal(mu = Some(1f), sigma = Some(2f), + shape = Some(Shape(2, 2))) + // scalastyle:off println + println(s"Symbol.random.sample_normal debug info: ${rnd.debugStr}") + println(s"Symbol.random.random_normal debug info: ${rnd2.debugStr}") + // scalastyle:on println + } } diff --git a/scala-package/macros/src/main/scala/org/apache/mxnet/APIDocGenerator.scala b/scala-package/macros/src/main/scala/org/apache/mxnet/APIDocGenerator.scala index b4efa659443c..8dbab4a984de 100644 --- a/scala-package/macros/src/main/scala/org/apache/mxnet/APIDocGenerator.scala +++ b/scala-package/macros/src/main/scala/org/apache/mxnet/APIDocGenerator.scala @@ -17,8 +17,6 @@ package org.apache.mxnet -import org.apache.mxnet.init.Base._ -import org.apache.mxnet.utils.CToScalaUtils import java.io._ import java.security.MessageDigest @@ -29,77 +27,115 @@ import scala.collection.mutable.{ArrayBuffer, ListBuffer} * Two file namely: SymbolAPIBase.scala and NDArrayAPIBase.scala * The code will be executed during Macros stage and file live in Core stage */ -private[mxnet] object APIDocGenerator{ - case class absClassArg(argName : String, argType : String, argDesc : String, isOptional : Boolean) - case class absClassFunction(name : String, desc : String, - listOfArgs: List[absClassArg], returnType : String) +private[mxnet] object APIDocGenerator extends GeneratorBase { + type absClassArg = Arg + type absClassFunction = Func - - def main(args: Array[String]) : Unit = { + def main(args: Array[String]): Unit = { val FILE_PATH = args(0) val hashCollector = ListBuffer[String]() hashCollector += absClassGen(FILE_PATH, true) hashCollector += absClassGen(FILE_PATH, false) + hashCollector += absRndClassGen(FILE_PATH, true) + hashCollector += absRndClassGen(FILE_PATH, false) hashCollector += nonTypeSafeClassGen(FILE_PATH, true) hashCollector += nonTypeSafeClassGen(FILE_PATH, false) val finalHash = hashCollector.mkString("\n") } - def MD5Generator(input : String) : String = { + def MD5Generator(input: String): String = { val md = MessageDigest.getInstance("MD5") md.update(input.getBytes("UTF-8")) val digest = md.digest() org.apache.commons.codec.binary.Base64.encodeBase64URLSafeString(digest) } - def absClassGen(FILE_PATH : String, isSymbol : Boolean) : String = { - // scalastyle:off - val absClassFunctions = getSymbolNDArrayMethods(isSymbol) - // Defines Operators that should not generated + def absRndClassGen(FILE_PATH: String, isSymbol: Boolean): String = { + val funcs = buildRandomFunctionList(isSymbol) + + val body = funcs.map(func => { + val scalaDoc = generateAPIDocFromBackend(func) + val decl = generateAPISignature(func, isSymbol) + s"$scalaDoc\n$decl" + }) + writeFile( + FILE_PATH, + if (isSymbol) "SymbolRandomAPIBase" else "NDArrayRandomAPIBase", + body) + } + + def absClassGen(FILE_PATH: String, isSymbol: Boolean): String = { val notGenerated = Set("Custom") - // TODO: Add Filter to the same location in case of refactor - val absFuncs = absClassFunctions.filterNot(_.name.startsWith("_")) + val funcs = buildFunctionList(isSymbol) + .filterNot(_.name.startsWith("_")) .filterNot(ele => notGenerated.contains(ele.name)) - .map(absClassFunction => { - val scalaDoc = generateAPIDocFromBackend(absClassFunction) - val defBody = generateAPISignature(absClassFunction, isSymbol) - s"$scalaDoc\n$defBody" + val body = funcs.map(func => { + val scalaDoc = generateAPIDocFromBackend(func) + val decl = generateAPISignature(func, isSymbol) + s"$scalaDoc\n$decl" }) - val packageName = if (isSymbol) "SymbolAPIBase" else "NDArrayAPIBase" - val apacheLicence = "/*\n* Licensed to the Apache Software Foundation (ASF) under one or more\n* contributor license agreements. See the NOTICE file distributed with\n* this work for additional information regarding copyright ownership.\n* The ASF licenses this file to You under the Apache License, Version 2.0\n* (the \"License\"); you may not use this file except in compliance with\n* the License. You may obtain a copy of the License at\n*\n* http://www.apache.org/licenses/LICENSE-2.0\n*\n* Unless required by applicable law or agreed to in writing, software\n* distributed under the License is distributed on an \"AS IS\" BASIS,\n* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n* See the License for the specific language governing permissions and\n* limitations under the License.\n*/\n" - val scalaStyle = "// scalastyle:off" - val packageDef = "package org.apache.mxnet" - val imports = "import org.apache.mxnet.annotation.Experimental" - val absClassDef = s"abstract class $packageName" - val finalStr = s"$apacheLicence\n$scalaStyle\n$packageDef\n$imports\n$absClassDef {\n${absFuncs.mkString("\n")}\n}" - val pw = new PrintWriter(new File(FILE_PATH + s"$packageName.scala")) - pw.write(finalStr) - pw.close() - MD5Generator(finalStr) + writeFile( + FILE_PATH, + if (isSymbol) "SymbolAPIBase" else "NDArrayAPIBase", + body) } - def nonTypeSafeClassGen(FILE_PATH : String, isSymbol : Boolean) : String = { - // scalastyle:off - val absClassFunctions = getSymbolNDArrayMethods(isSymbol) - val absFuncs = absClassFunctions.map(absClassFunction => { - val scalaDoc = generateAPIDocFromBackend(absClassFunction, false) - if (isSymbol) { - val defBody = s"def ${absClassFunction.name}(name : String = null, attr : Map[String, String] = null)(args : org.apache.mxnet.Symbol*)(kwargs : Map[String, Any] = null): org.apache.mxnet.Symbol" - s"$scalaDoc\n$defBody" - } else { - val defBodyWithKwargs = s"def ${absClassFunction.name}(kwargs: Map[String, Any] = null)(args: Any*) : org.apache.mxnet.NDArrayFuncReturn" - val defBody = s"def ${absClassFunction.name}(args: Any*) : org.apache.mxnet.NDArrayFuncReturn" - s"$scalaDoc\n$defBodyWithKwargs\n$scalaDoc\n$defBody" - } - }) + def nonTypeSafeClassGen(FILE_PATH: String, isSymbol: Boolean): String = { + val absClassFunctions = buildFunctionList(isSymbol) + val absFuncs = absClassFunctions + .filterNot(_.name.startsWith("_")) + .map(absClassFunction => { + val scalaDoc = generateAPIDocFromBackend(absClassFunction, false) + if (isSymbol) { + val defBody = + s"def ${absClassFunction.name}(name : String = null, attr : Map[String, String] = null)" + + s"(args : org.apache.mxnet.Symbol*)(kwargs : Map[String, Any] = null): " + + s"org.apache.mxnet.Symbol" + s"$scalaDoc\n$defBody" + } else { + val defBodyWithKwargs = s"def ${absClassFunction.name}(kwargs: Map[String, Any] = null)" + + s"(args: Any*): " + + s"org.apache.mxnet.NDArrayFuncReturn" + val defBody = s"def ${absClassFunction.name}(args: Any*): " + + s"org.apache.mxnet.NDArrayFuncReturn" + s"$scalaDoc\n$defBodyWithKwargs\n$scalaDoc\n$defBody" + } + }) val packageName = if (isSymbol) "SymbolBase" else "NDArrayBase" - val apacheLicence = "/*\n* Licensed to the Apache Software Foundation (ASF) under one or more\n* contributor license agreements. See the NOTICE file distributed with\n* this work for additional information regarding copyright ownership.\n* The ASF licenses this file to You under the Apache License, Version 2.0\n* (the \"License\"); you may not use this file except in compliance with\n* the License. You may obtain a copy of the License at\n*\n* http://www.apache.org/licenses/LICENSE-2.0\n*\n* Unless required by applicable law or agreed to in writing, software\n* distributed under the License is distributed on an \"AS IS\" BASIS,\n* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n* See the License for the specific language governing permissions and\n* limitations under the License.\n*/\n" + writeFile(FILE_PATH, packageName, absFuncs) + } + + def writeFile(FILE_PATH: String, packageName: String, body: Seq[String]): String = { + val apacheLicence = + """/* + |* Licensed to the Apache Software Foundation (ASF) under one or more + |* contributor license agreements. See the NOTICE file distributed with + |* this work for additional information regarding copyright ownership. + |* The ASF licenses this file to You 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. + |*/ + |""".stripMargin val scalaStyle = "// scalastyle:off" val packageDef = "package org.apache.mxnet" val imports = "import org.apache.mxnet.annotation.Experimental" val absClassDef = s"abstract class $packageName" - val finalStr = s"$apacheLicence\n$scalaStyle\n$packageDef\n$imports\n$absClassDef {\n${absFuncs.mkString("\n")}\n}" - import java.io._ + val finalStr = + s"""$apacheLicence + |$scalaStyle + |$packageDef + |$imports + |$absClassDef { + |${body.mkString("\n")} + |}""".stripMargin val pw = new PrintWriter(new File(FILE_PATH + s"$packageName.scala")) pw.write(finalStr) pw.close() @@ -107,20 +143,15 @@ private[mxnet] object APIDocGenerator{ } // Generate ScalaDoc type - def generateAPIDocFromBackend(func : absClassFunction, withParam : Boolean = true) : String = { + def generateAPIDocFromBackend(func: absClassFunction, withParam: Boolean = true): String = { val desc = ArrayBuffer[String]() desc += " *
" - func.desc.split("\n").foreach({ currStr => + func.desc.split("\n").foreach({ currStr => desc += s" * $currStr" }) desc += " *" val params = func.listOfArgs.map({ absClassArg => - val currArgName = absClassArg.argName match { - case "var" => "vari" - case "type" => "typeOf" - case _ => absClassArg.argName - } - s" * @param $currArgName\t\t${absClassArg.argDesc}" + s" * @param ${absClassArg.safeArgName}\t\t${absClassArg.argDesc}" }) val returnType = s" * @return ${func.returnType}" if (withParam) { @@ -130,65 +161,23 @@ private[mxnet] object APIDocGenerator{ } } - def generateAPISignature(func : absClassFunction, isSymbol : Boolean) : String = { - var argDef = ListBuffer[String]() - func.listOfArgs.foreach(absClassArg => { - val currArgName = absClassArg.argName match { - case "var" => "vari" - case "type" => "typeOf" - case _ => absClassArg.argName - } - if (absClassArg.isOptional) { - argDef += s"$currArgName : Option[${absClassArg.argType}] = None" - } - else { - argDef += s"$currArgName : ${absClassArg.argType}" - } - }) - var returnType = func.returnType + def generateAPISignature(func: absClassFunction, isSymbol: Boolean): String = { + val argDef = ListBuffer[String]() + + argDef ++= buildArgDefs(func) + if (isSymbol) { argDef += "name : String = null" argDef += "attr : Map[String, String] = null" } else { argDef += "out : Option[NDArray] = None" - returnType = "org.apache.mxnet.NDArrayFuncReturn" } + + val returnType = func.returnType + val experimentalTag = "@Experimental" s"$experimentalTag\ndef ${func.name} (${argDef.mkString(", ")}) : $returnType" } - // List and add all the atomic symbol functions to current module. - private def getSymbolNDArrayMethods(isSymbol : Boolean): List[absClassFunction] = { - val opNames = ListBuffer.empty[String] - val returnType = if (isSymbol) "Symbol" else "NDArray" - _LIB.mxListAllOpNames(opNames) - // TODO: Add '_linalg_', '_sparse_', '_image_' support - // TODO: Add Filter to the same location in case of refactor - opNames.map(opName => { - val opHandle = new RefLong - _LIB.nnGetOpHandle(opName, opHandle) - makeAtomicSymbolFunction(opHandle.value, opName, "org.apache.mxnet." + returnType) - }).toList.filterNot(_.name.startsWith("_")) - } - - // Create an atomic symbol function by handle and function name. - private def makeAtomicSymbolFunction(handle: SymbolHandle, aliasName: String, returnType : String) - : absClassFunction = { - val name = new RefString - val desc = new RefString - val keyVarNumArgs = new RefString - val numArgs = new RefInt - val argNames = ListBuffer.empty[String] - val argTypes = ListBuffer.empty[String] - val argDescs = ListBuffer.empty[String] - - _LIB.mxSymbolGetAtomicSymbolInfo( - handle, name, desc, numArgs, argNames, argTypes, argDescs, keyVarNumArgs) - val argList = argNames zip argTypes zip argDescs map { case ((argName, argType), argDesc) => - val typeAndOption = CToScalaUtils.argumentCleaner(argName, argType, returnType) - new absClassArg(argName, typeAndOption._1, argDesc, typeAndOption._2) - } - new absClassFunction(aliasName, desc.value, argList.toList, returnType) - } } diff --git a/scala-package/macros/src/main/scala/org/apache/mxnet/GeneratorBase.scala b/scala-package/macros/src/main/scala/org/apache/mxnet/GeneratorBase.scala new file mode 100644 index 000000000000..84b27ea79fbf --- /dev/null +++ b/scala-package/macros/src/main/scala/org/apache/mxnet/GeneratorBase.scala @@ -0,0 +1,176 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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.apache.mxnet + +import org.apache.mxnet.init.Base.{RefInt, RefLong, RefString, _LIB} +import org.apache.mxnet.utils.{CToScalaUtils, OperatorBuildUtils} + +import scala.collection.mutable.ListBuffer +import scala.reflect.macros.blackbox + +abstract class GeneratorBase { + type Handle = Long + + case class Arg(argName: String, argType: String, argDesc: String, isOptional: Boolean) { + def safeArgName: String = argName match { + case "var" => "vari" + case "type" => "typeOf" + case _ => argName + } + } + + case class Func(name: String, desc: String, listOfArgs: List[Arg], returnType: String) + + protected def buildFunctionList(isSymbol: Boolean): List[Func] = { + val opNames = ListBuffer.empty[String] + _LIB.mxListAllOpNames(opNames) + opNames.map(opName => { + val opHandle = new RefLong + _LIB.nnGetOpHandle(opName, opHandle) + makeAtomicFunction(opHandle.value, opName, isSymbol) + }).toList + } + + protected def buildRandomFunctionList(isSymbol: Boolean): List[Func] = { + buildFunctionList(isSymbol) + .filter(f => f.name.startsWith("_sample_") || f.name.startsWith("_random_")) + .map(f => f.copy(name = f.name.stripPrefix("_"))) + // unify _random and _sample + .map(f => unifyRandom(f, isSymbol)) + // deduplicate + .groupBy(_.name) + .mapValues(_.head) + .values + .toList + } + + protected def unifyRandom(func: Func, isSymbol: Boolean): Func = { + var typeConv = if (isSymbol) + Map( + "org.apache.mxnet.Symbol" -> "Any", + "org.apache.mxnet.Base.MXFloat" -> "Any", + "Int" -> "Any" + ) + else + Map( + "org.apache.mxnet.NDArray" -> "Any", + "org.apache.mxnet.Base.MXFloat" -> "Any", + "Int" -> "Any" + ) + + func.copy( + name = func.name.replaceAll("(random|sample)_", ""), + listOfArgs = func.listOfArgs + .map { arg => + // This is hack to manage the fact that random_normal and sample_normal have + // non-consistent parameter naming in the back-end + if (arg.argName == "loc") arg.copy(argName = "mu") + else if (arg.argName == "scale") arg.copy(argName = "sigma") + else arg + } + .map { arg => + arg.copy(argType = typeConv.getOrElse(arg.argType, arg.argType)) + } + ) + } + + protected def makeAtomicFunction(handle: Handle, aliasName: String, isSymbol: Boolean): Func = { + val name = new RefString + val desc = new RefString + val keyVarNumArgs = new RefString + val numArgs = new RefInt + val argNames = ListBuffer.empty[String] + val argTypes = ListBuffer.empty[String] + val argDescs = ListBuffer.empty[String] + + _LIB.mxSymbolGetAtomicSymbolInfo( + handle, name, desc, numArgs, argNames, argTypes, argDescs, keyVarNumArgs) + val paramStr = OperatorBuildUtils.ctypes2docstring(argNames, argTypes, argDescs) + val extraDoc: String = if (keyVarNumArgs.value != null && keyVarNumArgs.value.length > 0) { + s"This function support variable length of positional input (${keyVarNumArgs.value})." + } else { + "" + } + val realName = if (aliasName == name.value) "" else s"(a.k.a., ${name.value})" + val docStr = s"$aliasName $realName\n${desc.value}\n\n$paramStr\n$extraDoc\n" + // scalastyle:off println + if (System.getenv("MXNET4J_PRINT_OP_DEF") != null + && System.getenv("MXNET4J_PRINT_OP_DEF").toLowerCase == "true") { + println("Function definition:\n" + docStr) + } + // scalastyle:on println + val argList = argNames zip argTypes zip argDescs map { case ((argName, argType), argDesc) => + val family = if (isSymbol) "org.apache.mxnet.Symbol" else "org.apache.mxnet.NDArray" + val typeAndOption = + CToScalaUtils.argumentCleaner(argName, argType, family) + Arg(argName, typeAndOption._1, argDesc, typeAndOption._2) + } + val returnType = if (isSymbol) "org.apache.mxnet.Symbol" else "org.apache.mxnet.NDArrayFuncReturn" + Func(aliasName, desc.value, argList.toList, returnType) + } + + /** + * Generate class structure for all function APIs + * + * @param c + * @param funcDef DefDef type of function definitions + * @param annottees + * @return + */ + protected def structGeneration(c: blackbox.Context) + (funcDef: List[c.universe.DefDef], annottees: c.Expr[Any]*) + : c.Expr[Any] = { + import c.universe._ + val inputs = annottees.map(_.tree).toList + // pattern match on the inputs + val modDefs = inputs map { + case ClassDef(mods, name, something, template) => + val q = template match { + case Template(superMaybe, emptyValDef, defs) => + Template(superMaybe, emptyValDef, defs ++ funcDef) + case ex => + throw new IllegalArgumentException(s"Invalid template: $ex") + } + ClassDef(mods, name, something, q) + case ModuleDef(mods, name, template) => + val q = template match { + case Template(superMaybe, emptyValDef, defs) => + Template(superMaybe, emptyValDef, defs ++ funcDef) + case ex => + throw new IllegalArgumentException(s"Invalid template: $ex") + } + ModuleDef(mods, name, q) + case ex => + throw new IllegalArgumentException(s"Invalid macro input: $ex") + } + // wrap the result up in an Expr, and return it + val result = c.Expr(Block(modDefs, Literal(Constant()))) + result + } + + protected def buildArgDefs(func: Func): List[String] = { + func.listOfArgs.map(arg => + if (arg.isOptional) + s"${arg.safeArgName} : Option[${arg.argType}] = None" + else + s"${arg.safeArgName} : ${arg.argType}" + ) + } + + +} diff --git a/scala-package/macros/src/main/scala/org/apache/mxnet/NDArrayMacro.scala b/scala-package/macros/src/main/scala/org/apache/mxnet/NDArrayMacro.scala index 2d3a1c7ec5af..8b7ad67ba8b8 100644 --- a/scala-package/macros/src/main/scala/org/apache/mxnet/NDArrayMacro.scala +++ b/scala-package/macros/src/main/scala/org/apache/mxnet/NDArrayMacro.scala @@ -17,11 +17,8 @@ package org.apache.mxnet -import org.apache.mxnet.init.Base._ -import org.apache.mxnet.utils.{CToScalaUtils, OperatorBuildUtils} - import scala.annotation.StaticAnnotation -import scala.collection.mutable.{ArrayBuffer, ListBuffer} +import scala.collection.mutable.ListBuffer import scala.language.experimental.macros import scala.reflect.macros.blackbox @@ -33,21 +30,37 @@ private[mxnet] class AddNDArrayAPIs(isContrib: Boolean) extends StaticAnnotation private[mxnet] def macroTransform(annottees: Any*) = macro NDArrayMacro.typeSafeAPIDefs } -private[mxnet] object NDArrayMacro { - case class NDArrayArg(argName: String, argType: String, isOptional : Boolean) - case class NDArrayFunction(name: String, listOfArgs: List[NDArrayArg]) +private[mxnet] class AddNDArrayRandomAPIs(isContrib: Boolean) extends StaticAnnotation { + private[mxnet] def macroTransform(annottees: Any*) = macro NDArrayMacro.typeSafeRandomAPIDefs +} + + +private[mxnet] object NDArrayMacro extends GeneratorBase { + type NDArrayArg = Arg + type NDArrayFunction = Func // scalastyle:off havetype def addDefs(c: blackbox.Context)(annottees: c.Expr[Any]*) = { impl(c)(annottees: _*) } + def typeSafeAPIDefs(c: blackbox.Context)(annottees: c.Expr[Any]*) = { - typeSafeAPIImpl(c)(annottees: _*) + typedAPIImpl(c)(annottees: _*) } + + def typeSafeRandomAPIDefs(c: blackbox.Context)(annottees: c.Expr[Any]*) = { + typedRandomAPIImpl(c)(annottees: _*) + } + // scalastyle:off havetype - private val ndarrayFunctions: List[NDArrayFunction] = initNDArrayModule() + private val ndarrayFunctions = buildFunctionList(false) + private val rndFunctions = buildRandomFunctionList(false) + + /** + * Implementation for fixed input API structure + */ private def impl(c: blackbox.Context)(annottees: c.Expr[Any]*): c.Expr[Any] = { import c.universe._ @@ -60,26 +73,29 @@ private[mxnet] object NDArrayMacro { else ndarrayFunctions.filterNot(_.name.startsWith("_")) } - val functionDefs = newNDArrayFunctions flatMap { NDArrayfunction => - val funcName = NDArrayfunction.name - val termName = TermName(funcName) - Seq( - // scalastyle:off - // (yizhi) We are investigating a way to make these functions type-safe - // and waiting to see the new approach is stable enough. - // Thus these functions may be deprecated in the future. - // e.g def transpose(kwargs: Map[String, Any] = null)(args: Any*) - q"def $termName(kwargs: Map[String, Any] = null)(args: Any*) = {genericNDArrayFunctionInvoke($funcName, args, kwargs)}".asInstanceOf[DefDef], - // e.g def transpose(args: Any*) - q"def $termName(args: Any*) = {genericNDArrayFunctionInvoke($funcName, args, null)}".asInstanceOf[DefDef] - // scalastyle:on - ) - } - - structGeneration(c)(functionDefs, annottees : _*) + val functionDefs = newNDArrayFunctions flatMap { NDArrayfunction => + val funcName = NDArrayfunction.name + val termName = TermName(funcName) + Seq( + // scalastyle:off + // (yizhi) We are investigating a way to make these functions type-safe + // and waiting to see the new approach is stable enough. + // Thus these functions may be deprecated in the future. + // e.g def transpose(kwargs: Map[String, Any] = null)(args: Any*) + q"def $termName(kwargs: Map[String, Any] = null)(args: Any*) = {genericNDArrayFunctionInvoke($funcName, args, kwargs)}".asInstanceOf[DefDef], + // e.g def transpose(args: Any*) + q"def $termName(args: Any*) = {genericNDArrayFunctionInvoke($funcName, args, null)}".asInstanceOf[DefDef] + // scalastyle:on + ) + } + + structGeneration(c)(functionDefs, annottees: _*) } - private def typeSafeAPIImpl(c: blackbox.Context)(annottees: c.Expr[Any]*) : c.Expr[Any] = { + /** + * Implementation for Dynamic typed API NDArray.api.