Fix useless instantiation of zeros

[xalelax]

This line is problematic. We are instantiating zeros and then don't do anything with them. Why, you might ask 🙂

Let's start with the basics: jax has a funny way of keeping track of which arguments are traced in multi-input functions. Have a look at it yourself: in test_grad, when argnums is either 0 or 1, if you look at which primals and tangents are sent to, say, the tesseract_dispatch_jvp_rule function, you will see that only one argument is non-static (i.e., len(in_args)==1 and len(tan_args)==1), even if the tesseract we are using in the test accepts multiple inputs.

However, when the gradient of a jitted function is calculated, like we do in test_grad_jit, even if we try to do the same thing as above we can see that we are passing around two primals and two tangents. Here's what my debugger shows inside tesseract_dispatch_jvp_rule for argnums=2:

>>> tan_args
(Zero(ShapedArray(float32[2])), Traced<ShapedArray(float32[2])>with<DynamicJaxprTrace(level=3/0)>)

i.e.: we have a tuple of two vectors -- one of which, the one we are not using in differentiation, marked as a symbolic zero.

We could make use of this info when we build jvp_inputs in the TesseractClient.jacobian_vector_product... but we don't.
Even if Tesseracts would support calculating jvp only w.r.t. some inputs, this forces the client to always ask for derivatives w.r.t. everything.

[dionhaefner]

Note: we now do the same for outputs here, with the same opportunities for optimization.

[jpbrodrick89]

This could also help remove nans when using the tesseract-core vectoradd_jax Tesseract (i.e. without safe norm). For example, the follow returns an array of nans because "vector_min" has nans which affects the result even though it's not used directly.:

t = Tesseract.from_image("vectoradd_jax")
t.serve()

# Run it with JAX
x = jnp.ones((1000,))
y = jnp.ones((1000,))

def vector_sum(x, y):
    res = apply_tesseract(t, {"a": {"v": x}, "b": {"v": y}})
    return res["vector_add"]["result"].sum()

jax.grad(vector_sum)(x, y)

Note this issue is inherent to the Tesseract, as the following returns the same result:

t.vector_jacobian_product(
    inputs= {"a": {"v": x}, "b": {"v": y}}, 
    vjp_inputs=["a.v", "b.v"],
    vjp_outputs=["vector_min.normed_result", "vector_add.result"],
    cotangent_vector={
        "vector_min.normed_result": jnp.zeros(1000), 
        "vector_add.result":  jnp.ones(1000)
    })

Hoewver, if we are able to detect which cotangent vectors are all (non-instantiated?) zeros then we can simplify it to the following which returns the expect array of ones:

t.vector_jacobian_product(
    inputs= {"a": {"v": x}, "b": {"v": y}},
    vjp_inputs=["a.v", "b.v"],
    vjp_outputs=["vector_add.result"],
    cotangent_vector={"vector_add.result": jnp.ones(1000)}
    )

Also, see jax-ml/jax#1052 (comment) for some relevant discussion.