<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Describe your goal, and if possible provide a code snippet with a motivating example.</li>
</ul>
<p dir="auto">This grew oddly specific, but really I just want to make a generalized function benchmark, we give it data and functions and a maximum timeout, it gives us the grid of results = function(data) and elapsed times, and stops the functions if they run longer than the maximum timeout. Could be cool if we could jit-compile that sort of "outer loop" since longer function runtimes would seem worse for users and the environment and faster benchmarks could also yield better feedback on performance tuning both Jax and stuff written with Jax.</p>
<p dir="auto">Inspired by <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="3670702" data-permission-text="Title is private" data-url="https://github.com/google/jax/discussions/8477" data-hovercard-type="discussion" data-hovercard-url="/google/jax/discussions/8477/hovercard" href="https://github.com/google/jax/discussions/8477">#8477</a> with <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jakevdp/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jakevdp">@jakevdp</a> - Here's how it might look. 1st benchmark is python. 2nd benchmark w/ jit doesn't work now but i wish i could figure out how to make it work)</p>
<p dir="auto">How could we make it easier to compile benchmarks like these and work with elapsed time / function timeouts inside jax.jit?</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&quot;goal: play with compiled benchmarks for compiled functions in jax&quot;
from functools import partial
from itertools import product
from typing import Callable, List, Tuple
from time import perf_counter

from func_timeout import func_timeout
from pydantic import BaseSettings
import pandas as pd
import jax

jnp = jax.numpy
jit = jax.jit

# config
class Environment(BaseSettings):
    n_x: List[int] = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
    min_weight: float = -1.0
    max_weight: float = 1.0
    n_trials: int = 10
    timeout_seconds: int = 30
    seed: int = 42


env = Environment()

# types
PRNGKey = jax._src.random.PRNGKey
Array = jnp.ndarray


# functions - 'benchmark'
def benchmark_no_jit(
    x: Array,
    functions: List[Callable],
    timeout_seconds: int = env.timeout_seconds,
    min_weight: float = env.min_weight,
    max_weight: float = env.max_weight,
) -&gt; Tuple[Array, Array]:
    n_x = x.shape[0]
    n_functions = len(functions)
    outputs = jnp.empty((n_functions + 1, n_x)).at[-1, :].set(1.0)
    timings = jnp.empty(n_functions)
    for i, fun in enumerate(functions):
        try:
            # Without JIT, we can use func_timeout to set timeouts
            # and measure runtime with time.perf_counter
            start = perf_counter()
            y = func_timeout(timeout_seconds, fun, args=(x,))
            elapsed = perf_counter() - start
            y = y if isinstance(y, Array) and y.shape[0] == n_x else 0
        except Exception as e:
            print(e)
            y, elapsed = 0, timeout
        outputs = outputs.at[i, :].set(y)
        timings = timings.at[i].set(elapsed)
    outputs = jnp.nan_to_num(
        jnp.clip(outputs, a_min=min_weight, a_max=max_weight)
    ).transpose()
    return outputs, timings


# something like this...
@partial(jit, static_argnums=(0,1))
def _benchmark_with_jit(
    n_x: int,
    n_functions: int,
    x: Array,
    functions: List[Callable],
    timeout_seconds: int = env.timeout_seconds,
    min_weight: float = env.min_weight,
    max_weight: float = env.max_weight,
) -&gt; Tuple[Array, Array]:
    outputs = jnp.empty((n_functions + 1, n_x)).at[-1, :].set(1.0)
    timings = jnp.empty(n_functions)
    for i, fun in enumerate(functions):
        try:
            with jax.profiler.trace() as trace:
                # 1: set function timeouts / time limits (otherwise, search time is infinite)
                y = trace.timeout(timeout_seconds, fun, args=(x,))
                # 2: measure how long functions take to run (to find faster functions)
                elapsed = trace.elapsed()
                y = y if isinstance(y, Array) and y.shape[0] == n_x else 0
        except Exception as _:
            y, elapsed = 0, timeout
        outputs = outputs.at[i, :].set(y)
        timings = timings.at[i].set(elapsed)
    outputs = jnp.nan_to_num(
        jnp.clip(outputs, a_min=min_weight, a_max=max_weight)
    ).transpose()
    return outputs, timings


