Jax grad complains about ints but my data contains no ints

[pabloduque0]

I have a silly question. I am getting the following error:

TypeError: grad requires real- or complex-valued inputs (input dtype that is a sub-dtype of np.inexact), but got int32. If you want to use Boolean- or integer-valued inputs, use vjp or set allow_int to True.

Which I can easily avoid by just setting allow_int=True in jax.value_and_grad. But I am wondering where these ints are coming from since all my data matrices (data and labels) are dtype float32.

Here is a scrapy version of the code that is currently failing:

def apply_model_grads(state, batch_images, batch_labels):
  def dice_coefficient_loss(state, smooth=.0000001):
    prediction = state.apply_fn(state.params, batch_images)
    y_true_flat = jnp.ravel(batch_labels)
    y_pred_flat = jnp.ravel(prediction)
    intersection = jnp.sum(y_true_flat * y_pred_flat)
    return -(2. * intersection + smooth) / (jnp.sum(y_true_flat) + jnp.sum(y_pred_flat) + smooth), prediction
  grad_function = jax.value_and_grad(dice_coefficient_loss, has_aux=True)
  (train_loss, logits), grads = grad_function(state)
  return grads, train_loss

Here is a jax.tree_map(lambda x: x.dtype, state.params) just in case it provides more info (where state is a flax.train_state). I initially thought the ints where somewhere in my params pytree but they are not:

FrozenDict({
    params: {
        concat_back_dim_1: {
            bias: dtype('float32'),
            kernel: dtype('float32'),
        },
        concat_back_dim_2: {
            bias: dtype('float32'),
            kernel: dtype('float32'),
        },
        concat_back_dim_3: {
            bias: dtype('float32'),
            kernel: dtype('float32'),
        },
        layers_down_0: {
            blocks_0: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
            blocks_1: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
            downsample: {
                normalization: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                reduction: {
                    kernel: dtype('float32'),
                },
            },
        },
        layers_down_1: {
            blocks_0: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
            blocks_1: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
            downsample: {
                normalization: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                reduction: {
                    kernel: dtype('float32'),
                },
            },
        },
        layers_down_2: {
            blocks_0: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
            blocks_1: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
            downsample: {
                normalization: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                reduction: {
                    kernel: dtype('float32'),
                },
            },
        },
        layers_down_3: {
            blocks_0: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
            blocks_1: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
        },
        layers_up_0: {
            expand: {
                kernel: dtype('float32'),
            },
            normalization: {
                bias: dtype('float32'),
                scale: dtype('float32'),
            },
        },
        layers_up_1: {
            blocks_0: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
            blocks_1: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
            upsample: {
                expand: {
                    kernel: dtype('float32'),
                },
                normalization: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
            },
        },
        layers_up_2: {
            blocks_0: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
            blocks_1: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
            upsample: {
                expand: {
                    kernel: dtype('float32'),
                },
                normalization: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
            },
        },
        layers_up_3: {
            blocks_0: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
            blocks_1: {
                mlp: {
                    fc1: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    fc2: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                },
                normalization1: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                normalization2: {
                    bias: dtype('float32'),
                    scale: dtype('float32'),
                },
                window_attn: {
                    proj: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    qkv: {
                        bias: dtype('float32'),
                        kernel: dtype('float32'),
                    },
                    relative_position_bias_table: dtype('float32'),
                },
            },
        },
        norm: {
            bias: dtype('float32'),
            scale: dtype('float32'),
        },
        norm_up: {
            bias: dtype('float32'),
            scale: dtype('float32'),
        },
        output: {
            kernel: dtype('float32'),
        },
        patch_embed: {
            norm: {
                bias: dtype('float32'),
                scale: dtype('float32'),
            },
            proj: {
                bias: dtype('float32'),
                kernel: dtype('float32'),
            },
        },
        up: {
            expand: {
                kernel: dtype('float32'),
            },
            normalization: {
                bias: dtype('float32'),
                scale: dtype('float32'),
            },
        },
    },
})

[jakevdp]

I believe flax.TrainState includes a step parameter which is an integer. If you're attempting to take the gradient with respect to the state, that would cause this error. Perhaps you meant to take the gradient with respect to just state.params?

[pabloduque0]

Oh you are right, that is an error there. Apologies for the silly one :)