jax.vjp() with abstract-only primals (i.e. ShapedArray)?

[jacobhess118]

Wondering is there a way to make jax.vjp() work with abstract-only primals (i.e. ShapedArray)? I tried the following:

import jax
from jax.abstract_arrays import ShapedArray
import numpy as np

def func(x, y):
    return x * y

print(jax.make_jaxpr(
    jax.vjp(
        func,
        ShapedArray([3, 4], np.float32),
        ShapedArray([3, 4], np.float32)
    )
))

But it throws error:

Traceback (most recent call last):
  File "playground/vjp_abstract_repro.py", line 9, in <module>
    jax.vjp(
  File "jax/jax/_src/api.py", line 2512, in vjp
    return _vjp(
  File "jax/jax/_src/api.py", line 2518, in _vjp
    for arg in primals_flat: _check_arg(arg)
  File "jax/jax/_src/api.py", line 2928, in _check_arg
    raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid JAX type.")
TypeError: Argument 'ShapedArray(float32[3,4])' of type <class 'jax.core.ShapedArray'> is not a valid JAX type.

Ideally I would like to use this approach to get a Jaxpr-based graph for the backward pass of func(), which should have x, y, grad_out as input and grad_x, grad_y as output.

The reason I want to avoid needing to pass in ConcreteArray to jax.vjp() is because some inputs can be too large to fit in memory, and if there is a way to get the backward graph without needing to materialize the inputs, it would be super awesome.

Thanks!

[YouJiacheng]

You do not use jax.make_jaxpr correctly. It transform a function to another function with jaxpr as its output, and cannot transform a "expression"(Actually, python will eagerly evaluate the expression and pass the value to jax.make_jaxpr, not the expression/AST, and you code actually failed at the evaluation of the vjp expression which need concrete array)
Try following code:

import jax
import jax.numpy as jnp
from jax.abstract_arrays import ShapedArray

def func(x, y):
    return x * y

print(jax.make_jaxpr(lambda x, y: jax.vjp(func, x, y))(ShapedArray((3, 4), jnp.float32), ShapedArray((3, 4), jnp.float32)))

and output:

{ lambda ; a:f32[3,4] b:f32[3,4]. let c:f32[3,4] = mul a b in (c, b, a) }

However, you want:

Jaxpr-based graph for the backward pass of func, which should have x, y, grad_out as input and grad_x, grad_y as output.

Thus you should use following code:

import jax
import jax.numpy as jnp
from jax.abstract_arrays import ShapedArray

def func(x, y):
    return x * y

abs_arr = ShapedArray((3, 4), jnp.float32)
print(jax.make_jaxpr(lambda x, y, grad_out: jax.vjp(func, x, y)[1](grad_out))(abs_arr, abs_arr, abs_arr))

and output:

{ lambda ; a:f32[3,4] b:f32[3,4] c:f32[3,4]. let
    _:f32[3,4] = mul a b
    d:f32[3,4] = mul a c
    e:f32[3,4] = mul c b
  in (e, d) }

exactly what you need.
WDYT?