def f1(x: Array) -&gt; Array:
    return x ** 2 - 1


def f2(x: Array) -&gt; Array:
    return x / 2


def f3(x: Array) -&gt; Array:
    return x ** 3 * jnp.sin(x)


functions = [jit(f) for f in (f1, f2, f3)]

# because we can
def benchmark_benchmark(
    benchmark: Callable = benchmark_no_jit,
    words: str = &quot;Python + func_timeout + time.perf_counter&quot;,
    functions: List[Callable] = functions,
    N_x: List[int] = env.n_x,
    seed: int = env.seed,
) -&gt; float:
    n_functions = len(functions)
    begin = perf_counter()
    rng = PRNGKey(seed)
    df = []  # i admit to sucking at pandas
    for i, n_x in product(range(env.n_trials), N_x):
        rng, key = jax.random.split(rng)
        x = jax.random.normal(key, (n_x,))
        if isinstance(benchmark, dict):
            print(&quot;n_x&quot;, n_x)
            benchmark = benchmark[n_x]
            print(&quot;benchmark&quot;, benchmark)
            print(&quot;x&quot;, x)
        outputs, timings = benchmark(x, functions)
        # print(timings)
        # print(outputs.shape)
        # print(i, n_x, timings)
        assert outputs.shape == (n_x, n_functions + 1)
        assert timings.shape == (n_functions,)
        assert (timings &gt; 0).all()
        assert (timings &lt; env.timeout).all()
        assert (env.min_weight &lt;= outputs).all()
        assert (env.max_weight &gt;= outputs).all()
        # print(f&quot;All assertions passed.&quot;)
        df.append(
            {
                &quot;n_x&quot;: n_x,
                &quot;i&quot;: i,
                **{f&quot;f{i+1} elapsed&quot;: t for i, t in enumerate(timings)},
            }
        )
    df = pd.DataFrame(df).set_index([&quot;n_x&quot;, &quot;i&quot;])
    elapsed = perf_counter() - begin
    print(df)
    print(f&quot;{words} took {elapsed:.2f} seconds&quot;)
    return elapsed


def test_jax_benchmarks():
    elapsed_no_jit = benchmark_benchmark()
    elapsed_with_jit = benchmark_benchmark(
        benchmark={
            n: partial(_benchmark_with_jit, n_x=n, n_functions=len(functions))
            for n in env.n_x
        },
        words=&quot;XLA (jax.jit) + jax.profiler.trace&quot;,
    )
    print(&quot;done!&quot;)


if __name__ == &quot;__main__&quot;:
    test_jax_benchmarks()"><pre class="notranslate"><span class="pl-s">"goal: play with compiled benchmarks for compiled functions in jax"</span>
<span class="pl-k">from</span> <span class="pl-s1">functools</span> <span class="pl-k">import</span> <span class="pl-s1">partial</span>
<span class="pl-k">from</span> <span class="pl-s1">itertools</span> <span class="pl-k">import</span> <span class="pl-s1">product</span>
<span class="pl-k">from</span> <span class="pl-s1">typing</span> <span class="pl-k">import</span> <span class="pl-v">Callable</span>, <span class="pl-v">List</span>, <span class="pl-v">Tuple</span>
<span class="pl-k">from</span> <span class="pl-s1">time</span> <span class="pl-k">import</span> <span class="pl-s1">perf_counter</span>

<span class="pl-k">from</span> <span class="pl-s1">func_timeout</span> <span class="pl-k">import</span> <span class="pl-s1">func_timeout</span>
<span class="pl-k">from</span> <span class="pl-s1">pydantic</span> <span class="pl-k">import</span> <span class="pl-v">BaseSettings</span>
<span class="pl-k">import</span> <span class="pl-s1">pandas</span> <span class="pl-k">as</span> <span class="pl-s1">pd</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>

<span class="pl-s1">jnp</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span>
<span class="pl-s1">jit</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">jit</span>

