<p dir="auto">I am writing a probabilistic package and I am using jax as the backend. I am looking for Owens T function for calculating the CDF of  a  skewed Normal likelihood function. I searched through lax and jax.scipy.special and I think Owens T function is not implemented yet. Please let me know if this is the case. If so, is there any possiblity in the near future that you implement the function? I really don't want to switch to either SciPy pf tensorflow and integrate those in my code.</p>
<p dir="auto">Thanks,</p>
<p dir="auto"><a href="https://www.tensorflow.org/probability/api_docs/python/tfp/math/owens_t" rel="nofollow">https://www.tensorflow.org/probability/api_docs/python/tfp/math/owens_t</a><br>
<a href="https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.owens_t.html" rel="nofollow">https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.owens_t.html</a></p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> 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">In one of my networks, I have a large (64 MB) NumPy constant that I matmul against 160 different weight matrices.<br>
I am noticing that JAX is repeatedly calling <code class="notranslate">_raw_device_put</code> on it, seemingly taking up &gt;10G of accelerator HBM during tracing (resulting in an OOM at trace-time).</p>
<p dir="auto">A repro is below:</p>
<ol dir="auto">
<li>Add a print statement in  <code class="notranslate">convert_element_type</code> when <code class="notranslate">_device_put_raw</code> is called.</li>
<li>Run the below script.</li>
</ol>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index eaf94ced..0f3b5483 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -436,6 +436,7 @@ def convert_element_type(operand: Array, new_dtype: DType) -&gt; Array:
     if isinstance(operand, (core.Tracer, xla.DeviceArray)):
       return operand
     else:
+      print('_device_put_raw inside convert_element_type.')
       return _device_put_raw(np.asarray(operand))
   if (dtypes.issubdtype(old_dtype, np.complexfloating) and
       not dtypes.issubdtype(new_dtype, np.complexfloating)):"><pre class="notranslate"><code class="notranslate">diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index eaf94ced..0f3b5483 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -436,6 +436,7 @@ def convert_element_type(operand: Array, new_dtype: DType) -&gt; Array:
     if isinstance(operand, (core.Tracer, xla.DeviceArray)):
       return operand
     else:
+      print('_device_put_raw inside convert_element_type.')
       return _device_put_raw(np.asarray(operand))
   if (dtypes.issubdtype(old_dtype, np.complexfloating) and
       not dtypes.issubdtype(new_dtype, np.complexfloating)):
</code></pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import functools
import jax
import jax.numpy as jnp
import numpy as np
 
@jax.jit
def make_ws(key):
  return [k for k in jax.random.normal(key, (6, 4, 4))]
 
@functools.partial(jax.jit, static_argnums=1)
def run_10x(ws, wrap_array):
  # Create some NumPy constant.
  x = np.arange(8).reshape((2, 4)).astype(np.float32)
  if wrap_array:
    x = jnp.array(x)
  print('x type: ', type(x))
 
  # Use the same constant repeatedly.
  # If not wrap_array, it gets _raw_device_put repeatedly.
  # Real case: 160 layers, x is 64M -&gt; 10G used from this!
  print('Running jnp.dot(x, w) 6x.')
  ys = [jnp.dot(x, w) for w in ws]
  print('Done.')
  return sum(ys)
 
ws = make_ws(jax.random.PRNGKey(428))
print('w types: ', [type(w) for w in ws])
run_10x(ws, False)
run_10x(ws, True)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">functools</span>
<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">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>
 
<span class="pl-en">@<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span></span>
<span class="pl-k">def</span> <span class="pl-en">make_ws</span>(<span class="pl-s1">key</span>):
  <span class="pl-k">return</span> [<span class="pl-s1">k</span> <span class="pl-k">for</span> <span class="pl-s1">k</span> <span class="pl-c1">in</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-s1">key</span>, (<span class="pl-c1">6</span>, <span class="pl-c1">4</span>, <span class="pl-c1">4</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">1</span>)</span>
