Expand Java definition extraction for TensorEncoderDef and InputSelec…

…torDef (#71) * Expand java extractor capability for InputSelectorDef and TensorEncoderDef * Format file --------- Co-authored-by: Steve Lee <[email protected]>
bazeltools · Jul 11, 2023 · e9dd379 · e9dd379
1 parent 41c4374
commit e9dd379
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/...xtractor/src/main/scala/io/bazeltools/buildgen/javadefref/JavaSourceEntityExtractor.scala b/...xtractor/src/main/scala/io/bazeltools/buildgen/javadefref/JavaSourceEntityExtractor.scala
@@ -61,7 +61,13 @@ object JavaSourceEntityExtractor {
 
     // TODO there needs to be some way to configure this
     val generators: Set[String] =
-      Set("InputTransformDef", "FeatureEncoderDef", "DataTransformDef")
+      Set(
+        "InputTransformDef",
+        "FeatureEncoderDef",
+        "DataTransformDef",
+        "TensorEncoderDef",
+        "InputSelectorDef"
+      )
 
     val topLevelDefsTypes: SortedSet[Entity] =
       compUnit