<span class="pl-c"># config</span>
<span class="pl-k">class</span> <span class="pl-v">Environment</span>(<span class="pl-v">BaseSettings</span>):
    <span class="pl-s1">n_x</span>: <span class="pl-v">List</span>[<span class="pl-s1">int</span>] <span class="pl-c1">=</span> [<span class="pl-c1">4</span>, <span class="pl-c1">8</span>, <span class="pl-c1">16</span>, <span class="pl-c1">32</span>, <span class="pl-c1">64</span>, <span class="pl-c1">128</span>, <span class="pl-c1">256</span>, <span class="pl-c1">512</span>, <span class="pl-c1">1024</span>]
    <span class="pl-s1">min_weight</span>: <span class="pl-s1">float</span> <span class="pl-c1">=</span> <span class="pl-c1">-</span><span class="pl-c1">1.0</span>
    <span class="pl-s1">max_weight</span>: <span class="pl-s1">float</span> <span class="pl-c1">=</span> <span class="pl-c1">1.0</span>
    <span class="pl-s1">n_trials</span>: <span class="pl-s1">int</span> <span class="pl-c1">=</span> <span class="pl-c1">10</span>
    <span class="pl-s1">timeout_seconds</span>: <span class="pl-s1">int</span> <span class="pl-c1">=</span> <span class="pl-c1">30</span>
    <span class="pl-s1">seed</span>: <span class="pl-s1">int</span> <span class="pl-c1">=</span> <span class="pl-c1">42</span>


<span class="pl-s1">env</span> <span class="pl-c1">=</span> <span class="pl-v">Environment</span>()

<span class="pl-c"># types</span>
<span class="pl-v">PRNGKey</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">_src</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>
<span class="pl-v">Array</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-s1">ndarray</span>


<span class="pl-c"># functions - 'benchmark'</span>
<span class="pl-k">def</span> <span class="pl-en">benchmark_no_jit</span>(
    <span class="pl-s1">x</span>: <span class="pl-v">Array</span>,
    <span class="pl-s1">functions</span>: <span class="pl-v">List</span>[<span class="pl-v">Callable</span>],
    <span class="pl-s1">timeout_seconds</span>: <span class="pl-s1">int</span> <span class="pl-c1">=</span> <span class="pl-s1">env</span>.<span class="pl-s1">timeout_seconds</span>,
    <span class="pl-s1">min_weight</span>: <span class="pl-s1">float</span> <span class="pl-c1">=</span> <span class="pl-s1">env</span>.<span class="pl-s1">min_weight</span>,
    <span class="pl-s1">max_weight</span>: <span class="pl-s1">float</span> <span class="pl-c1">=</span> <span class="pl-s1">env</span>.<span class="pl-s1">max_weight</span>,
) <span class="pl-c1">-&gt;</span> <span class="pl-v">Tuple</span>[<span class="pl-v">Array</span>, <span class="pl-v">Array</span>]:
    <span class="pl-s1">n_x</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">0</span>]
    <span class="pl-s1">n_functions</span> <span class="pl-c1">=</span> <span class="pl-en">len</span>(<span class="pl-s1">functions</span>)
    <span class="pl-s1">outputs</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">empty</span>((<span class="pl-s1">n_functions</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>, <span class="pl-s1">n_x</span>)).<span class="pl-s1">at</span>[<span class="pl-c1">-</span><span class="pl-c1">1</span>, :].<span class="pl-en">set</span>(<span class="pl-c1">1.0</span>)
    <span class="pl-s1">timings</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">empty</span>(<span class="pl-s1">n_functions</span>)
    <span class="pl-k">for</span> <span class="pl-s1">i</span>, <span class="pl-s1">fun</span> <span class="pl-c1">in</span> <span class="pl-en">enumerate</span>(<span class="pl-s1">functions</span>):
        <span class="pl-k">try</span>:
            <span class="pl-c"># Without JIT, we can use func_timeout to set timeouts</span>
            <span class="pl-c"># and measure runtime with time.perf_counter</span>
            <span class="pl-s1">start</span> <span class="pl-c1">=</span> <span class="pl-en">perf_counter</span>()
            <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-en">func_timeout</span>(<span class="pl-s1">timeout_seconds</span>, <span class="pl-s1">fun</span>, <span class="pl-s1">args</span><span class="pl-c1">=</span>(<span class="pl-s1">x</span>,))
            <span class="pl-s1">elapsed</span> <span class="pl-c1">=</span> <span class="pl-en">perf_counter</span>() <span class="pl-c1">-</span> <span class="pl-s1">start</span>
            <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">y</span> <span class="pl-k">if</span> <span class="pl-en">isinstance</span>(<span class="pl-s1">y</span>, <span class="pl-v">Array</span>) <span class="pl-c1">and</span> <span class="pl-s1">y</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">0</span>] <span class="pl-c1">==</span> <span class="pl-s1">n_x</span> <span class="pl-k">else</span> <span class="pl-c1">0</span>
        <span class="pl-k">except</span> <span class="pl-v">Exception</span> <span class="pl-k">as</span> <span class="pl-s1">e</span>:
            <span class="pl-en">print</span>(<span class="pl-s1">e</span>)
            <span class="pl-s1">y</span>, <span class="pl-s1">elapsed</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>, <span class="pl-s1">timeout</span>
        <span class="pl-s1">outputs</span> <span class="pl-c1">=</span> <span class="pl-s1">outputs</span>.<span class="pl-s1">at</span>[<span class="pl-s1">i</span>, :].<span class="pl-en">set</span>(<span class="pl-s1">y</span>)
        <span class="pl-s1">timings</span> <span class="pl-c1">=</span> <span class="pl-s1">timings</span>.<span class="pl-s1">at</span>[<span class="pl-s1">i</span>].<span class="pl-en">set</span>(<span class="pl-s1">elapsed</span>)
    <span class="pl-s1">outputs</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">nan_to_num</span>(
        <span class="pl-s1">jnp</span>.<span class="pl-en">clip</span>(<span class="pl-s1">outputs</span>, <span class="pl-s1">a_min</span><span class="pl-c1">=</span><span class="pl-s1">min_weight</span>, <span class="pl-s1">a_max</span><span class="pl-c1">=</span><span class="pl-s1">max_weight</span>)
    ).<span class="pl-en">transpose</span>()
    <span class="pl-k">return</span> <span class="pl-s1">outputs</span>, <span class="pl-s1">timings</span>


