# Copyright 2024 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

from functools import partial
import jax.numpy as jnp
from jax import jit

a = jnp.asarray([1, 2, 3, 4, 5, 6, 7, 8, 9])


@partial(jit, compiler_options={"iree-scheduling-dump-statistics-format": "csv"})
def f(a, b):
    return a + b


print(f(a, a))