Buffers not moving to jax device

[mitkotak]

Hey I am getting the following error in my conversion which seems to be coming from the buffers which is weird since it was able to move params (also a dict) without any issues. Wondering whether I need to use a more explicit API.

  File "/orcd/home/002/mkotak/y/envs/matbench-speed/lib/python3.11/site-packages/nequix/calculator.py", line 185, in calculate
    energy_per_atom, forces, stress = self.model_func(
                                      ^^^^^^^^^^^^^^^^
  File "/orcd/home/002/mkotak/y/envs/matbench-speed/lib/python3.11/site-packages/torchax/interop.py", line 219, in call_jax
    args, kwargs = jax_view((args, kwargs))
                   ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/orcd/home/002/mkotak/y/envs/matbench-speed/lib/python3.11/site-packages/jax/_src/tree_util.py", line 361, in tree_map
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/orcd/home/002/mkotak/y/envs/matbench-speed/lib/python3.11/site-packages/jax/_src/tree_util.py", line 361, in <genexpr>
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
                             ^^^^^^
  File "/orcd/home/002/mkotak/y/envs/matbench-speed/lib/python3.11/site-packages/torchax/interop.py", line 202, in _jax_view
    assert isinstance(t, tensor.Tensor) or isinstance(t, tensor.View), type(t)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: <class 'torch.Tensor'>

import torch
from torchax.interop import JittableModule, jax_jit, call_jax

graph = dict_to_pytorch_geometric(processed_graph)
graph.n_graph = torch.zeros(graph.x.shape[0], dtype=torch.int32).to(self.device)
graph = graph.to(self.device)
model_jittable = JittableModule(self.model).to(self.device)
self.model_func = jax_jit(functools.partial(model_jittable.functional_call, 'forward'),
                        kwargs_for_jax_jit={'donate_argnums': (0,)})
self.weights = model_jittable.params
self.buffers = model_jittable.buffers

energy_per_atom, forces, stress = self.model_func(
                self.weights,
                self.buffers,
                (graph.x,
                graph.positions,
                graph.edge_attr,
                graph.edge_index,
                getattr(graph, "cell", None),
                graph.n_node,
                graph.n_edge,
                graph.n_graph)
            )

[qihqi]

I suspect that some variable in this list:

                (graph.x,
                graph.positions,
                graph.edge_attr,
                graph.edge_index,
                getattr(graph, "cell", None),
                graph.n_node,
                graph.n_edge,
                graph.n_graph)

is a torch.Tensor (regular torch CPU tensor) instead of a torchax.tensor.Tensor (what .to('jax') would produce).

Currently we dont support mixed math with torch.Tensor and torchax.tensor.Tensor as this is the same as mixed math of CPU and CUDA tensors (which is not supported neither).

If that happens, then likely graph = graph.to(self.device) didn't move it. This usually happens if the variable is assigned to self instead of registered with self.register_buffer in torch.