Issue 1:
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.

Thanks,

https://www.tensorflow.org/probability/api_docs/python/tfp/math/owens_t

https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.owens_t.html

Check for duplicate requests.

Describe your goal, and if possible provide a code snippet with a motivating example.

Issue 2:
In one of my networks, I have a large (64 MB) NumPy constant that I matmul against 160 different weight matrices.

I am noticing that JAX is repeatedly calling _raw_device_put on it, seemingly taking up >10G of accelerator HBM during tracing (resulting in an OOM at trace-time).

A repro is below:

Add a print statement in convert_element_type when _device_put_raw is called.

Run the below script.

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) -> 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)):

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 -> 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)

Output:

tycai:~/jax$ python test.py
...
w types: [<class 'jax.interpreters.xla._DeviceArray'>, <class 'jax.interpreters.xla._DeviceArray'>, ..., <class 'jax.interpreters.xla._DeviceArray'>]
x type: <class 'numpy.ndarray'>
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: <class 'jax.interpreters.xla._DeviceArray'>
Running jnp.dot(x, w) 6x.
Done.
_device_put_raw inside convert_element_type.

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).

I did not observe this prior to omnistaging, as I had tie_in-ed the NumPy constant, but tie_in was turned into a no-op.

In the good wrap_array case, jnp.array is basically functioning as a tie_in replacement.