-
Notifications
You must be signed in to change notification settings - Fork 0
Handle verbosity #9
Description
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: responsePerhaps 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.