<p dir="auto">I'm writing a model in jax/flax and I use pytorch <code class="notranslate">utils.data.DataLoader</code> for multiprocessed data loading. The dataset functions however do not use pytorch tensors but numpy and jax, because some of the code I use for my model is also relevant for data preprocessing (various 3d geometry routines). This works with just CPU but when I run it on a GPU server I run into</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="2021-12-09 14:23:32.808066: F external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:153] Failed setting context: CUDA_ERROR_NOT_INITIALIZED: initialization error"><pre class="notranslate"><code class="notranslate">2021-12-09 14:23:32.808066: F external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:153] Failed setting context: CUDA_ERROR_NOT_INITIALIZED: initialization error
</code></pre></div>
<p dir="auto">which I believe to be because of spawning multiple CUDA-using processes. If I duplicate the geometry-related functions to use <code class="notranslate">np</code> instead of <code class="notranslate">jnp</code>, everything works fine. Ideally, I would like to have a way to instruct jax to execute all instructions in some context on CPU and avoid the code duplication. I envision it as something like</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="class Dataset:
   ...
   def _getitem_impl(self, index):
      # uses np and jnp code intermixed
      return data_item

   def __getitem__(self, index):
      with jax.default_device_placement('cpu'):
         return self._getitem_impl(index)"><pre class="notranslate"><span class="pl-k">class</span> <span class="pl-v">Dataset</span>:
   ...
   <span class="pl-k">def</span> <span class="pl-en">_getitem_impl</span>(<span class="pl-s1">self</span>, <span class="pl-s1">index</span>):
      <span class="pl-c"># uses np and jnp code intermixed</span>
      <span class="pl-k">return</span> <span class="pl-s1">data_item</span>

   <span class="pl-k">def</span> <span class="pl-en">__getitem__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">index</span>):
      <span class="pl-k">with</span> <span class="pl-s1">jax</span>.<span class="pl-en">default_device_placement</span>(<span class="pl-s">'cpu'</span>):
         <span class="pl-k">return</span> <span class="pl-s1">self</span>.<span class="pl-en">_getitem_impl</span>(<span class="pl-s1">index</span>)</pre></div>
<p dir="auto">I saw that there are other issues asking about similar features, however in my case <code class="notranslate">jax.device_put</code> won't really cut it because of the intermixing of jax and numpy code, which would force me to add <code class="notranslate">device_put</code> at every API boundary. It is also not something that can be solved with global flags because I still want the model in the main process to run on GPU. Would it be difficult to add such a feature?</p>
<p dir="auto">JAX detects my GPUs correctly and I am able to put data and parameters on the GPU, however, when I run my code, the GPU seems to idle (1% GPU utility) while the CPU is utilized at 140%.</p>
<p dir="auto">Is this behavior intended?</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from flax import linen as nn
from flax import optim
import jax.numpy as jnp
import jax

import numpy as np

def load_data():
    data = np.load('mnist.npz')
    x_train = data['x_train']
    x_train = (x_train-128.0)/128.0
    x_train = np.reshape(x_train,(60000, 28**2))
    y_train = data['y_train'][:,None]
    
    return x_train, y_train


class fnet(nn.Module):
    
    @nn.compact        
    def __call__(self, x):
        x = nn.Dense(features=25)(x)
        for i in range(10):
            x = nn.relu(nn.Dense(features=25)(x))
        return nn.Dense(features=1)(x)
    
    def initialize(self, key, x):
        return fn.init(key, x)
        
    def loss(self, x, y):
        def loss_fn(params):
            yhat = self.apply(params, x)
            loss = jnp.mean((yhat-y)**2)
            return loss
        return jax.jit(loss_fn)
            
def train_epoch(net, x_train, y_train, optimizer):

    #for i in range(len(x_train)//32):
    for i in range(300):
        
        x_tr = x_train[i*32:(i+1)*32]
        y_tr = y_train[i*32:(i+1)*32]
    
        loss = net.loss(x_tr, y_tr)
    
        v_grad = jax.value_and_grad(loss)
        val, grad = v_grad(optimizer.target)
        
        print(i, val) 
        optimizer = optimizer.apply_gradient(grad)

    return optimizer



fn = fnet()
key = jax.random.PRNGKey(1)
params = fn.initialize(key, jnp.zeros((32,28**2)))

jax.tree_map(lambda x: print(x.device_buffer.device()), params)

x_train, y_train = load_data()

rp = optim.Adam(0.0005).create(params)
x,y = jnp.array(x_train), jnp.array(y_train)

print(x.device_buffer.device(), y.device_buffer.device())
opt = train_epoch(fn, x, y, rp)"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">flax</span> <span class="pl-k">import</span> <span class="pl-s1">linen</span> <span class="pl-k">as</span> <span class="pl-s1">nn</span>
<span class="pl-k">from</span> <span class="pl-s1">flax</span> <span class="pl-k">import</span> <span class="pl-s1">optim</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">jax</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-k">def</span> <span class="pl-en">load_data</span>():
    <span class="pl-s1">data</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">load</span>(<span class="pl-s">'mnist.npz'</span>)
    <span class="pl-s1">x_train</span> <span class="pl-c1">=</span> <span class="pl-s1">data</span>[<span class="pl-s">'x_train'</span>]
    <span class="pl-s1">x_train</span> <span class="pl-c1">=</span> (<span class="pl-s1">x_train</span><span class="pl-c1">-</span><span class="pl-c1">128.0</span>)<span class="pl-c1">/</span><span class="pl-c1">128.0</span>
    <span class="pl-s1">x_train</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">reshape</span>(<span class="pl-s1">x_train</span>,(<span class="pl-c1">60000</span>, <span class="pl-c1">28</span><span class="pl-c1">**</span><span class="pl-c1">2</span>))
    <span class="pl-s1">y_train</span> <span class="pl-c1">=</span> <span class="pl-s1">data</span>[<span class="pl-s">'y_train'</span>][:,<span class="pl-c1">None</span>]
    
    <span class="pl-k">return</span> <span class="pl-s1">x_train</span>, <span class="pl-s1">y_train</span>


<span class="pl-k">class</span> <span class="pl-s1">fnet</span>(<span class="pl-s1">nn</span>.<span class="pl-v">Module</span>):
    
    <span class="pl-en">@<span class="pl-s1">nn</span>.<span class="pl-s1">compact</span></span>        
    <span class="pl-k">def</span> <span class="pl-en">__call__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>):
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Dense</span>(<span class="pl-s1">features</span><span class="pl-c1">=</span><span class="pl-c1">25</span>)(<span class="pl-s1">x</span>)
        <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-c1">10</span>):
            <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-en">relu</span>(<span class="pl-s1">nn</span>.<span class="pl-v">Dense</span>(<span class="pl-s1">features</span><span class="pl-c1">=</span><span class="pl-c1">25</span>)(<span class="pl-s1">x</span>))
        <span class="pl-k">return</span> <span class="pl-s1">nn</span>.<span class="pl-v">Dense</span>(<span class="pl-s1">features</span><span class="pl-c1">=</span><span class="pl-c1">1</span>)(<span class="pl-s1">x</span>)
    
    <span class="pl-k">def</span> <span class="pl-en">initialize</span>(<span class="pl-s1">self</span>, <span class="pl-s1">key</span>, <span class="pl-s1">x</span>):
        <span class="pl-k">return</span> <span class="pl-s1">fn</span>.<span class="pl-en">init</span>(<span class="pl-s1">key</span>, <span class="pl-s1">x</span>)
        
    <span class="pl-k">def</span> <span class="pl-en">loss</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>, <span class="pl-s1">y</span>):
        <span class="pl-k">def</span> <span class="pl-en">loss_fn</span>(<span class="pl-s1">params</span>):
            <span class="pl-s1">yhat</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">apply</span>(<span class="pl-s1">params</span>, <span class="pl-s1">x</span>)
            <span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">mean</span>((<span class="pl-s1">yhat</span><span class="pl-c1">-</span><span class="pl-s1">y</span>)<span class="pl-c1">**</span><span class="pl-c1">2</span>)
            <span class="pl-k">return</span> <span class="pl-s1">loss</span>
        <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">loss_fn</span>)
            