<span class="pl-c"># something like this...</span>
<span class="pl-en">@<span class="pl-en">partial</span>(<span class="pl-s1">jit</span>, <span class="pl-s1">static_argnums</span><span class="pl-c1">=</span>(<span class="pl-c1">0</span>,<span class="pl-c1">1</span>))</span>
<span class="pl-k">def</span> <span class="pl-en">_benchmark_with_jit</span>(
    <span class="pl-s1">n_x</span>: <span class="pl-s1">int</span>,
    <span class="pl-s1">n_functions</span>: <span class="pl-s1">int</span>,
    <span class="pl-s1">x</span>: <span class="pl-v">Array</span>,
    <span class="pl-s1">functions</span>: <span class="pl-v">List</span>[<span class="pl-v">Callable</span>],
    <span class="pl-s1">timeout_seconds</span>: <span class="pl-s1">int</span> <span class="pl-c1">=</span> <span class="pl-s1">env</span>.<span class="pl-s1">timeout_seconds</span>,
    <span class="pl-s1">min_weight</span>: <span class="pl-s1">float</span> <span class="pl-c1">=</span> <span class="pl-s1">env</span>.<span class="pl-s1">min_weight</span>,
    <span class="pl-s1">max_weight</span>: <span class="pl-s1">float</span> <span class="pl-c1">=</span> <span class="pl-s1">env</span>.<span class="pl-s1">max_weight</span>,
) <span class="pl-c1">-&gt;</span> <span class="pl-v">Tuple</span>[<span class="pl-v">Array</span>, <span class="pl-v">Array</span>]:
    <span class="pl-s1">outputs</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">empty</span>((<span class="pl-s1">n_functions</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>, <span class="pl-s1">n_x</span>)).<span class="pl-s1">at</span>[<span class="pl-c1">-</span><span class="pl-c1">1</span>, :].<span class="pl-en">set</span>(<span class="pl-c1">1.0</span>)
    <span class="pl-s1">timings</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">empty</span>(<span class="pl-s1">n_functions</span>)
    <span class="pl-k">for</span> <span class="pl-s1">i</span>, <span class="pl-s1">fun</span> <span class="pl-c1">in</span> <span class="pl-en">enumerate</span>(<span class="pl-s1">functions</span>):
        <span class="pl-k">try</span>:
            <span class="pl-k">with</span> <span class="pl-s1">jax</span>.<span class="pl-s1">profiler</span>.<span class="pl-en">trace</span>() <span class="pl-k">as</span> <span class="pl-s1">trace</span>:
                <span class="pl-c"># 1: set function timeouts / time limits (otherwise, search time is infinite)</span>
                <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">trace</span>.<span class="pl-en">timeout</span>(<span class="pl-s1">timeout_seconds</span>, <span class="pl-s1">fun</span>, <span class="pl-s1">args</span><span class="pl-c1">=</span>(<span class="pl-s1">x</span>,))
                <span class="pl-c"># 2: measure how long functions take to run (to find faster functions)</span>
                <span class="pl-s1">elapsed</span> <span class="pl-c1">=</span> <span class="pl-s1">trace</span>.<span class="pl-en">elapsed</span>()
                <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">y</span> <span class="pl-k">if</span> <span class="pl-en">isinstance</span>(<span class="pl-s1">y</span>, <span class="pl-v">Array</span>) <span class="pl-c1">and</span> <span class="pl-s1">y</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">0</span>] <span class="pl-c1">==</span> <span class="pl-s1">n_x</span> <span class="pl-k">else</span> <span class="pl-c1">0</span>
        <span class="pl-k">except</span> <span class="pl-v">Exception</span> <span class="pl-k">as</span> <span class="pl-s1">_</span>:
            <span class="pl-s1">y</span>, <span class="pl-s1">elapsed</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>, <span class="pl-s1">timeout</span>
        <span class="pl-s1">outputs</span> <span class="pl-c1">=</span> <span class="pl-s1">outputs</span>.<span class="pl-s1">at</span>[<span class="pl-s1">i</span>, :].<span class="pl-en">set</span>(<span class="pl-s1">y</span>)
        <span class="pl-s1">timings</span> <span class="pl-c1">=</span> <span class="pl-s1">timings</span>.<span class="pl-s1">at</span>[<span class="pl-s1">i</span>].<span class="pl-en">set</span>(<span class="pl-s1">elapsed</span>)
    <span class="pl-s1">outputs</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">nan_to_num</span>(
        <span class="pl-s1">jnp</span>.<span class="pl-en">clip</span>(<span class="pl-s1">outputs</span>, <span class="pl-s1">a_min</span><span class="pl-c1">=</span><span class="pl-s1">min_weight</span>, <span class="pl-s1">a_max</span><span class="pl-c1">=</span><span class="pl-s1">max_weight</span>)
    ).<span class="pl-en">transpose</span>()
    <span class="pl-k">return</span> <span class="pl-s1">outputs</span>, <span class="pl-s1">timings</span>


