Incorrect branch handling during derivative computation for Max function

[akhadke-bdai]

Describe the bug
I am attempting to generate PyTorch code for computing the derivative/jacobian for the following function $w = \sqrt{\max(0, x^2 - z^2)}$. The expected derivative/jacobian for this function is as follows

$\frac{\partial w}{\partial x} = 0$ if $x^2 < z^2$
undefined if $x^2 = z^2$
$\frac{\partial w}{\partial x} = \frac{x}{\sqrt{x^2 - z^2}}$ if $x^2 > z^2$

However, symforce codegen computes the following
$\left(\frac{\partial w}{\partial x}\right)_{\text{codegen}} = \frac{x}{2} \left(\frac{sign(x^2 - z^2) + 1}{\sqrt{\max(0, x^2 - z^2)}}\right)$
which leads to nan values when evaluating any $x: x^2 < z^2$

To Reproduce
My script to generate the symbolic code

import os
import symforce

symforce.set_symbolic_api("symengine")
symforce.set_log_level("warning")
symforce.set_epsilon_to_symbol()

import symforce.symbolic as se
from symforce import codegen
from symforce.codegen.backends.pytorch.pytorch_config import PyTorchConfig
from symforce.values import Values

def gen_py(inputs, outputs, name):
    gen = codegen.Codegen(
        inputs=inputs,
        outputs=outputs,
        config=PyTorchConfig(),
        name=name,
    )
    data = gen.generate_function()

    # Print what we generated
    print("Files generated in {}:\n".format(data.output_dir))
    for f in data.generated_files:
        print("  |- {}".format(os.path.relpath(f, data.output_dir)))

def main():
    x = se.Symbol('x')
    z = se.Symbol('z')
    y = se.Max(se.Scalar(0.0), (x**2 - z**2))
    w = se.sqrt(y)

    dwdx = [se.diff(w, x), se.diff(w, z)]

    inputs = Values(x=se.Matrix([x, z]))
    outputs = Values(dwdx=se.Matrix(dwdx))
    gen_py(inputs=inputs, outputs=outputs, name='dwdx')

if __name__=='__main__':
    main()

Expected behavior
See issue description

Environment (please complete the following information):

• OS and version: Ubuntu 22.04.5 LTS
• Python version 3.10.12
• SymForce Version 0.10.1

[aaron-skydio]

The underlying symbolic library (sympy or symengine) is what's generating the expression for the derivative. From sympy, for example, I get:

In [1]: import sympy as sm

In [2]: x, z = sm.symbols('x z')

In [3]: sm.sqrt(sm.Max(0, x**2 - z**2)).diff(x)
Out[3]: x*Heaviside(x**2 - z**2)/sqrt(Max(0, x**2 - z**2))

And the heaviside there is the $\frac{\mathrm{sign}(x^2 - z^2) + 1}{2}$ that you see in the expression [1]. When $x^2 < z^2$, that should evaluate to $\frac{x}{2} \left(\frac{0}{\sqrt{0}}\right)$ which gives you your 0/0 singularity.

This is sorta interesting, and doesn't fit into the typical examples here. I think sqrt(max(epsilon, x^2 - z^2)) should work, like this (switching to the python backend to skip messing with pytorch for the moment):

import os
import symforce

symforce.set_symbolic_api("symengine")
symforce.set_log_level("warning")
symforce.set_epsilon_to_symbol()

import symforce.symbolic as se
from symforce import codegen
from symforce.codegen.backends.pytorch.pytorch_config import PyTorchConfig
from symforce.codegen.backends.python.python_config import PythonConfig
from symforce.values import Values

def gen_py(inputs, outputs, name):
    gen = codegen.Codegen(
        inputs=inputs,
        outputs=outputs,
        config=PythonConfig(),
        name=name,
    )
    data = gen.generate_function()

    # Print what we generated
    print("Files generated in {}:\n".format(data.output_dir))
    for f in data.generated_files:
        print("  |- {}".format(os.path.relpath(f, data.output_dir)))

def main():
    x = se.Symbol('x')
    z = se.Symbol('z')
    y = se.Max(se.epsilon(), (x**2 - z**2))
    w = se.sqrt(y)

    dwdx = [se.diff(w, x), se.diff(w, z)]

    inputs = Values(x=se.Matrix([x, z]), epsilon=se.epsilon())
    outputs = Values(dwdx=se.Matrix(dwdx))
    gen_py(inputs=inputs, outputs=outputs, name='dwdx')

if __name__=='__main__':
    main()

which produces this:

In [10]: import sym

In [11]: dwdx.dwdx(np.array([-1, -1]), sym.epsilon)
Out[11]: array([-0.,  0.])

In [12]: dwdx.dwdx(np.array([-1, 3]), sym.epsilon)
Out[12]: array([-0., -0.])

In [13]: dwdx.dwdx(np.array([10, 3]), sym.epsilon)
Out[13]: array([ 1.04828484, -0.31448545])

In [14]: dwdx.dwdx(np.array([1, 1]), sym.epsilon)
Out[14]: array([ 0., -0.])

[1] $\mathrm{sign}(0)$ is 0 in sympy which might be a separate issue but I don't think is the bug you're seeing