Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

RFC: specify jit static args via Static annotation #24705

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

jakevdp
Copy link
Collaborator

@jakevdp jakevdp commented Nov 4, 2024

An idea inspired by chats with @superbobry.

Example:

import jax
from jax.typing import Static

@jax.jit
def f(x: jax.Array, square: Static[bool]):
  return x ** 2 if square else x

print(f(2, True))  # 4
print(f(2, False))  # 2

The current way to define this would be

from functools import partial

@partial(jax.jit, static_argnames=['square'])
def f(x: jax.Array, square: bool):
  return x ** 2 if square else x

@jakevdp jakevdp self-assigned this Nov 4, 2024
@jakevdp jakevdp marked this pull request as draft November 4, 2024 23:24
@jakevdp jakevdp added the enhancement New feature or request label Nov 4, 2024
@froystig
Copy link
Member

froystig commented Nov 5, 2024

How should this behave in the absence of an immediate jit? In particular, adapting the example, what do we expect this code to do?

@jax.jit
def f(x, square):
  return g(x, square)

def g(x: jax.Array, square: Static[bool]):
  return x ** 2 if square else x

print(f(2, True))
print(f(2, False))

@jakevdp
Copy link
Collaborator Author

jakevdp commented Nov 5, 2024

How should this behave in the absence of an immediate jit? In particular, adapting the example, what do we expect this code to do?

My vision would be for the Static annotation to be ignored unless used in the definition of a function that is wrapped by jit.

@shoyer
Copy link
Collaborator

shoyer commented Nov 5, 2024

This looks great to me! It's definitely a step-up in readability from the way we currently annotate static arguments.

Copy link
Collaborator

@superbobry superbobry left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like it!

def static_argnames_from_annotations(fun: Callable[..., Any]) -> tuple[str, ...]:
try:
hints = get_type_hints(fun, include_extras=True)
except (TypeError, ValueError, NameError):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure if we should suppress NameError. IIUC this can only happen if a forward reference is local or if a type annotation references an undefined name.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I got a NameError when running the JAX test suite and this was my fix.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm just concerned that we would ignore Static annotations if we handle NameError unconditionally. Would it be a hack to check if __annotations__ doesn't contain Static[ as a substring on NameError and re-raise if that's the case?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, I'll give that a try.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just pushed a commit that tries to do smarter error handling, with detailed comments. It's a bit messy to be honest... let me know what you think.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants