diff --git a/.github/workflows/continuous-integration.yml b/.github/workflows/continuous-integration.yml index 539ca59ec0..d9f59c244c 100644 --- a/.github/workflows/continuous-integration.yml +++ b/.github/workflows/continuous-integration.yml @@ -14,6 +14,8 @@ # ============================================================================ name: Tests on: [push, pull_request] +permissions: + contents: read jobs: lints: name: Lints @@ -23,11 +25,12 @@ jobs: python-version: [3.12] steps: - name: Checkout - uses: actions/checkout@v1 + uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 with: fetch-depth: 20 + persist-credentials: false - name: Setup Python - uses: actions/setup-python@v2 + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 with: python-version: ${{ matrix.python-version }} - name: Lints @@ -45,11 +48,12 @@ jobs: NUM_SHARDS: 5 steps: - name: Checkout - uses: actions/checkout@v1 + uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 with: fetch-depth: 1 + persist-credentials: false - name: Setup Python - uses: actions/setup-python@v2 + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 with: python-version: ${{ matrix.python-version }} - name: Tests @@ -57,7 +61,7 @@ jobs: ./testing/run_github_tests.sh - name: Upload test logs if: failure() - uses: actions/upload-artifact@v1 + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 with: name: testlogs-${{ matrix.shard }} path: bazel-testlogs diff --git a/debug_test.py b/debug_test.py new file mode 100644 index 0000000000..383d66d509 --- /dev/null +++ b/debug_test.py @@ -0,0 +1,23 @@ +import tensorflow as tf +import tensorflow_probability as tfp + +def quadratic_loss_and_grad(x): + """f(x) = (x-2)^2 + (y-3)^2, minimum at [2, 3]""" + diff = x - tf.constant([2.0, 3.0]) + loss = tf.reduce_sum(tf.square(diff)) + grad = 2.0 * diff + return loss, grad + +print("Testing line_search_kwargs feature...") + +# Test 1: Default behavior +start = tf.constant([0.0, 0.0]) +print("About to call bfgs_minimize...") +try: + results_default = tfp.optimizer.bfgs_minimize( + quadratic_loss_and_grad, start, tolerance=1e-10) + print("Default result:", results_default.position.numpy()) +except Exception as e: + print("Error in bfgs_minimize:", e) + import traceback + traceback.print_exc() \ No newline at end of file diff --git a/tensorflow_probability/python/optimizer/bfgs.py b/tensorflow_probability/python/optimizer/bfgs.py index ea03553800..d00d588ff9 100644 --- a/tensorflow_probability/python/optimizer/bfgs.py +++ b/tensorflow_probability/python/optimizer/bfgs.py @@ -62,10 +62,15 @@ # below the tolerance. 'inverse_hessian_estimate', # A tensor containing the inverse of the # estimated Hessian. - 'scale_initial_inverse_hessian' # Should the initial inverse Hessian + 'scale_initial_inverse_hessian', # Should the initial inverse Hessian # be rescaled on the first iteration, # as per Chapter 6 of Nocedal and # Wright. + 'line_search_kwargs', # A dict of keyword arguments to pass to the + # line search algorithm. These arguments are + # passed directly to the underlying line search + # implementation (e.g., the Hager-Zhang line + # search). ]) @@ -81,6 +86,7 @@ def minimize(value_and_gradients_function, stopping_condition=None, validate_args=True, max_line_search_iterations=50, + line_search_kwargs=None, f_absolute_tolerance=0, name=None): """Applies the BFGS algorithm to minimize a differentiable function. @@ -174,6 +180,14 @@ def quadratic_loss_and_gradient(x): outputs. max_line_search_iterations: Python int. The maximum number of iterations for the `hager_zhang` line search algorithm. + line_search_kwargs: A dict of keyword arguments to pass to the line + search algorithm. These arguments are passed directly to the underlying + line search implementation (e.g., the Hager-Zhang line search). + Common parameters include initial_step_size, value_at_initial_step, + value_at_zero, threshold_use_approximate_wolfe_condition, + shrinkage_param, expansion_param, sufficient_decrease_param, + curvature_param, and max_iterations. If not supplied, the line + search uses its default parameters. f_absolute_tolerance: Scalar `Tensor` of real dtype. If the absolute change in the objective value between one iteration and the next is smaller than this value, the algorithm is stopped. @@ -286,7 +300,7 @@ def _body(state): next_state = bfgs_utils.line_search_step( current_state, value_and_gradients_function, actual_search_direction, tolerance, f_relative_tolerance, x_tolerance, stopping_condition, - max_line_search_iterations, f_absolute_tolerance) + max_line_search_iterations, f_absolute_tolerance, line_search_kwargs=current_state.line_search_kwargs) # Update the inverse Hessian if needed and continue. return [_update_inv_hessian(current_state, next_state)] @@ -298,6 +312,7 @@ def _body(state): control_inputs) kwargs['inverse_hessian_estimate'] = initial_inv_hessian kwargs['scale_initial_inverse_hessian'] = scale_initial_inverse_hessian + kwargs['line_search_kwargs'] = line_search_kwargs if line_search_kwargs is not None else {} initial_state = BfgsOptimizerResults(**kwargs) return tf.while_loop( cond=_cond, diff --git a/tensorflow_probability/python/optimizer/bfgs_utils.py b/tensorflow_probability/python/optimizer/bfgs_utils.py index 0495651e3f..c22fd11c14 100644 --- a/tensorflow_probability/python/optimizer/bfgs_utils.py +++ b/tensorflow_probability/python/optimizer/bfgs_utils.py @@ -61,14 +61,14 @@ def get_initial_state_args(value_and_gradients_function, Returns: An dictionary with values for the following keys: converged: True if the convergence check finds that the initial position - is already an argmin of the objective function. + is already an argmin of the objective function. failed: Initialized to False. num_objective_evaluations: Initialized to 1. position: Initialized to the initial position. objective_value: Initialized to the value of the objective function at - the initial position. + the initial position. objective_gradient: Initialized to the gradient of the objective - function at the initial position. + function at the initial position. """ if control_inputs: with tf.control_dependencies(control_inputs): @@ -148,7 +148,8 @@ def _is_negative_inf(x): def line_search_step(state, value_and_gradients_function, search_direction, grad_tolerance, f_relative_tolerance, x_tolerance, - stopping_condition, max_iterations, f_absolute_tolerance): + stopping_condition, max_iterations, f_absolute_tolerance, + line_search_kwargs=None): """Performs the line search step of the BFGS search procedure. Uses hager_zhang line search procedure to compute a suitable step size @@ -164,7 +165,7 @@ def line_search_step(state, value_and_gradients_function, search_direction, value_and_gradients_function: A Python callable that accepts a point as a real `Tensor` of shape `[..., n]` and returns a tuple of two tensors of the same dtype: the objective function value, a real `Tensor` of shape - `[...]`, and its derivative, another real `Tensor` of shape `[..., n]`. + `[...]`, and its derivative, another real `Tensor` of shape `[..., n]`. search_direction: A real `Tensor` of shape `[..., n]`. The direction along which to perform line search. grad_tolerance: Scalar `Tensor` of real dtype. Specifies the gradient @@ -182,6 +183,9 @@ def line_search_step(state, value_and_gradients_function, search_direction, iterations of the hager_zhang line search algorithm f_absolute_tolerance: Scalar `Tensor` of real dtype. Specifies the tolerance for the absolute change in the objective value. + line_search_kwargs: A dict of keyword arguments to pass to the line + search algorithm. These arguments are passed directly to the underlying + line search implementation (e.g., the Hager-Zhang line search). Returns: A copy of the input state with the following fields updated: @@ -197,6 +201,14 @@ def line_search_step(state, value_and_gradients_function, search_direction, updated by computing the new position and evaluating the objective function at that position. """ + # Extract line_search_kwargs from state + line_search_kwargs = state.line_search_kwargs + # Remove parameters that are set explicitly in the hager_zhang call to avoid conflicts + filtered_line_search_kwargs = {k: v for k, v in line_search_kwargs.items() + if k not in ['initial_step_size', 'value_at_initial_step', 'value_at_zero', + 'converged', 'threshold_use_approximate_wolfe_condition', + 'shrinkage_param', 'expansion_param', 'sufficient_decrease_param', + 'curvature_param', 'max_iterations', 'name']} line_search_value_grad_func = _restrict_along_direction( value_and_gradients_function, state.position, search_direction) derivative_at_start_pt = tf.reduce_sum( @@ -211,7 +223,8 @@ def line_search_step(state, value_and_gradients_function, search_direction, initial_step_size=_broadcast(1, state.position), value_at_zero=val_0, converged=inactive, - max_iterations=max_iterations) # No search needed for these. + max_iterations=max_iterations, + **(filtered_line_search_kwargs or {})) # No search needed for these. state_after_ls = update_fields( state, diff --git a/test_line_search_kwargs.py b/test_line_search_kwargs.py new file mode 100644 index 0000000000..e9221669c9 --- /dev/null +++ b/test_line_search_kwargs.py @@ -0,0 +1,33 @@ +import tensorflow as tf +import tensorflow_probability as tfp + +def quadratic_loss_and_grad(x): + """f(x) = (x-2)^2 + (y-3)^2, minimum at [2, 3]""" + diff = x - tf.constant([2.0, 3.0]) + loss = tf.reduce_sum(tf.square(diff)) + grad = 2.0 * diff + return loss, grad + +print("Testing line_search_kwargs feature...") + +# Test 1: Default behavior +start = tf.constant([0.0, 0.0]) +results_default = tfp.optimizer.bfgs_minimize( + quadratic_loss_and_grad, start, tolerance=1e-10) +print("Default result:", results_default.position.numpy()) + +# Test 2: With custom line search parameters +custom_kwargs = { + 'initial_step_size': 0.5, + 'threshold_use_approximate_wolfe_condition': 1e-4 +} +results_custom = tfp.optimizer.bfgs_minimize( + quadratic_loss_and_grad, start, tolerance=1e-10, + line_search_kwargs=custom_kwargs) +print("Custom result:", results_custom.position.numpy()) + +# Both should be close to [2.0, 3.0] +expected = tf.constant([2.0, 3.0]) +assert tf.reduce_all(tf.abs(results_default.position - expected) < 1e-6) +assert tf.reduce_all(tf.abs(results_custom.position - expected) < 1e-6) +print("✅ All tests passed!") \ No newline at end of file