Drop @defmodel decorator; plain functions work directly

@defmodel was a no-op marker that wrapped a (params, u) -> Node
function in a ModelFn container which compile_model immediately
unwrapped via _resolve_model_fn.  compile_model already accepts
plain functions identically, so the decorator added an import and
a layer of ceremony with zero runtime effect.

Remove ModelFn, defmodel, and _resolve_model_fn from compile.py;
narrow compile_model's signature to plain Callable; drop the
exports from __init__.py; update the migration comment in core.py;
update README quickstart to use compile_model directly (the
old set_params/log_likelihood example was bit-rotted to the
removed Model class anyway); migrate tests and examples to the
plain-function form.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: thisiscam

README.md

@@ -34,15 +34,14 @@ Uniform leaves, and gradient-based MLE in twenty lines:
 import jax
 import jax.numpy as jnp
 from tensorflow_probability.substrates import jax as tfp
-from acopula import copula, defmodel, marginal
+from acopula import compile_model, copula, marginal
 
 @copula
 class Frank:
     theta: float
     def generator(self, u):
         return -jnp.log1p(jnp.expm1(-self.theta) * jnp.exp(-u)) / self.theta
 
-@defmodel
 def model(params, obs):
     outer = Frank(params[0])
     inner = Frank(params[1])
@@ -54,15 +53,12 @@ def model(params, obs):
 obs = jnp.array([[0.3, 0.7],
                  [0.4, 0.8]])
 
-model.set_params(jnp.array([2.0, 5.0]))
-print(model.log_likelihood(obs))                # scalar log-likelihood
+cm = compile_model(model, template=jnp.array([1.0, 1.0]))
+params = jnp.array([2.0, 5.0])
+print(cm.eval(obs, params))                     # scalar log-likelihood
 
-def neg_ll(params, obs):
-    model.set_params(params)
-    return -model.log_likelihood(obs)
-
-grad = jax.grad(neg_ll)(jnp.array([2.0, 5.0]), obs)
-print(grad)                                     # ∂(-ll)/∂params, exact
+grad = jax.grad(cm.eval, argnums=1)(obs, params)
+print(grad)                                     # ∂ll/∂params, exact
 ```
 
 ### With Weibull marginals and right-censoring
@@ -72,7 +68,6 @@ Use `float64` for the distribution parameters since `acopula` enables
 `jax_enable_x64` at import time.
 
 ```python
-@defmodel
 def survival_model(params, obs):
     outer = Frank(params[0])
     inner = Frank(params[1])
@@ -86,8 +81,9 @@ def survival_model(params, obs):
 
 `@copula` registers the parameters declared as type-annotated class
 attributes and derives the inverse generator symbolically via `oryx`.
-`@defmodel` traces the function into a copula tree, flattens parameters
-into a single array, and exposes `log_likelihood`, all `jit`/`grad`-compatible.
+`compile_model` traces the model function into a copula tree, flattens
+parameters into a single array, and returns a `CompiledModel` exposing
+a jit'd log-likelihood that is `jax.grad`/`jax.vmap`-compatible.
 `marginal` pairs each leaf with a distribution and an optional
 per-observation censoring flag.
 

acopula/__init__.py

@@ -1,13 +1,13 @@
 from . import _oryx_patch  # noqa: F401 — registers oryx ILDJ rule for lax.copy_p
 from .core import Copula, copula, marginal
