jax.debug.print doesn't print anything

[snat1505027]

I am trying to print the intermediate ShapedArray denom_logits with jax.debug.print but it doesn't print anything when I run the function.

def train_loss_fn(state, params, batch):
    preds = state.apply_fn({'params': params}, batch)
    loss_info = {}

    for c_type, c_dict in preds.items():
        numer_logits = (c_dict['x'] * c_dict['y']).sum(-1) # (945)
        loss_info[c_type] = 0.0

        # For both directions (average the result)
        for k1, k2 in ['xy', 'yx']:
            x = c_dict[k1] # (945,768)
            y = c_dict[k2] # (945,768)

            y_allgather = jax.lax.all_gather(y, 'batch').reshape(-1, x.shape[-1]) # (7560, 768)

            print(f"{c_type} {k1}->{k2} dot product sim:  shaped [{x.shape}] -> [{y_allgather.shape}]", flush=True) 
            
            denom_logits = jnp.einsum('lh,vh->lv', x, y_allgather) # (945,7560)
            jax.debug.print("denom_logits: {}", denom_logits)
            
            dlse_key = '_'+c_type+'_'+k1+'_'+k2+'_'+'denom_logits'
            loss_info[dlse_key] = denom_logits
            
            denom_lse = jax.nn.logsumexp(denom_logits.astype(jnp.float32), axis=-1) # (945)
            loss_info[c_type] += (denom_lse - numer_logits).mean() / 2.0 
            
            dlse_key = '_'+c_type+'_'+k1+'_'+k2+'_'+'denom_lse'
            loss_info[dlse_key] = denom_lse
            dlse_key = '_'+c_type+'_'+k1+'_'+k2+'_'+'numer_logits'
            loss_info[dlse_key] = numerologist
            
    loss = sum([v for k, v in loss_info.items() if not k.startswith('_')])
    return loss, loss_info

[sharadmv]

It's hard to tell what's going here from just the provided snippet. Could you raise an issue with a minimal repro and platform/version?

Also is this function being jit-ted?