Skip to content
This repository was archived by the owner on Sep 9, 2025. It is now read-only.
This repository was archived by the owner on Sep 9, 2025. It is now read-only.

Handle verbosity #9

@arcruz0

Description

@arcruz0

I think users should be able to do

library(marginaleffects)
library(marginaleffectsJAX)
enable_JAX_backend()

and get the JAX speedup whenever possible, without much interference in unsupported calls.

For example, if marginaleffectsJAX only supports lm models, I think the current behavior is confusing. Users get nothing when the JAX backend is used, but get a warning when it is not used.

library(marginaleffects)
library(marginaleffectsJAX)
enable_JAX_backend(verbose = F)
#> JAX is now a backend for `marginaleffects`. Run `disable_JAX_backend()` to disable.

m_lm <- lm(mpg ~ hp, mtcars)
predictions(m_lm, by = T)
#> 
#>  Estimate Std. Error    z Pr(>|z|)     S 2.5 % 97.5 %
#>      20.1      0.683 29.4   <0.001 629.6  18.8   21.4
#> 
#> Type: response
m_glm <- glm(am ~ hp, mtcars, family = binomial())
predictions(m_glm, by = T)
#> Warning: `marginaleffectsJAX` only supports models of class `lm`. Reverting to
#> standard computation.
#> 
#>  Estimate Std. Error    z Pr(>|z|)    S 2.5 % 97.5 %
#>     0.406     0.0842 4.83   <0.001 19.5 0.241  0.571
#> 
#> Type: response

The current verbose = T argument is clear but too cumbersome:

enable_JAX_backend(verbose = T)
#> JAX is now a backend for `marginaleffects`. Run `disable_JAX_backend()` to disable.

m_lm <- lm(mpg ~ hp, mtcars)
predictions(m_lm, by = T)
#> --Calling JAX backend
#> --Executing JAX function
#> --Succesfully executed JAX function
#> 
#>  Estimate Std. Error    z Pr(>|z|)     S 2.5 % 97.5 %
#>      20.1      0.683 29.4   <0.001 629.6  18.8   21.4
#> 
#> Type: response
m_glm <- glm(am ~ hp, mtcars, family = binomial())
predictions(m_glm, by = T)
#> --Calling JAX backend
#> Warning: `marginaleffectsJAX` only supports models of class `lm`. Reverting to
#> standard computation.
#> 
#>  Estimate Std. Error    z Pr(>|z|)    S 2.5 % 97.5 %
#>     0.406     0.0842 4.83   <0.001 19.5 0.241  0.571
#> 
#> Type: response

Perhaps there's a compromise: a one-liner message any time jax_jacobian.R is called, saying whether the JAX backend was used or not. That could be the default verbosity. Users could alternatively get nothing or get a more descriptive "debug" message. See this rOpenSci blog post for implementation details—apparently the best practice here is to set up verbosity using options(), not function arguments.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions