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

Refactoring special function implementations #112

Merged
merged 4 commits into from
Sep 30, 2023
Merged

Refactoring special function implementations #112

merged 4 commits into from
Sep 30, 2023

Conversation

hatemhelal
Copy link
Contributor

This patch adds additional implementations of the special functions used in the integral evaluation methods. Also revisited the binom_factor implementation to avoid the fori_loop.

The motivation to provide different implementations is to allow comparing across different metrics such as:

  • jit tracing time
  • run time
  • compile time for ipu backend
  • ipu execution

A separate PR to come will introduce measurements across these metrics.

@hatemhelal hatemhelal requested a review from awf September 29, 2023 15:37
@hatemhelal hatemhelal self-assigned this Sep 29, 2023
Copy link
Collaborator

@awf awf left a comment

Choose a reason for hiding this comment

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

Looks good


@pytest.mark.parametrize("binom_func", [binom_beta, binom_fori, binom_lookup])
def test_binom(binom_func):
x = jnp.array([4, 4, 4, 4])
Copy link
Collaborator

Choose a reason for hiding this comment

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

Maybe test some edge cases too (1,1) (3,3) (3,0) (0,-1) etc

Copy link
Contributor Author

Choose a reason for hiding this comment

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

good idea, I think $3 \choose 3$ and $3 \choose 0$ should already be covered but I've added some more cases involving ones, zeros, and minus ones.

@hatemhelal hatemhelal merged commit 023e114 into main Sep 30, 2023
4 checks passed
@hatemhelal hatemhelal deleted the specials branch September 30, 2023 12:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants