Add missing PyTorch/JAX export for logical_or, logical_and, and relu

[MilesCranmer]

Thanks to @j-thib for pointing this out, I didn't realize the SymPy maps were not built-in.

However this gets fairly complicated as we need to map sympy.Piecewise into torch/jax code.

TODO:

• Add sympy2torch for Piecewise
• Add sympy2jax for Piecewise
• Add tests for all new branches. We should reach 100% coverage on the diff because the code gets fairly convoluted.

[tbuckworth]

In order to avoid the following error:
KeyError: 'Function ITE was not found in Torch function mappings.Please add it to extra_torch_mappings in the format, e.g., {sympy.sqrt: torch.sqrt}.'

This function should be added to the mappings:

def if_then_else(*conds):
    a, b, c = conds
    return torch.where(a, torch.where(b, True, False), torch.where(c, True, False))

extra_torch_mappings = {sympy.logic.boolalg.ITE: if_then_else}

Additionally, lines 87-89 can be replaced with:

output += torch.where(
                    cond.bool() & ~already_used, expr, torch.zeros_like(expr)
                )
                already_used = already_used | cond.bool()

as when cond is a float (1.) the code fails without cond.bool()

[MilesCranmer]

If you go to the files tab, you can click the + next to a line of code where you want to add it. Then click "suggest". Then I can add your code as a commit on the PR. If you can't figure it out no worries!

[tbuckworth]

Added if_then_else as ITE mapping.

Changed cond to cond.bool() to avoid error when cond it Float (1.).