<span class="pl-k">def</span> <span class="pl-en">f1</span>(<span class="pl-s1">x</span>: <span class="pl-v">Array</span>) <span class="pl-c1">-&gt;</span> <span class="pl-v">Array</span>:
    <span class="pl-k">return</span> <span class="pl-s1">x</span> <span class="pl-c1">**</span> <span class="pl-c1">2</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>


<span class="pl-k">def</span> <span class="pl-en">f2</span>(<span class="pl-s1">x</span>: <span class="pl-v">Array</span>) <span class="pl-c1">-&gt;</span> <span class="pl-v">Array</span>:
    <span class="pl-k">return</span> <span class="pl-s1">x</span> <span class="pl-c1">/</span> <span class="pl-c1">2</span>


<span class="pl-k">def</span> <span class="pl-en">f3</span>(<span class="pl-s1">x</span>: <span class="pl-v">Array</span>) <span class="pl-c1">-&gt;</span> <span class="pl-v">Array</span>:
    <span class="pl-k">return</span> <span class="pl-s1">x</span> <span class="pl-c1">**</span> <span class="pl-c1">3</span> <span class="pl-c1">*</span> <span class="pl-s1">jnp</span>.<span class="pl-en">sin</span>(<span class="pl-s1">x</span>)


<span class="pl-s1">functions</span> <span class="pl-c1">=</span> [<span class="pl-en">jit</span>(<span class="pl-s1">f</span>) <span class="pl-k">for</span> <span class="pl-s1">f</span> <span class="pl-c1">in</span> (<span class="pl-s1">f1</span>, <span class="pl-s1">f2</span>, <span class="pl-s1">f3</span>)]

<span class="pl-c"># because we can</span>
<span class="pl-k">def</span> <span class="pl-en">benchmark_benchmark</span>(
    <span class="pl-s1">benchmark</span>: <span class="pl-v">Callable</span> <span class="pl-c1">=</span> <span class="pl-s1">benchmark_no_jit</span>,
    <span class="pl-s1">words</span>: <span class="pl-s1">str</span> <span class="pl-c1">=</span> <span class="pl-s">"Python + func_timeout + time.perf_counter"</span>,
    <span class="pl-s1">functions</span>: <span class="pl-v">List</span>[<span class="pl-v">Callable</span>] <span class="pl-c1">=</span> <span class="pl-s1">functions</span>,
    <span class="pl-v">N_x</span>: <span class="pl-v">List</span>[<span class="pl-s1">int</span>] <span class="pl-c1">=</span> <span class="pl-s1">env</span>.<span class="pl-s1">n_x</span>,
    <span class="pl-s1">seed</span>: <span class="pl-s1">int</span> <span class="pl-c1">=</span> <span class="pl-s1">env</span>.<span class="pl-s1">seed</span>,
) <span class="pl-c1">-&gt;</span> <span class="pl-s1">float</span>:
    <span class="pl-s1">n_functions</span> <span class="pl-c1">=</span> <span class="pl-en">len</span>(<span class="pl-s1">functions</span>)
    <span class="pl-s1">begin</span> <span class="pl-c1">=</span> <span class="pl-en">perf_counter</span>()
    <span class="pl-s1">rng</span> <span class="pl-c1">=</span> <span class="pl-v">PRNGKey</span>(<span class="pl-s1">seed</span>)
    <span class="pl-s1">df</span> <span class="pl-c1">=</span> []  <span class="pl-c"># i admit to sucking at pandas</span>
    <span class="pl-k">for</span> <span class="pl-s1">i</span>, <span class="pl-s1">n_x</span> <span class="pl-c1">in</span> <span class="pl-en">product</span>(<span class="pl-en">range</span>(<span class="pl-s1">env</span>.<span class="pl-s1">n_trials</span>), <span class="pl-v">N_x</span>):
        <span class="pl-s1">rng</span>, <span class="pl-s1">key</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">split</span>(<span class="pl-s1">rng</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-s1">key</span>, (<span class="pl-s1">n_x</span>,))
        <span class="pl-k">if</span> <span class="pl-en">isinstance</span>(<span class="pl-s1">benchmark</span>, <span class="pl-s1">dict</span>):
            <span class="pl-en">print</span>(<span class="pl-s">"n_x"</span>, <span class="pl-s1">n_x</span>)
            <span class="pl-s1">benchmark</span> <span class="pl-c1">=</span> <span class="pl-s1">benchmark</span>[<span class="pl-s1">n_x</span>]
            <span class="pl-en">print</span>(<span class="pl-s">"benchmark"</span>, <span class="pl-s1">benchmark</span>)
            <span class="pl-en">print</span>(<span class="pl-s">"x"</span>, <span class="pl-s1">x</span>)
        <span class="pl-s1">outputs</span>, <span class="pl-s1">timings</span> <span class="pl-c1">=</span> <span class="pl-en">benchmark</span>(<span class="pl-s1">x</span>, <span class="pl-s1">functions</span>)
        <span class="pl-c"># print(timings)</span>
        <span class="pl-c"># print(outputs.shape)</span>
        <span class="pl-c"># print(i, n_x, timings)</span>
        <span class="pl-k">assert</span> <span class="pl-s1">outputs</span>.<span class="pl-s1">shape</span> <span class="pl-c1">==</span> (<span class="pl-s1">n_x</span>, <span class="pl-s1">n_functions</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>)
        <span class="pl-k">assert</span> <span class="pl-s1">timings</span>.<span class="pl-s1">shape</span> <span class="pl-c1">==</span> (<span class="pl-s1">n_functions</span>,)
        <span class="pl-k">assert</span> (<span class="pl-s1">timings</span> <span class="pl-c1">&gt;</span> <span class="pl-c1">0</span>).<span class="pl-en">all</span>()
        <span class="pl-k">assert</span> (<span class="pl-s1">timings</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">env</span>.<span class="pl-s1">timeout</span>).<span class="pl-en">all</span>()
        <span class="pl-k">assert</span> (<span class="pl-s1">env</span>.<span class="pl-s1">min_weight</span> <span class="pl-c1">&lt;=</span> <span class="pl-s1">outputs</span>).<span class="pl-en">all</span>()
        <span class="pl-k">assert</span> (<span class="pl-s1">env</span>.<span class="pl-s1">max_weight</span> <span class="pl-c1">&gt;=</span> <span class="pl-s1">outputs</span>).<span class="pl-en">all</span>()
        <span class="pl-c"># print(f"All assertions passed.")</span>
        <span class="pl-s1">df</span>.<span class="pl-en">append</span>(
            {
                <span class="pl-s">"n_x"</span>: <span class="pl-s1">n_x</span>,
                <span class="pl-s">"i"</span>: <span class="pl-s1">i</span>,
                <span class="pl-c1">**</span>{<span class="pl-s">f"f<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">i</span><span class="pl-c1">+</span><span class="pl-c1">1</span><span class="pl-kos">}</span></span> elapsed"</span>: <span class="pl-s1">t</span> <span class="pl-k">for</span> <span class="pl-s1">i</span>, <span class="pl-s1">t</span> <span class="pl-c1">in</span> <span class="pl-en">enumerate</span>(<span class="pl-s1">timings</span>)},
            }
        )
    <span class="pl-s1">df</span> <span class="pl-c1">=</span> <span class="pl-s1">pd</span>.<span class="pl-v">DataFrame</span>(<span class="pl-s1">df</span>).<span class="pl-en">set_index</span>([<span class="pl-s">"n_x"</span>, <span class="pl-s">"i"</span>])
    <span class="pl-s1">elapsed</span> <span class="pl-c1">=</span> <span class="pl-en">perf_counter</span>() <span class="pl-c1">-</span> <span class="pl-s1">begin</span>
    <span class="pl-en">print</span>(<span class="pl-s1">df</span>)
    <span class="pl-en">print</span>(<span class="pl-s">f"<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">words</span><span class="pl-kos">}</span></span> took <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">elapsed</span>:.2f<span class="pl-kos">}</span></span> seconds"</span>)
    <span class="pl-k">return</span> <span class="pl-s1">elapsed</span>


