Debugging a "ComplexWarning: Casting complex values to real discards the imaginary part"

[GeoffNN]

I'm implementing a Fourier Neural Operator model in jax, and am getting this warning. I've been trying to debug this line by line to figure out where the implicit casting happens, but I can't seem to quite get it. Is there anyway to have this warning point to the line of code that triggers it?

Here's my model code.

from typing import List
import jax
import jax.numpy as jnp
import flax.linen as nn
import tensorflow as tf
import numpy as np

def compl_mul2d(a, b):
    # ( in_channel, x,y,t ), (in_channel, out_channel, x,y,t) -> (out_channel, x,y,t)
    return jnp.einsum("ixy,ioxy->oxy", a, b)


################################################################
# 2d Fourier Layer
################################################################


class SpectralConv2d(nn.Module):
    in_channels: int
    out_channels: int
    mode1: int
    mode2: int

    @nn.compact
    def __call__(self, x):
        def initializer(key):
            scale = 1 / (self.in_channels * self.out_channels)
            key_real, key_imaginary, key = jax.random.split(key, 3)
            real = jax.random.uniform(
                key_real,
                (
                    2,
                    self.in_channels,
                    self.out_channels,
                    self.mode1,
                    self.mode2,
                ),
            )
            imaginary = jax.random.uniform(
                key_imaginary,
                (
                    2,
                    self.in_channels,
                    self.out_channels,
                    self.mode1,
                    self.mode2,
                ),
            )
            return scale * jax.lax.complex(real, imaginary)

        weights = self.param("weights", initializer)
        weights1 = weights[0]
        weights2 = weights[1]

        # Compute Fourier coeffcients up to factor of e^(- something constant)
        x_ft = jnp.fft.rfftn(x, axes=[1, 2])

        # Removed gridy, since it's not used in the training code

        # Multiply relevant Fourier modes
        out_ft = jnp.zeros(
            (self.out_channels, x.shape[-2], x.shape[-1] // 2 + 1),
            dtype=jnp.complex64,
        )
        out_ft = out_ft.at[:, : self.mode1, : self.mode2].set(
            compl_mul2d(x_ft[:, : self.mode1, : self.mode2], weights1)
        )
        out_ft = out_ft.at[:, -self.mode1 :, : self.mode2].set(
            compl_mul2d(x_ft[:, -self.mode1 :, : self.mode2], weights2)
        )

        # Return to physical space
        x = jnp.fft.irfftn(out_ft, s=(x.shape[-2], x.shape[-1]), axes=[1, 2])

        return x


#######################################
# Fourier Neural Operator in 2d
#######################################


class FNN2D(nn.Module):
    modes1: List[int]
    modes2: List[int]
    layers: List[int]
    fc_dim: int = 128
    in_dim: int = 3
    out_dim: int = 1
    activation: str = "tanh"

    def setup(self):

        self.sp_convs = [
            SpectralConv2d(in_size, out_size, mode1, mode2)
            for in_size, out_size, mode1, mode2 in zip(
                self.layers, self.layers[1:], self.modes1, self.modes2
            )
        ]
        self.convs = [
            nn.Conv(out_size, [1], padding="VALID")
            for out_size in self.layers[1:]
        ]

        self.fc0 = nn.Dense(self.layers[0])
        self.fc1 = nn.Dense(self.fc_dim)
        self.fc2 = nn.Dense(self.out_dim)

        if self.activation == "tanh":
            self.activation_fn = nn.tanh
        elif self.activation == "gelu":
            self.activation_fn = nn.gelu
        elif self.activation == "relu":
            self.activation_fn = nn.relu
        else:
            raise ValueError(f"{self.activation} is not supported.")

    def __call__(self, x):
        """
        Args:
          x: (x_grid, y_grid, 2)
        Returns:
          prediction: (x_grid, y_grid, 1)
        """

        size_x = x.shape[0]
        size_y = x.shape[1]

        n_convs = len(self.convs)
        import pdb

        pdb.set_trace()
        x = self.fc0(x)
        x = x.transpose((2, 0, 1))

        for i, (speconv, conv) in enumerate(zip(self.sp_convs, self.convs)):
            x1 = speconv(x)
            x2 = conv(x.reshape((1, -1, self.layers[i])))
            x2 = x2.reshape((self.layers[i + 1], size_x, size_y))
            x = x1 + x2
            if i != n_convs - 1:
                x = self.activation_fn(x)

        x = x.transpose((1, 2, 0))
        x = self.fc1(x)
        x = self.fc2(x)
        return x

The code is a translation of https://github.com/devzhk/PINO/tree/master/models : files basics.py and fourier2d.py.

[jakevdp]

Can you add a main() function to the example that includes the call that generates the warning?

[GeoffNN]

Here's the full reproducing script.

from typing import List
import jax
import jax.numpy as jnp
import flax.linen as nn
import tensorflow as tf
import optax
import numpy as np

# For FNO

gpus = tf.config.experimental.list_physical_devices("GPU")
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)