<span class="pl-k">def</span> <span class="pl-en">run_10x</span>(<span class="pl-s1">ws</span>, <span class="pl-s1">wrap_array</span>):
  <span class="pl-c"># Create some NumPy constant.</span>
  <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">arange</span>(<span class="pl-c1">8</span>).<span class="pl-en">reshape</span>((<span class="pl-c1">2</span>, <span class="pl-c1">4</span>)).<span class="pl-en">astype</span>(<span class="pl-s1">np</span>.<span class="pl-s1">float32</span>)
  <span class="pl-k">if</span> <span class="pl-s1">wrap_array</span>:
    <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>(<span class="pl-s1">x</span>)
  <span class="pl-en">print</span>(<span class="pl-s">'x type: '</span>, <span class="pl-en">type</span>(<span class="pl-s1">x</span>))
 
  <span class="pl-c"># Use the same constant repeatedly.</span>
  <span class="pl-c"># If not wrap_array, it gets _raw_device_put repeatedly.</span>
  <span class="pl-c"># Real case: 160 layers, x is 64M -&gt; 10G used from this!</span>
  <span class="pl-en">print</span>(<span class="pl-s">'Running jnp.dot(x, w) 6x.'</span>)
  <span class="pl-s1">ys</span> <span class="pl-c1">=</span> [<span class="pl-s1">jnp</span>.<span class="pl-en">dot</span>(<span class="pl-s1">x</span>, <span class="pl-s1">w</span>) <span class="pl-k">for</span> <span class="pl-s1">w</span> <span class="pl-c1">in</span> <span class="pl-s1">ws</span>]
  <span class="pl-en">print</span>(<span class="pl-s">'Done.'</span>)
  <span class="pl-k">return</span> <span class="pl-en">sum</span>(<span class="pl-s1">ys</span>)
 
<span class="pl-s1">ws</span> <span class="pl-c1">=</span> <span class="pl-en">make_ws</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">428</span>))
<span class="pl-en">print</span>(<span class="pl-s">'w types: '</span>, [<span class="pl-en">type</span>(<span class="pl-s1">w</span>) <span class="pl-k">for</span> <span class="pl-s1">w</span> <span class="pl-c1">in</span> <span class="pl-s1">ws</span>])
<span class="pl-en">run_10x</span>(<span class="pl-s1">ws</span>, <span class="pl-c1">False</span>)
<span class="pl-en">run_10x</span>(<span class="pl-s1">ws</span>, <span class="pl-c1">True</span>)</pre></div>
<p dir="auto">Output:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="tycai:~/jax$ python test.py 
...
w types:  [&lt;class 'jax.interpreters.xla._DeviceArray'&gt;, &lt;class 'jax.interpreters.xla._DeviceArray'&gt;, ..., &lt;class 'jax.interpreters.xla._DeviceArray'&gt;]
x type:  &lt;class 'numpy.ndarray'&gt;
Running jnp.dot(x, w) 6x.
_device_put_raw inside convert_element_type.
_device_put_raw inside convert_element_type.
_device_put_raw inside convert_element_type.
_device_put_raw inside convert_element_type.
_device_put_raw inside convert_element_type.
_device_put_raw inside convert_element_type.
Done.
_device_put_raw inside convert_element_type.
x type:  &lt;class 'jax.interpreters.xla._DeviceArray'&gt;
Running jnp.dot(x, w) 6x.
Done.
_device_put_raw inside convert_element_type."><pre class="notranslate"><code class="notranslate">tycai:~/jax$ python test.py 
...
w types:  [&lt;class 'jax.interpreters.xla._DeviceArray'&gt;, &lt;class 'jax.interpreters.xla._DeviceArray'&gt;, ..., &lt;class 'jax.interpreters.xla._DeviceArray'&gt;]
x type:  &lt;class 'numpy.ndarray'&gt;
Running jnp.dot(x, w) 6x.
_device_put_raw inside convert_element_type.
_device_put_raw inside convert_element_type.
_device_put_raw inside convert_element_type.
_device_put_raw inside convert_element_type.
_device_put_raw inside convert_element_type.
_device_put_raw inside convert_element_type.
Done.
_device_put_raw inside convert_element_type.
x type:  &lt;class 'jax.interpreters.xla._DeviceArray'&gt;
Running jnp.dot(x, w) 6x.
Done.
_device_put_raw inside convert_element_type.
</code></pre></div>
<p dir="auto">This seems like a sharp edge that might trip others up; it took me a couple hours to track down (in context of a larger model).<br>
I did not observe this prior to omnistaging, as I had <code class="notranslate">tie_in</code>-ed the NumPy constant, but <code class="notranslate">tie_in</code> was turned into a no-op.<br>
In the good <code class="notranslate">wrap_array</code> case, <code class="notranslate">jnp.array</code> is basically functioning as a <code class="notranslate">tie_in</code> replacement.</p>