Add LM precompile workload

Project.toml

@@ -8,6 +8,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
 ManifoldsBase = "3362f125-f0bb-47a3-aa74-596ffd7ef2fb"
 Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
@@ -18,6 +19,7 @@ JuliaFormatter = "=2.10.1"
 Manifolds = "0.11.28"
 ManifoldsBase = "2.3.5"
 Manopt = "0.6"
+PrecompileTools = "1"
 ProgressMeter = "1.11.0"
 RecursiveArrayTools = "4.3"
 TensorOperations = "5.6"

docs/src/PIPELINE.md

@@ -42,7 +42,8 @@ Most optimization APIs share this core pattern:
 3. Convert to a public result struct
 
 `_solve_model` lives in `src/solvers/solve_dispatch.jl` and is the common
-symbol-to-solver dispatch layer (`:rgd`, `:rcg`, `:lbfgs`, `:als`, `:btd_tsd`).
+symbol-to-solver dispatch layer (`:rgd`, `:rgd_fixed`, `:rcg`, `:lbfgs`, `:lm`,
+`:als`, `:btd_tsd`).
 
 ## API flows
 
@@ -58,7 +59,7 @@ symbol-to-solver dispatch layer (`:rgd`, `:rcg`, `:lbfgs`, `:als`, `:btd_tsd`).
 Notes:
 
 - `:als` means CP-ALS.
-- Manifold solvers (`:rgd`, `:rgd_fixed`, `:rcg`, `:lbfgs`) share dispatch with other pipelines.
+- Manifold solvers (`:rgd`, `:rgd_fixed`, `:rcg`, `:lbfgs`, `:lm`) share dispatch with other pipelines.
 - For `solver != :als`, `init = :auto` resolves to `:alswarm`, so CPD and NNCPD start from an ALS warm point before manifold refinement.
 - Generic `approx(...)` does not use CPD's ALS warm-start path unless it auto-routes to `cpd(...)`.
 

src/TensorKitchen.jl

@@ -6,6 +6,7 @@ using Random
 using ManifoldsBase
 using Manifolds
 using Manopt
+using PrecompileTools
 import ProgressMeter
 using RecursiveArrayTools
 using TensorOperations
@@ -19,6 +20,7 @@ include("manifolds.jl")
 include("decompositions.jl")
 include("backend.jl")
 include("api.jl")
+include("precompile.jl")
 
 function __init__()
     return nothing

src/api/cpd.jl

@@ -988,20 +988,20 @@ If `r` is omitted, uses the smallest tensor mode as a heuristic rank.
         - `:random`: Uses a random initial point.
         - `:hosvd`: Uses a HOSVD initial point.
 * `solver = :rgd`: Sets the algorithm for refinement. Possible options are:
-    - `rgd` (default): Riemannian gradient descent
-    - `rgd_fixed`: Riemannian gradient descent with fixed step size
-    - `rcg`: Riemannian conjugate gradient
-    - `lbfgs`: Limited-memory Riemannian quasi-Newton
-    - `lm`: Levenberg-Marquardt using residual/Jacobian least squares
-    - `als`: Alternating Least Squares
+  - `rgd` (default): Riemannian gradient descent
+  - `rgd_fixed`: Riemannian gradient descent with fixed step size
+  - `rcg`: Riemannian conjugate gradient
+  - `lbfgs`: Limited-memory Riemannian quasi-Newton
+  - `lm`: Levenberg-Marquardt using residual/Jacobian least squares
+  - `als`: Alternating Least Squares
 
 ## Extended Options
 * `p0 = nothing`: Explicit initial point. If provided, it overrides the default initial point.
 * `:alswarm`: ALS warm start option.
     - `warm_init = TuckerInit()`: Before finding the warm start initial point, this sets the good starting point for ALS.
     - `warm_steps = 500`: Once finding the best initial point from warm_init, it runs this many ALS iterations to refine the initial point.