-from .compile import compile_model, defmodel, CompiledModel, ModelFn
+from .compile import compile_model, CompiledModel
 from .compose import register_composition, set_composition_fallback
 from .bell import set_stable_log
 from ._compile_cache import set_compile_cache_dir
 
 __all__ = [
     "Copula", "copula", "marginal",
-    "compile_model", "defmodel", "CompiledModel", "ModelFn",
+    "compile_model", "CompiledModel",
     "register_composition", "set_composition_fallback",
     "set_stable_log",
     "set_compile_cache_dir",

acopula/compile.py

@@ -16,15 +16,14 @@
 
 Typical usage::
 
-    from acopula import compile_model, defmodel, marginal, copula
+    from acopula import compile_model, marginal, copula
 
     @copula
     class Clayton:
         theta: float
         def generator(self, t):
             return jnp.power(1.0 + t, -1.0 / self.theta)
 
-    @defmodel
     def model(p, u):
         return Clayton(p["theta"])([
             marginal(Uniform(), obs=u[i]) for i in range(5)
@@ -43,48 +42,14 @@ def model(p, u):
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Callable, Optional, Tuple, Union
+from typing import Any, Callable, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
 from jax import random as jrandom
 from oryx import core as oryx_core
 
 
-@dataclass(frozen=True)
-class ModelFn:
-    """Marker wrapper produced by ``@defmodel``.
-
-    Holding the user function in a typed container lets ``compile_model``
-    distinguish "function defined as a copula model" from arbitrary
-    callables, while still supporting plain function passing for users
-    who skip the decorator.
-    """
-
-    fn: Callable
-
-    def __call__(self, params, u):
-        return self.fn(params, u)
-
-
-def defmodel(fn: Callable) -> ModelFn:
-    """Mark a function as a copula composition model.
-
-    Decorated functions take ``(params, u)`` and must return a Node
-    produced by an outer copula call (see the package README for the
-    composition syntax).  The decorator itself does no work besides
-    wrapping; the actual graph construction happens inside
-    :func:`compile_model`.
-    """
-    return ModelFn(fn=fn)
-
-
-def _resolve_model_fn(model_fn: Union[Callable, ModelFn]) -> Callable:
-    if isinstance(model_fn, ModelFn):
-        return model_fn.fn
-    return model_fn
-
-
 def _harvest_graph_and_params(
     fn: Callable, template: Any
 ) -> Tuple[Any, dict]:
@@ -112,7 +77,7 @@ def wrapped_fn(p):
 
     if not isinstance(graph, Node):
         raise TypeError(
-            "@defmodel function must return a Copula call producing a Node"
+            "Model function must return a Copula call producing a Node"
         )
 
     leaves_in_order = _collect_leaves_in_order(graph)
@@ -356,7 +321,7 @@ def visualize(
 
 
 def compile_model(
-    model_fn: Union[Callable, ModelFn],
+    model_fn: Callable,
     *,
     template: Any,
     method: str = "auto",
@@ -381,8 +346,8 @@ def compile_model(
     invalidating cache) is avoided by construction.
 
     Args:
-        model_fn: A function ``(params, u) -> Node`` (typically wrapped
-            with ``@defmodel``).  Plain functions also work.
+        model_fn: A function ``(params, u) -> Node`` returning a Copula
+            call that produces the structural graph.
         template: A parameter pytree whose structure matches what
             ``model_fn`` expects.  Values are used only to drive oryx's
             tracer; pass any sensible defaults (e.g.
@@ -407,10 +372,8 @@ def compile_model(
         A :class:`CompiledModel` holding the frozen graph, flattener,
         and jit'd log-likelihood.
     """
-    fn = _resolve_model_fn(model_fn)
-
-    graph, _harvested, obs_shape = _harvest_graph_and_params(fn, template)
-    flatten, param_shape = _make_flatten(fn, template)
+    graph, _harvested, obs_shape = _harvest_graph_and_params(model_fn, template)
+    flatten, param_shape = _make_flatten(model_fn, template)
 
     raw_ll = _build_ll_fn(
         graph,
@@ -428,7 +391,7 @@ def compile_model(
         param_shape=param_shape,
         obs_shape=obs_shape,
         with_censored_mask=with_censored_mask,
-        _model_fn=fn,
+        _model_fn=model_fn,
         _flatten_fn=flatten,
         _ll_fn_jit=ll_fn_jit,
     )

acopula/core.py

@@ -595,9 +595,8 @@ def _reconstruct_obs_from_marginals(
 #
 # Migrate Model() / set_params() / make_ll_fn() callers to:
 #
-#     from acopula import compile_model, defmodel
+#     from acopula import compile_model
 #
-#     @defmodel
 #     def m(p, u): ...
 #
 #     cm = compile_model(m, template={'theta': 1.0}, method='bell')

examples/sample_and_plot_pairs.py

@@ -1,6 +1,6 @@
 """Generate samples using vmap and plot pairwise distributions."""
 
-from acopula import copula, defmodel, marginal
+from acopula import copula, marginal
 import jax
 import jax.numpy as jnp
 import jax.random as jrandom
@@ -25,7 +25,6 @@ def generator(self, u: jax.Array) -> jax.Array:
         return -jnp.log1p(jnp.expm1(-self.theta) * jnp.exp(-u)) / self.theta
 
 
-@defmodel
 def model(params, obs):
     c1 = Frank(params[0])
     c2 = Frank(params[1])

examples/two_level.py

@@ -1,4 +1,4 @@
-from acopula import Copula, defmodel
+from acopula import Copula
 import jax
 import jax.numpy as jnp
 
@@ -17,7 +17,6 @@ def generator(self, u: jax.Array) -> jax.Array:
         return -jax.log(1 - (jax.exp(-self.theta * u) - 1) / (jax.exp(-self.theta) - 1))
 
 
-@defmodel
 def model(params, u):
     f = Frank(params[0])
     c = Clayton(params[1])
@@ -30,7 +29,7 @@ def model(params, u):
     jax.random.PRNGKey(0), 100
 )  # this will generate 100 samples from the model
 # under the hood
-# defmodel calls the user defined `model` and builds a complete graph of the model
+# compile_model traces the user-defined `model` and builds a complete graph
 # then, we need to find a "plan" that traverses the graph
 # In particular, a plan is a sequence of stages, where each stage is a set of nodes of the same family and all ancestor families.
 # For each stage, all the nodes that are inputs to the current staged are already computed and thus available
@@ -49,7 +48,7 @@ def model(params, u):
     jnp.ones((10, 20))
 )  # this will compute the log likelihood of the model
 # under the hood
-# defmodel calls the user defined `model` and builds a complete graph of
+# compile_model traces the user-defined `model` and builds a complete graph
 # if the graph is a single layer tree (there's only one node of copula),
 # then the log likelihood is simply found by finding the d^th order mixed derivative of the outer generator, multiplied by the product of the derivatives of the inner generators.
 # if the graph is a multi-layer tree,

examples/visualize_sampling.py

@@ -1,6 +1,6 @@
 """Example showing step-by-step visualization of sampling computation."""
 
-from acopula import copula, defmodel
+from acopula import copula
 import jax
 import jax.numpy as jnp
 
@@ -21,7 +21,6 @@ def generator(self, u: jax.Array) -> jax.Array:
         return -jnp.log1p(jnp.expm1(-self.theta) * jnp.exp(-u)) / self.theta
 
 
-@defmodel
 def model(params, u):
     f = Clayton(params[0])
     c = Clayton(params[1])

examples/visualize_tree.py

@@ -1,4 +1,4 @@
-from acopula import Copula, defmodel
+from acopula import Copula
 import jax
 import jax.numpy as jnp
 
@@ -25,7 +25,6 @@ def generator(self, u: jax.Array) -> jax.Array:
         )
 
 
-@defmodel
 def model(params, u):
     f = Frank(params[0])
     c = Clayton(params[1])

tests/test_3level_quick.py

@@ -10,7 +10,7 @@
 
 jax.config.update("jax_enable_x64", True)
 
-from acopula import compile_model, copula, defmodel, marginal
+from acopula import compile_model, copula, marginal
 
 @copula
 class Clayton:
@@ -44,7 +44,6 @@ def main():
     n_groups, n_sectors, n_leaves = 2, 2, 2
     d = n_groups * n_sectors * n_leaves
 
-    @defmodel
     def model(params, obs):
         root = Clayton(params[0])
         groups = []