<span class="pl-k">def</span> <span class="pl-en">test_jax_benchmarks</span>():
    <span class="pl-s1">elapsed_no_jit</span> <span class="pl-c1">=</span> <span class="pl-en">benchmark_benchmark</span>()
    <span class="pl-s1">elapsed_with_jit</span> <span class="pl-c1">=</span> <span class="pl-en">benchmark_benchmark</span>(
        <span class="pl-s1">benchmark</span><span class="pl-c1">=</span>{
            <span class="pl-s1">n</span>: <span class="pl-en">partial</span>(<span class="pl-s1">_benchmark_with_jit</span>, <span class="pl-s1">n_x</span><span class="pl-c1">=</span><span class="pl-s1">n</span>, <span class="pl-s1">n_functions</span><span class="pl-c1">=</span><span class="pl-en">len</span>(<span class="pl-s1">functions</span>))
            <span class="pl-k">for</span> <span class="pl-s1">n</span> <span class="pl-c1">in</span> <span class="pl-s1">env</span>.<span class="pl-s1">n_x</span>
        },
        <span class="pl-s1">words</span><span class="pl-c1">=</span><span class="pl-s">"XLA (jax.jit) + jax.profiler.trace"</span>,
    )
    <span class="pl-en">print</span>(<span class="pl-s">"done!"</span>)


<span class="pl-k">if</span> <span class="pl-s1">__name__</span> <span class="pl-c1">==</span> <span class="pl-s">"__main__"</span>:
    <span class="pl-en">test_jax_benchmarks</span>()</pre></div>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> Check for duplicate requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> Describe your goal, and if possible provide a code snippet with a motivating example.<br>
Support complete <code class="notranslate">numpy</code> functions and some <code class="notranslate">SciPy</code> functions</li>
</ul>