jax.grad yields nan of jnp.sqrt(array)?

[TomClarkMassSpec]

working with 64bit and isolated a problem down to one line which I abbreviate here:

value_inside_exp = multiplier * jnp.sqrt(arr)
result = jnp.exp(value_inside_exp)

example input:

multiplier
Out[93]: DeviceArray(-3.55448731e-13, dtype=float64)

Out[94]: 
 jnp.sqrt(arr)
DeviceArray([0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             1.86344989, 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
...
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.02569974, 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.36218982,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ,
             0.        , 0.        , 0.        , 0.        , 0.        ],            dtype=float64)

So if there was problem with the value then I could use jnp.where or otherwise to retain a value at extremes. Assuming this is about extreme values, the grad returns nan. I've tried clipping the 'value_inside_exp' but that just affects the value while the grad still returns nan values in the result. Is there a way to control the grad output (ie set to zero when values or grad are really small)?

I do think nan are showing up in the gradient array due to scale of numbers as the input 'arr' above is operated on elsewhere without issue, while the variables in grad wrt variables array that deliver nan are the variables used in building this padded array.... unless I'm missing another reason for the nan's.

Every function is @jax.jit.

[jakevdp]

I don't quite understand the question – here's the code I constructed based on what you provided in your question:

import jax
import jax.numpy as jnp

jax.config.update('jax_enable_x64', True)

multiplier = jnp.float64(-3.55448731e-13)
arr = jnp.array([0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 1.86344989, 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.02569974, 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.36218982,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ,
                 0.        , 0.        , 0.        , 0.        , 0.        ])

value_inside_exp = multiplier * jnp.sqrt(arr)
result = jnp.exp(value_inside_exp)
print(result)
# [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.
#  1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.
#  1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.
#  1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.
#  1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]

This seems like reasonable output – what should I be looking for?

[TomClarkMassSpec]

I slashed and burned everything after, only returning the variables of interest:

@jax.jit
def protein_data_transformations_for_terminal_speed(custom_tree):
    # Unpack relevant pytree variables and constants

    ....   
    
    #micelle_speed_term2 = jnp.exp(-C5 * radius_micelle *... * jnp.sqrt(array_of_micelles))
    part_term_2 = (array_of_micelles)
    return part_term_2

def distance_prediction(custom_tree):
        # terminal speed of protein-SDS complex with ot without gel fiber interaction
    micelles_speed_term_2 = protein_data_transformations_for_terminal_speed(custom_tree) 
    distance_prediction = jnp.sum(micelles_speed_term_2)
    return distance_prediction

def calc_individual_y(variables_list, constants_tuple, padded_seq, protein_length):
    custom_tree = CustomPyTree(
        variables_list,
        constants_tuple,
        (
            padded_seq,
            protein_length
        )
    )
    y = distance_prediction(custom_tree)
    return y

def wrapped_y_function(variables, constants_tuple, padded_seq, protein_length):
    v0, v1, v2, v3, v4, v5, v6, v7, v8, v9, v10, v11, v12, v13, v14, v15, v16, v17, v18, v19, v20, v21, v22, v23, v24, v25, v26, v27, v28, v29 = variables
    return calc_individual_y([v0, v1, v2, v3, v4, v5, v6, v7, v8, v9, v10, v11, v12, v13, v14, v15, v16, v17, v18, v19, v20, v21, v22, v23, v24, v25, v26, v27, v28, v29], constants_tuple, padded_seq, protein_length)

# Value and gradient function
value_and_grad_fn = jax.value_and_grad(wrapped_y_function, argnums=0)

# Evaluate value and gradient at some point

#value, gradient = value_and_grad_fn(variables_at_some_point, constants_at_some_point, data_array_at_some_point)
value, gradient = value_and_grad_fn(variables_at_some_point, constants_tuple, padded_seq_at_some_point, protein_length_at_some_point)

print("Value:", value)
print("Gradient:", gradient)

[jakevdp]

There are several undefined variables in your code: array_of_micelles, CustomPyTree, variables_at_some_point, constants_tuple, padded_seq_at_some_point, protein_length_at_some_point... etc.

I'd love to help, but I can't help if I don't know what code you're running!

[TomClarkMassSpec]

I don't think I can drop a thousand lines of code. But I figured it out and will leave the answer. In short I have non-functional zeros such as due to padding to maintain array shape. Sqrt of 0. is 0. and gradient of that is inf. Further down the code, the inf. yields nan for jax.grad. The fix:

    masked_sds_cluster_quant = sdscluster_quant>0
    multiplier = -C5 * radius_micelle *...
    exponential_term_2 = jnp.where(masked_sds_cluster_quant, multiplier*jnp.sqrt(sdscluster_quant),0.)
    micelle_speed_term2 = jnp.exp(exponential_term_2)

Thanks for looking

[jakevdp]

That makes sense - glad you figured out the issue!

[keithlinndowning]

Thanks Tom. I should have searched for this 4 hours ago, but it took me that long to localize my error to jnp.sqrt.