Binding primitives with ad.Zero arguments

[cgel]

I would like to understand why binding primitive ops with ad.Zero doesn't work.

First, I constructed my own primitive and defined an implementation rule that properly handles ad.Zero. It all works fine. When the top trace is EvalTrace(level=0/0) then I can bind my primitive op with an argument list that contains ad.Zeros. But when the top trace is DynamicJaxprTrace(level=1/0), then binding with the same arguments gives the error

TypeError: Value Zero(ShapedArray(float32[1,128])) with type <class 'jax._src.ad_util.Zero'> is not a valid JAX type

This seems to occur because binding withDynamicJaxprTrace relies on get_aval on the ad.Zero argument, which gives the above error. I find it weird that get_aval crashes on ad.Zero given that ad.Zero literally has an attribute called aval.

My question is:

• Is there ever a circumstance where one should be binding a primitive with ad.Zero arguemnts?
• If not, why not? What is the reason to not see ad.Zero as a valid jax_type?

[jakevdp]

ad.Zero is not a valid JAX type, in the sense that creating an ad.Zero object and passing it directly to a JAX function is not allowed. Similarly, you should not directly bind primitives with it. It is a non-public object used to represent symbolic zeros within autodiff rules, and is handled consistently by JAX in that context.

Can you say more about what you're trying to accomplish when you were attempting to bind a primitive with a raw ad.Zero object?

[patrick-kidger]

As Jake says, by default JAX only supports symbolic zeros in the limited context of autodiff rules (which is a private API). Symbolic zeros are not supported in general programs.

If you'd like to use them in general programs, then try Quax instead. It even includes a toy quax.examples.zero symbolic zero library as a demonstration -- you could perhaps expand that out to suit your needs.

[davisyoshida]

Here's an example of how to get a symbolic zero from Qax:

import jax
import jax.numpy as jnp

import qax
from qax.symbols import SymbolicConstant

@qax.use_implicit_args
def f(x, y):
    return jnp.sin(x * y) + x

x = jnp.ones(3)
zero = SymbolicConstant(0, shape=(5, 3), dtype=jnp.float32)

print(f(x, zero))
# Output:
# [[1. 1. 1.]
# [1. 1. 1.]
# [1. 1. 1.]
# [1. 1. 1.]
# [1. 1. 1.]]

print(jax.make_jaxpr(f)(x, zero))
# Output:
# { lambda ; a:f32[3]. let
# _:f32[1,3] = broadcast_in_dim[broadcast_dimensions=(1,) shape=(1, 3)] a
# b:f32[1,3] = broadcast_in_dim[broadcast_dimensions=(1,) shape=(1, 3)] a
# c:f32[3] = squeeze[dimensions=(0,)] b
# d:f32[5,3] = broadcast_in_dim[broadcast_dimensions=(1,) shape=(5, 3)] c
# in (d,) }

You can see from the generated jaxpr that zero kills off the multiplication, sin, and addition so all that's left is broadcasting ops to get the right output shape. A symbolic zero is actually the first example in the README, if you want to take a look or implement similar things yourself.

If you have a specific primitive, my_primitive, and want to define how it operates with constants you could do:

@qax.primitive_handler(my_primitive)
def handler(primitive, x : SymbolicConstant):
    if x.value != 0:
        # If the value is non-zero, just create a dense array and run your primitive
        return primitive.bind(x.materialize())

    return my_special_logic_for_zero(x.shape)

[cgel]

Can you say more about what you're trying to accomplish when you were attempting to bind a primitive with a raw ad.Zero object?

I have my own primitive and I'm trying to make jvp and vjp work with it. Now I think I understand the intended way to do so. I am supposed to "intercept" ad.Zero when defining the jvp and also, if the jvp involves custom primitives, I need to make sure their transposition rule also handles ad.Zero objects. Basically, I need to write an implementation of jvp and transpose for every possible combination of ad.Zero on the input arguments. This is what I was doing wrong. I was using the arguments directly, which resulted in binding primitives with non valid jax types.

@jakevdp I'm a bit confused about the "philosophy" of ad.Zero. I'm a big fan of the way backpropagation is implemented in jax. There is a single transformation jvp that is inherently about differentiation and the rest consists of generic mathematical/computational concepts like partial evaluation and transposition, which have legitimate use outside of automatic differentiation. But killing off parts of the computation when an array of 0s is passed also feels like a very generic idea with uses outside of AD. It is very similar to optimizing the computations where there is an elementwise multiplication by 1s, adding 0s or when doing a matmul by the identity matrix. Is there any reason to not have generic objects like Zero and Identity as valid jax types? Wouldn't it remove the need to handle all combinations of ad.Zero in the arguments of the jvp and transposition rules?

Also, thanks for the pointer to Qax @patrick-kidger and @davisyoshida. It is a very interesting usage of all the tracing machinery to perform these symbolic optimizations!

[jakevdp]

Is there any reason to not have generic objects like Zero and Identity as valid jax types?

It's a useful idea, but it would come at the cost of pretty significant added complexity (every API, both primitive and non-primitive, would have to support these symbolic arrays in place of every argument, both in JAX itself and in every downstream library built on JAX). Someone would need to demonstrate that the potential benefits are worth the potential cost.