message: Copy `function.register` helper to Keras.
diff: @@ -839,7 +839,7 @@ def gru_with_backend_selection(inputs, init_h, kernel, recurrent_kernel, bias,
  # Call the normal GRU impl and register the CuDNN impl function. The
  # grappler will kick in during session execution to optimize the graph.
  last_output, outputs, new_h, runtime = defun_standard_gru(**params)
  - function.register(defun_gpu_gru, **params)
  + _function_register(defun_gpu_gru, **params)
  return last_output, outputs, new_h, runtime
  @@ -1653,7 +1653,7 @@ def lstm_with_backend_selection(inputs, init_h, init_c, kernel,
  # Call the normal LSTM impl and register the CuDNN impl function. The
  # grappler will kick in during session execution to optimize the graph.
  last_output, outputs, new_h, new_c, runtime = defun_standard_lstm(**params)
  - function.register(defun_gpu_lstm, **params)
  + _function_register(defun_gpu_lstm, **params)
  return last_output, outputs, new_h, new_c, runtime
  @@ -1765,3 +1765,30 @@ def _read_variable_value(v):
  if isinstance(v, variables.Variable):
  return v.read_value()
  return v
  +
  +
  +def _function_register(func, *args, **kwargs):
  + """Register a specialization of a `Function` into the graph.
  +
  + This won't actually call the function with the inputs, and only put the
  + function definition into graph. Register function with different input param
  + will result into multiple version of functions registered in graph.
  +
  + Args:
  + func: the `Function` instance that generated by a @defun
  + *args: input arguments for the Python function.
  + **kwargs: input keyword arguments for the Python function.
  +
  + Returns:
  + a `ConcreteFunction` object specialized to inputs and execution context.
  +
  + Raises:
  + ValueError: When the input function is not a defun wrapped python function.
  + """
  + if not isinstance(func, function.Function):
  + raise ValueError('Only defun function is allowed to be registered. '
  + 'Got type: %s' % type(func))
  + concrete_func = func.get_concrete_function(*args, **kwargs)
  + concrete_func.add_to_graph()
  + concrete_func.add_gradient_functions_to_graph()
  + return concrete_func
  