Question on abstract_eval_fun

[DarrenZhang01]

Hi, it seems that jax.interpreters.partial_eval.abstract_eval_fun (https://github.com/google/jax/blob/master/jax/interpreters/partial_eval.py#L326) is only doing shape inference at the abstract level? But why is that the case that sometimes the output avals_out will contain PartialVal objects with known values like ConcreteArray? Thanks.

[hawkinsp]

By the way, this is probably better as a GitHub discussion, not an issue. There's no action for us to take here (other than answer the question.). I suggest opening a discussion thread?

[hawkinsp]

Actually apparently I can convert it to a discussion. Let me do that!

[apaszke]

@DarrenZhang01 Even though all inputs to a partially evaluated function might be unknown, it does not imply that all outputs are unknown, because some of those outputs might not have any dependency on the inputs! A simple example would be a function like this one:

def f(x):
  return x * 2, jnp.ones((10,))

In this case, running abstract_eval_fun on f would produce one ShapedArray eval for x * 2 and one ConcreteArray aval for the jnp.ones result.

[DarrenZhang01]

Thanks, Adam @apaszke !