Expose cached graph function names

Author: nfeybesse

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java

@@ -404,7 +404,7 @@ public GraphOperationBuilder opBuilder(String type, String name, Scope scope) {
    * immediately without re-registering it.
    *
    * <p>The function is also stored in an internal cache to speed up subsequent lookups performed by
-   * {@link #getFunction(String)} and {@link #getFunctionCached(String)}.
+   * {@link #getFunction(String)}.
    */
   @Override
   public void attachFunction(ConcreteFunction function) {
@@ -915,26 +915,15 @@ Set<Operation> initializers() {
       new ConcurrentHashMap<>();
 
   /**
-   * Returns a cached {@link ConcreteFunction} whose name starts with the provided prefix.
+   * Returns a read-only view of the function names cached by this graph.
    *
-   * <p>This is a lightweight lookup helper used when the exact function name is not known but
-   * follows a deterministic prefix (for example functions generated for control-flow constructs or
-   * custom gradient expansions).
+   * <p>This exposes only the function names so callers can resolve ambiguous matches themselves
+   * before calling {@link #getFunction(String)} with an exact name.
    *
-   * <p>The search is performed only in the local cache and does not query the native TensorFlow
-   * function library.
-   *
-   * @param prefix function name prefix
-   * @return a cached {@link ConcreteFunction} whose name starts with {@code prefix}, or {@code
-   *     null} if none is found
+   * @return a read-only view of cached function names
    */
-  public ConcreteFunction getFunctionCached(String prefix) {
-    for (Map.Entry<String, ConcreteFunction> e : functionCache.entrySet()) {
-      if (e.getKey().startsWith(prefix)) {
-        return e.getValue();
-      }
-    }
-    return null;
+  public Set<String> functionNames() {
+    return Collections.unmodifiableSet(functionCache.keySet());
   }
 
   /**

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/AttributeMetadata.java

@@ -30,6 +30,7 @@ public class AttributeMetadata {
 
   /** The size of the list if this attribute is a list, undefined otherwise. */
   public final long listSize;
+
   /**
    * The type of this attribute, or the type of the list values if it is a list.
    *

tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/IfGradientTest.java

@@ -144,6 +144,18 @@ private static Operand<?> sumAll(Ops tf, Operand<?> v) {
     return tf.reduceSum((Operand) v, axes);
   }
 
+  private static ConcreteFunction getSingleFunctionByPrefix(Graph graph, String prefix) {
+    List<String> matches =
+        graph.functionNames().stream()
+            .filter(name -> name.startsWith(prefix))
+            .collect(Collectors.toList());
+    if (matches.size() != 1) {
+      throw new IllegalStateException(
+          "Expected one cached function for prefix=" + prefix + ", found=" + matches);
+    }
+    return graph.getFunction(matches.get(0));
+  }
+
   @Test
   public void testStatefullIfGradient() {
     TensorFlow.registerCustomGradient(
@@ -199,8 +211,8 @@ public void testStatefullIfGradient() {
           final String thenPrefix = op.name() + "/then_grad"; // op has unique name
           final String elsePrefix = op.name() + "/else_grad";
 
-          ConcreteFunction thenGrad = op.env().getFunctionCached(thenPrefix);
-          ConcreteFunction elseGrad = op.env().getFunctionCached(elsePrefix);
+          ConcreteFunction thenGrad = getSingleFunctionByPrefix(op.env(), thenPrefix);
+          ConcreteFunction elseGrad = getSingleFunctionByPrefix(op.env(), elsePrefix);
 
           if (thenGrad == null || elseGrad == null) {
             throw new IllegalStateException("If grad functions not primed for op=" + op.name());