Use of PyTorch layers and tensors inside a flax network?

[rosikand]

I am trying to build a spatial transformation layer inside my network.

PyTorch has great abilities for this: F.affine_grid and F.grid_sample. Would it be possible to build a flax network and in the forward pass, call these PyTorch functions and cast the outputs back to jnp arrays upon return? Would this mess up the model parameters that are supposed to be learned during training?

Thanks!

[marcvanzee]

No, that is not possible if you are planning on jitting your code (which is strongly recommended). jit will compile your code and use Tracers, which do not work with numpy arrays. Here's a short example:

import jax
import numpy as np

def sum_it(x):
  x = np.array(x)
  x = x.sum()
  return jnp.array(x)

jax.jit(sum_it)(jax.numpy.array(1))  # Tracer error!

Error trace

UnfilteredStackTrace: jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on 
the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
While tracing the function sum_it at <ipython-input-129-c2c23ebbfd77>:4 for jit, this concrete value was not available in Python 
because it depends on the value of the argument 'x'.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

Currently I do not believe JAX provides those functions, perhaps you can file an issue with them and ask what their plans are?