-* `maxiter = 500`: Maximum number of Riemannian gradient descent iterations.
-* `stepsize`: Initial step size for line search in Riemannian gradient descent. Defaults to `1.0` for ordinary CPD and `0.01` for the nonnegative route.
+* `maxiter = 500`: Maximum number of refinement iterations.
+* `stepsize`: Initial step size for first-order line-search solvers. Defaults to `1.0` for ordinary CPD and `0.01` for the nonnegative route.
 * `armijo_alpha_min = 1e-8`: Minimum Armijo line-search step size for `solver = :rgd`.
 * `tol = 1e-6`: Convergence tolerance.
 * `gradient_mode = :riemannian`: Gradient rule for manifold solvers. 

src/api/nncpd.jl

@@ -14,6 +14,8 @@ end
 Computes a nonnegative rank-`r` CP approximation of `A` in two steps: (1) the first step finds an initial point; (2) the second step refines the initial point. Returns a [`CPDResult`](@ref). 
 If `r` is omitted, uses the smallest tensor mode as a heuristic rank.
 `cpd(A, r; nonnegative=true, ...)` routes here and adopts the same effective defaults.
+With the default manifold refinement path, this means an ALS warm start followed
+by `solver = :rgd`.
 
 ## Options 
 The options are the same as for [`cpd`](@ref).

src/precompile.jl

@@ -0,0 +1,113 @@
+if Base.VERSION >= v"1.9"
+    using PrecompileTools: @compile_workload, @setup_workload
+
+    @setup_workload begin
+        rng = MersenneTwister(7)
+        A = randn(rng, 4, 3, 2)
+        B = abs.(A)
+        A_med = randn(rng, 6, 5, 4)
+
+        model = JoinModel(A, 2; geometry = :canonical)
+        p0 = _solver_point(manifold(model), initial_point(model, :random; verbose = false))
+        model_med = JoinModel(A_med, 2; geometry = :canonical)
+        p0_med = _solver_point(
+            manifold(model_med),
+            initial_point(model_med, :random; verbose = false),
+        )
+
+        nn_model = JoinModel(B, 2; nonnegative = true, geometry = :softplus_metric)
+        p0_nn = _solver_point(
+            manifold(nn_model),
+            initial_point(nn_model, :random; verbose = false),
+        )
+
+        # The first LM call is dominated by JIT compilation in the operator-action
+        # path. Precompiling a minimal canonical and softplus NNCP workload moves
+        # that cost to package precompile time instead of first user execution.
+        @compile_workload begin
+            # Direct LM solve path on an already-prepared canonical CP model.
+            solve(
+                LMSolver(),
+                model;
+                init = :random,
+                p0 = p0,
+                maxiter = 1,
+                tol = 1e-6,
+                verbose = false,
+                return_stats = true,
+            )
+            solve(
+                LMSolver(),
+                model_med;
+                init = :random,
+                p0 = p0_med,
+                maxiter = 1,
+                tol = 1e-6,
+                verbose = false,
+                return_stats = true,
+            )
+            cpd(
+                A,
+                2;
+                solver = :lm,
+                init = :random,
+                p0 = p0,
+                maxiter = 1,
+                tol = 1e-6,
+                verbose = false,
+            )
+            cpd(
+                A_med,
+                2;
+                solver = :lm,
+                init = :random,
+                p0 = p0_med,
+                maxiter = 1,
+                tol = 1e-6,
+                verbose = false,
+            )
+            # Frontend LM route with ALS warm start, which is the path users hit
+            # most often when they request cpd(...; solver=:lm).
+            cpd(
+                A,
+                2;
+                solver = :lm,
+                init = :alswarm,
+                warm_steps = 2,
+                maxiter = 1,
+                tol = 1e-6,
+                verbose = false,
+            )
+            cpd(
+                A_med,
+                2;
+                solver = :lm,
+                init = :alswarm,
+                warm_steps = 2,
+                maxiter = 1,
+                tol = 1e-6,
+                verbose = false,
+            )
+            nncpd(
+                B,
+                2;
+                solver = :lm,
+                init = :random,
+                p0 = p0_nn,
+                maxiter = 1,
+                tol = 1e-6,
+                verbose = false,
+            )
+            nncpd(
+                B,
+                2;
+                solver = :lm,
+                init = :alswarm,
+                warm_steps = 2,
+                maxiter = 1,
+                tol = 1e-6,
+                verbose = false,
+            )
+        end
+    end
+end