<span class="pl-k">def</span> <span class="pl-en">train_epoch</span>(<span class="pl-s1">net</span>, <span class="pl-s1">x_train</span>, <span class="pl-s1">y_train</span>, <span class="pl-s1">optimizer</span>):

    <span class="pl-c">#for i in range(len(x_train)//32):</span>
    <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-c1">300</span>):
        
        <span class="pl-s1">x_tr</span> <span class="pl-c1">=</span> <span class="pl-s1">x_train</span>[<span class="pl-s1">i</span><span class="pl-c1">*</span><span class="pl-c1">32</span>:(<span class="pl-s1">i</span><span class="pl-c1">+</span><span class="pl-c1">1</span>)<span class="pl-c1">*</span><span class="pl-c1">32</span>]
        <span class="pl-s1">y_tr</span> <span class="pl-c1">=</span> <span class="pl-s1">y_train</span>[<span class="pl-s1">i</span><span class="pl-c1">*</span><span class="pl-c1">32</span>:(<span class="pl-s1">i</span><span class="pl-c1">+</span><span class="pl-c1">1</span>)<span class="pl-c1">*</span><span class="pl-c1">32</span>]
    
        <span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-s1">net</span>.<span class="pl-en">loss</span>(<span class="pl-s1">x_tr</span>, <span class="pl-s1">y_tr</span>)
    
        <span class="pl-s1">v_grad</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">value_and_grad</span>(<span class="pl-s1">loss</span>)
        <span class="pl-s1">val</span>, <span class="pl-s1">grad</span> <span class="pl-c1">=</span> <span class="pl-en">v_grad</span>(<span class="pl-s1">optimizer</span>.<span class="pl-s1">target</span>)
        
        <span class="pl-en">print</span>(<span class="pl-s1">i</span>, <span class="pl-s1">val</span>) 
        <span class="pl-s1">optimizer</span> <span class="pl-c1">=</span> <span class="pl-s1">optimizer</span>.<span class="pl-en">apply_gradient</span>(<span class="pl-s1">grad</span>)

    <span class="pl-k">return</span> <span class="pl-s1">optimizer</span>



<span class="pl-s1">fn</span> <span class="pl-c1">=</span> <span class="pl-en">fnet</span>()
<span class="pl-s1">key</span> <span class="pl-c1">=</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">params</span> <span class="pl-c1">=</span> <span class="pl-s1">fn</span>.<span class="pl-en">initialize</span>(<span class="pl-s1">key</span>, <span class="pl-s1">jnp</span>.<span class="pl-en">zeros</span>((<span class="pl-c1">32</span>,<span class="pl-c1">28</span><span class="pl-c1">**</span><span class="pl-c1">2</span>)))

<span class="pl-s1">jax</span>.<span class="pl-en">tree_map</span>(<span class="pl-k">lambda</span> <span class="pl-s1">x</span>: <span class="pl-en">print</span>(<span class="pl-s1">x</span>.<span class="pl-s1">device_buffer</span>.<span class="pl-en">device</span>()), <span class="pl-s1">params</span>)

<span class="pl-s1">x_train</span>, <span class="pl-s1">y_train</span> <span class="pl-c1">=</span> <span class="pl-en">load_data</span>()

<span class="pl-s1">rp</span> <span class="pl-c1">=</span> <span class="pl-s1">optim</span>.<span class="pl-v">Adam</span>(<span class="pl-c1">0.0005</span>).<span class="pl-en">create</span>(<span class="pl-s1">params</span>)
<span class="pl-s1">x</span>,<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>(<span class="pl-s1">x_train</span>), <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>(<span class="pl-s1">y_train</span>)

<span class="pl-en">print</span>(<span class="pl-s1">x</span>.<span class="pl-s1">device_buffer</span>.<span class="pl-en">device</span>(), <span class="pl-s1">y</span>.<span class="pl-s1">device_buffer</span>.<span class="pl-en">device</span>())
<span class="pl-s1">opt</span> <span class="pl-c1">=</span> <span class="pl-en">train_epoch</span>(<span class="pl-s1">fn</span>, <span class="pl-s1">x</span>, <span class="pl-s1">y</span>, <span class="pl-s1">rp</span>)</pre></div>
<p dir="auto">(I am aware that using MSE loss is stupid)</p>