En/review eg/losse val grad

[eduardolneto]

Modifying losses.py to allow custom_losses to have a val_and_grad function instead of only a grad function. This is necessary for the augmented lagragian modifications in the parallel pull request, since in lbfgsb the val_and_grad option is more efficient and the only viable option. Changes on losses.py.

[EstevaoMGomes]

Great changes!

[rogeriojorge]

Mostly 3 minor stuff:

1. losses.py:111, custom_loss.call_pytree():
   dofs_to_pytree now returns a named dict, but call_pytree() still does self.fun(*dofs_pytree, **self.kwargs). If dofs_pytree is a dict, this passes the keys as positional args. I think this should either be:

args = tuple(dofs_pytree[name] for name in self.args_names)
return self.fun(*args, **self.kwargs)

or, if the custom loss functions are expected to accept keyword args:

return self.fun(**dofs_pytree, **self.kwargs)

2. losses.py:132-136, custom_loss.grad_pytree()

same things. It currently does *dofs_pytree, which will be wrong if dofs_pytree is the new named dict. I would build

args = tuple(dofs_pytree[name] for name in self.args_names)

and differentiate with those positional args, then put the gradients back into buffer.

3. losses.py:80, _ensure_unravelers

I think "or" is actually the right condition here. Using "and" would only rebuild when all of them are None, which could leave _dofs_to_args or _dofs_to_pytree unset. But not sure.

---

A small test for custom_loss.value_and_grad() with two dependencies, and dofs_to_pytree() followed by call_pytree() / grad_pytree() would probably catch the dict-vs-tuple issue above.

[eduardolneto]

I think I addressed all the comments for this PR.

[rogeriojorge]

Can you just add one assertion for value_and_grad() itself in test_custom_loss_named_unraveler()?
Otherwise, it's good for me to merge

[rogeriojorge]

Can you just add one assertion for value_and_grad() itself in test_custom_loss_named_unraveler()?
Otherwise, it's good for me to merge

[eduardolneto]

Added assertion for val_and_grad

[EstevaoMGomes]

After fixing the test, this PR is ready to be merged.