def update_entry(arr, val, i, j):
    val = lax.reshape(val, [1, 1])
    return lax.dynamic_update_slice(arr, val, (i, j))
def update_entry(arr, val, i, j):
    val = lax.reshape(val, [1, 1])
    return lax.dynamic_update_slice(arr, val, (i, j))