<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> Check for duplicate requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> Describe your goal, and if possible provide a code snippet with a motivating example.</li>
</ul>
<p dir="auto">I would like to have a warning that tells me that I should consider increasing the value of <code class="notranslate">max_squarings</code> when <code class="notranslate">expm</code> outputs nan values. For example <code class="notranslate">Warning: expm outputs is nan because max_squarings was reached. Consider increasing max_squarings in order to solve this problem.</code><br>
I had this issue and did not realize that it was the cause of my problem (see <a href="https://github.com/google/jax/issues/13985" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/13985/hovercard">this</a>).</p>
<p dir="auto">Typically this would be useful in this example:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp
from jax.scipy.linalg import expm as jexpm

dimension_x = 10
t = 1e6


# the jax way
sqrt_H = jax.random.normal(jax.random.PRNGKey(1), (dimension_x, dimension_x))
H = jnp.dot(sqrt_H, sqrt_H.T)
print(jexpm(-t*H))"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">scipy</span>.<span class="pl-s1">linalg</span> <span class="pl-k">import</span> <span class="pl-s1">expm</span> <span class="pl-k">as</span> <span class="pl-s1">jexpm</span>

<span class="pl-s1">dimension_x</span> <span class="pl-c1">=</span> <span class="pl-c1">10</span>
<span class="pl-s1">t</span> <span class="pl-c1">=</span> <span class="pl-c1">1e6</span>


<span class="pl-c"># the jax way</span>
<span class="pl-s1">sqrt_H</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">1</span>), (<span class="pl-s1">dimension_x</span>, <span class="pl-s1">dimension_x</span>))
<span class="pl-v">H</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">dot</span>(<span class="pl-s1">sqrt_H</span>, <span class="pl-s1">sqrt_H</span>.<span class="pl-v">T</span>)
<span class="pl-en">print</span>(<span class="pl-en">jexpm</span>(<span class="pl-c1">-</span><span class="pl-s1">t</span><span class="pl-c1">*</span><span class="pl-v">H</span>))</pre></div>
<p dir="auto">The following code does not work,</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="class Foo:
  @jax.jit
  @staticmethod
  def foo(x):
    return x * 2"><pre class="notranslate"><span class="pl-k">class</span> <span class="pl-v">Foo</span>:
  <span class="pl-en">@<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span></span>
  <span class="pl-en">@<span class="pl-s1">staticmethod</span></span>
  <span class="pl-k">def</span> <span class="pl-en">foo</span>(<span class="pl-s1">x</span>):
    <span class="pl-k">return</span> <span class="pl-s1">x</span> <span class="pl-c1">*</span> <span class="pl-c1">2</span></pre></div>
<p dir="auto">with an error <code class="notranslate">TypeError: Expected a callable value, got &lt;staticmethod object at 0x7fe007e767f0&gt;</code>. This is because a <code class="notranslate">staticmethod</code> is NOT a <code class="notranslate">callable</code>, <a href="https://bugs.python.org/issue43682" rel="nofollow">until Python 3.10</a>. Of course, staticmethod is <em>callable</em> -- something we can call.</p>
<p dir="auto">In <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="485382843" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/1251" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/1251/hovercard" href="https://github.com/google/jax/issues/1251">#1251</a> one possible workaround is suggested:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="class Foo:
  @functools.partial(jax.jit, static_argnums=(0,))
  def foo(self, x):
    # this function should not make use of `self` to be pure-functional
    return x * 2"><pre class="notranslate"><span class="pl-k">class</span> <span class="pl-v">Foo</span>:
  <span class="pl-en">@<span class="pl-s1">functools</span>.<span class="pl-en">partial</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span>, <span class="pl-s1">static_argnums</span><span class="pl-c1">=</span>(<span class="pl-c1">0</span>,))</span>
  <span class="pl-k">def</span> <span class="pl-en">foo</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>):
    <span class="pl-c"># this function should not make use of `self` to be pure-functional</span>
    <span class="pl-k">return</span> <span class="pl-s1">x</span> <span class="pl-c1">*</span> <span class="pl-c1">2</span></pre></div>
<p dir="auto">which is a bit ugly, but basically does the same thing as long as <code class="notranslate">self</code> is not used.</p>
<p dir="auto">I don't see any reason staticmethod cannot be jit-ed as other non-method functions. I think this is just a matter of extending <code class="notranslate">_check_callable(fun)</code> to support <code class="notranslate">staticmethod</code>s (which will happen automatically in Python 3.10); any pitfalls I would've missed?</p>