def compl_mul2d(a, b):
    # ( in_channel, x,y,t ), (in_channel, out_channel, x,y,t) -> (out_channel, x,y,t)
    return jnp.einsum("ixy,ioxy->oxy", a, b)


################################################################
# 2d Fourier Layer
################################################################


class SpectralConv2d(nn.Module):
    in_channels: int
    out_channels: int
    mode1: int
    mode2: int

    @nn.compact
    def __call__(self, x):
        def initializer(key):
            scale = 1 / (self.in_channels * self.out_channels)
            key_real, key_imaginary, key = jax.random.split(key, 3)
            real = jax.random.uniform(
                key_real,
                (
                    2,
                    self.in_channels,
                    self.out_channels,
                    self.mode1,
                    self.mode2,
                ),
            )
            imaginary = jax.random.uniform(
                key_imaginary,
                (
                    2,
                    self.in_channels,
                    self.out_channels,
                    self.mode1,
                    self.mode2,
                ),
            )
            return scale * jax.lax.complex(real, imaginary)

        weights = self.param("weights", initializer)
        weights1 = weights[0]
        weights2 = weights[1]

        # Compute Fourier coeffcients up to factor of e^(- something constant)
        x_ft = jnp.fft.rfftn(x, axes=[1, 2])

        # Removed gridy, since it's not used in the training code

        # Multiply relevant Fourier modes
        out_ft = jnp.zeros(
            (self.out_channels, x.shape[-2], x.shape[-1] // 2 + 1),
            dtype=jnp.complex64,
        )
        out_ft = out_ft.at[:, : self.mode1, : self.mode2].set(
            compl_mul2d(x_ft[:, : self.mode1, : self.mode2], weights1)
        )
        out_ft = out_ft.at[:, -self.mode1 :, : self.mode2].set(
            compl_mul2d(x_ft[:, -self.mode1 :, : self.mode2], weights2)
        )

        # Return to physical space
        x = jnp.fft.irfftn(out_ft, s=(x.shape[-2], x.shape[-1]), axes=[1, 2])

        return x


#######################################
# Fouriel Neural Operator in 2d
#######################################


class FNN2D(nn.Module):
    modes1: List[int]
    modes2: List[int]
    layers: List[int]
    fc_dim: int = 128
    in_dim: int = 3
    out_dim: int = 1
    activation: str = "tanh"

    def setup(self):

        self.sp_convs = [
            SpectralConv2d(in_size, out_size, mode1, mode2)
            for in_size, out_size, mode1, mode2 in zip(
                self.layers, self.layers[1:], self.modes1, self.modes2
            )
        ]
        self.convs = [
            nn.Conv(out_size, [1], padding="VALID")
            for out_size in self.layers[1:]
        ]

        self.fc0 = nn.Dense(self.layers[0])
        self.fc1 = nn.Dense(self.fc_dim)
        self.fc2 = nn.Dense(self.out_dim)

        if self.activation == "tanh":
            self.activation_fn = nn.tanh
        elif self.activation == "gelu":
            self.activation_fn = nn.gelu
        elif self.activation == "relu":
            self.activation_fn = nn.relu
        else:
            raise ValueError(f"{self.activation} is not supported.")

    def __call__(self, x):
        """
        Args:
          x: (x_grid, y_grid, 2)
        Returns:
          prediction: (x_grid, y_grid, 1)
        """

        size_x = x.shape[0]
        size_y = x.shape[1]

        n_convs = len(self.convs)
        x = self.fc0(x)
        x = x.transpose((2, 0, 1))

        for i, (speconv, conv) in enumerate(zip(self.sp_convs, self.convs)):
            x1 = speconv(x)
            x2 = conv(x.reshape((1, -1, self.layers[i])))
            x2 = x2.reshape((self.layers[i + 1], size_x, size_y))
            x = x1 + x2
            if i != n_convs - 1:
                x = self.activation_fn(x)

        x = x.transpose((1, 2, 0))
        x = self.fc1(x)
        x = self.fc2(x)
        return x

#####################
### Main Script here  ###
#####################

rng = jax.random.PRNGKey(0)
x = jnp.ones((61, 61, 3))
model = FNN2D([20] * 4, [20] * 4, [64] * 5)
params = model.init(rng, x)

def forward(params, x):
    return model.apply(params, x).sum()

grad_fn = jax.grad(forward)

rng = jax.random.PRNGKey(0)
lr = 1e-3

scheduler = optax.exponential_decay(
        init_value=lr,  # initial LR
        transition_steps=130,
        decay_rate=.9,
        staircase=True,
        end_value=1e-4,
    )
opt = optax.chain(
        optax.clip_by_global_norm(
            1.0
        ),  # Clip by the gradient by the global norm
        optax.scale_by_adam(),  # Use updates from Adam
        optax.scale_by_schedule(scheduler),  # use the LR from the scheduler
        optax.scale(
            -1.0
        ),  # Scale updates by -1 since optax.apply_updates is additive and we want to descend on the loss
    )
opt_state = opt.init(params)

# Catch warnings as errors
import warnings
warnings.filterwarnings("error")

for _ in range(50):
    key , rng = jax.random.split(rng)
    x = jax.random.uniform(rng, (61,61,3))
    model_grad = grad_fn(params, x)
    params_update, opt_state = opt.update(model_grad, opt_state)
    params = optax.apply_updates(params, params_update)

jax.tree_map(lambda x: x.dtype, params_update)

jax.tree_map(lambda x: x.dtype, params)

[GeoffNN]

The above script produces the following traceback on Colab.

---------------------------------------------------------------------------
ComplexWarning                            Traceback (most recent call last)
<ipython-input-11-b5cca6539384> in <module>()
      4     model_grad = grad_fn(params, x)
      5     params_update, opt_state = opt.update(model_grad, opt_state)
----> 6     params = optax.apply_updates(params, params_update)

6 frames
/usr/local/lib/python3.7/dist-packages/optax/_src/update.py in apply_updates(params, updates)
     42   return jax.tree_multimap(
     43       lambda p, u: jnp.asarray(p + u).astype(jnp.asarray(p).dtype),
---> 44       params, updates)
     45 
     46 

/usr/local/lib/python3.7/dist-packages/jax/_src/tree_util.py in tree_map(f, tree, is_leaf, *rest)
    176   leaves, treedef = tree_flatten(tree, is_leaf)
    177   all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
--> 178   return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
    179 
    180 tree_multimap = tree_map

/usr/local/lib/python3.7/dist-packages/jax/_src/tree_util.py in <genexpr>(.0)
    176   leaves, treedef = tree_flatten(tree, is_leaf)
    177   all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
--> 178   return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
    179 
    180 tree_multimap = tree_map

/usr/local/lib/python3.7/dist-packages/optax/_src/update.py in <lambda>(p, u)
     41   """
     42   return jax.tree_multimap(
---> 43       lambda p, u: jnp.asarray(p + u).astype(jnp.asarray(p).dtype),
     44       params, updates)
     45 

/usr/local/lib/python3.7/dist-packages/jax/_src/numpy/lax_numpy.py in _astype(arr, dtype)
   6492 def _astype(arr, dtype):
   6493   lax._check_user_dtype_supported(dtype, "astype")
-> 6494   return lax.convert_element_type(arr, dtype)
   6495 
   6496 

/usr/local/lib/python3.7/dist-packages/jax/_src/lax/lax.py in convert_element_type(operand, new_dtype)
    449   if hasattr(operand, '__jax_array__'):
    450     operand = operand.__jax_array__()
--> 451   return _convert_element_type(operand, new_dtype, weak_type=False)
    452 
    453 def _convert_element_type(operand: Array, new_dtype: Optional[DType] = None,

/usr/local/lib/python3.7/dist-packages/jax/_src/lax/lax.py in _convert_element_type(operand, new_dtype, weak_type)
    464       not dtypes.issubdtype(new_dtype, np.complexfloating)):
    465     msg = "Casting complex values to real discards the imaginary part"
--> 466     warnings.warn(msg, np.ComplexWarning, stacklevel=2)
    467 
    468   # Python has big integers, but convert_element_type(2 ** 100, np.float32) need

ComplexWarning: Casting complex values to real discards the imaginary part

[GeoffNN]

I managed to narrow this down to the optax scheduler. It seems like using optax.adam works fine. I'm opening an issue on the optax github. google-deepmind/optax#287

[GeoffNN]

Edit: just noticed I wasn't conjugating the gradients before passing them to opt.update -- although this doesn't change the type problems above.

[jakevdp]

If you have opinions about conjugation of gradients, you might weigh in on #4891

[shakes76]

Looks like Complex valued optimization has been recently proposed for Optax (https://gist.github.com/wdphy16/118aef6fb5f82c49790d7678cf87da29) and may have been incorporated into the experimental module officially (https://optax.readthedocs.io/en/latest/api.html#experimental) using the split_real_and_imaginary approach. I haven't tried it yet though.

[YouJiacheng]

Looks like optax's problem.